sync the tsingmicro backend patches

Author: tsingmicro-public-e

third_party/tle/include/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(tle-dsa/Dialect/IR)
+

third_party/tle/include/tle-dsa/Dialect/IR/CMakeLists.txt

@@ -12,3 +12,4 @@ set(LLVM_TARGET_DEFINITIONS DsaOps.td)
 mlir_tablegen(DsaOps.h.inc -gen-op-decls)
 mlir_tablegen(DsaOps.cpp.inc -gen-op-defs)
 add_public_tablegen_target(TleDsaOpsIncGen)
+

third_party/tle/include/tle-dsa/Dialect/IR/DsaDialect.h

@@ -15,8 +15,8 @@
 
 // DsaOps.td uses TT_Tensor / TT_Ptr / TT_Int type constraints from the
 // Triton dialect, so the generated verifiers need these types visible.
-#include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Types.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
 
 #include "tle-dsa/Dialect/IR/DsaOpsDialect.h.inc"
 
@@ -31,3 +31,4 @@ class PatternRewriter;
 #include "tle-dsa/Dialect/IR/DsaOps.h.inc"
 
 #endif // TLE_DSA_DIALECT_IR_DSADIALECT_H
+

third_party/tle/include/tle-dsa/Dialect/IR/DsaDialect.td

@@ -66,3 +66,4 @@ class Dsa_Op<string mnemonic, list<Trait> traits = []> :
     Op<Dsa_Dialect, mnemonic, traits>;
 
 #endif // TLE_DSA_DIALECT
+

third_party/tle/include/tle-dsa/Dialect/IR/DsaOps.td

@@ -72,8 +72,31 @@ def Dsa_DistributedBarrierOp : Dsa_Op<"distributed_barrier",
 
 def Dsa_RemotePointersOp : Dsa_Op<"remote_pointers", [Pure]> {
   let arguments = (ins DsaRemotePointerType:$src, DsaRemoteShardIdType:$shard_id,
-                   OptionalAttr<DenseI32ArrayAttr>:$mesh_physical_ids);
+                   OptionalAttr<DenseI32ArrayAttr>:$mesh_physical_ids,
+                   OptionalAttr<DenseI32ArrayAttr>:$mesh_shape);
   let results = (outs DsaRemotePointerType:$result);
 }
 
+def DsaCumsumInputType : RankedTensorOf<[TT_Float, TT_Int]>;
+def DsaCumsumTotalResultType : AnyTypeOf<[DsaCumsumInputType, TT_Float, TT_Int]>;
+
+//===----------------------------------------------------------------------===//
+// dsa.cumsum
+//===----------------------------------------------------------------------===//
+
+def Dsa_CumsumOp : Dsa_Op<"cumsum", [Pure]> {
+  let summary = "exclusive cumsum marker lowered by Tsingmicro CRT";
+  let arguments = (ins
+    DsaCumsumInputType:$input,
+    I32Attr:$axis,
+    BoolAttr:$reverse,
+    DenseI64ArrayAttr:$shape,
+    I64Attr:$pad
+  );
+  let results = (outs
+    DsaCumsumInputType:$exclusive,
+    DsaCumsumTotalResultType:$total
+  );
+}
+
 #endif // TLE_DSA_OPS

third_party/tle/lib/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(Dialect/IR)
 add_subdirectory(Conversion/DsaToCore)
+

third_party/tle/lib/Conversion/DsaToCore/DsaToCore.cpp

@@ -26,7 +26,8 @@ struct DsaAllocToMemRefPattern : public OpRewritePattern<mlir::dsa::AllocOp> {
       return failure();
     // tx81-memref-to-llvm expects integer/default memref address spaces.
     // Canonicalize away non-integer memory-space attrs (e.g. "local").
-    if (Attribute ms = memrefTy.getMemorySpace(); ms && !isa<IntegerAttr>(ms)) {
+    if (Attribute ms = memrefTy.getMemorySpace();
+        ms && !isa<IntegerAttr>(ms)) {
       memrefTy = MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
                                  memrefTy.getLayout());
     }
@@ -56,7 +57,8 @@ struct DsaMemoryToCorePass
     RewritePatternSet patterns(&getContext());
     patterns.add<DsaAllocToMemRefPattern, DsaCopyToMemRefPattern>(
         &getContext());
-    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+    if (failed(applyPatternsGreedily(getOperation(),
+                                            std::move(patterns))))
       signalPassFailure();
   }
 };

third_party/tle/lib/Dialect/IR/CMakeLists.txt

@@ -17,3 +17,4 @@ target_include_directories(TleDsaIR
     ${TLE_INCLUDE_SOURCE}
     ${TLE_INCLUDE_BINARY}
 )
+

third_party/tle/lib/Dialect/IR/DsaDialect.cpp

@@ -37,3 +37,4 @@ void DsaDialect::registerTypes() {
 
 #define GET_TYPEDEF_CLASSES
 #include "tle-dsa/Dialect/IR/DsaOpsTypes.cpp.inc"
+

third_party/tle/python/CMakeLists.txt

@@ -5,3 +5,4 @@ if(TRITON_BUILD_PYTHON_MODULE)
   )
   target_link_libraries(TritonTleDsaTemplate PRIVATE Python3::Module pybind11::headers)
 endif()
+