[Tsingmicro] Sync backend code

include/triton-shared/Analysis/MaskAnalysis.h

@@ -85,6 +85,23 @@ struct dimInfo {
 //
 // Example of creating 2D mask:
 //  mask = (rows[:, None] < M) & (cols[None, :] < N)
+//
+// NOTE: Bool tensor mask could be saved into dims and record the dim in
+// nonContinuousDim in case that dimension failed MaskAnalysis.
+// These is to allow case where only one dimension failed while others passed. A
+// MakeGatherScatterTensorPtrOp operation could be generated for the failed
+// dimension. Only 3 patterns are supported for this.
+// 1. offsets[:, None] < n where the offsets is 1d tensor.
+//    It will in pattern of expandDims -> broadcast -> cmp
+// 2. mask[:, None] where mask is 1d bool tensor.
+//    It will in pattern of cmp -> expandDims -> broadcast
+// 3. scalar_mask[:, None] where scalar mask is scalar bool.
+//    It will in pattern of splat -> expandDims -> broadcast
+// These 3 patterns are only about how a bool tensor was created from 1D or
+// scalar bool. How the 1D and scalar bool were created is not important for the
+// unstructured mask.
+// Only one tensor mask is allowed. If multiple dimensions have failed
+// MaskAnalysis, then MaskAnalysis will still fail on the current operation.
 struct MaskState {
   OpFoldResult start;
   OpFoldResult end;
@@ -93,13 +110,24 @@ struct MaskState {
   const bool useUnsafeMask;
   ///ASCEND
   SmallVector<dimInfo> stateInfo;
-
+
+  SmallVector<int64_t> nonContinuousDim;
   void dump() const;
 
   MaskState(bool useUnsafeMask = false) : useUnsafeMask(useUnsafeMask) {}
 
   int64_t getRank() const { return dims.size(); }
 
+  bool hasNonContinuousDim() const { return !nonContinuousDim.empty(); }
+
+  std::optional<int64_t> getNonContinuousDim() const {
+    if (hasNonContinuousDim()) {
+      assert(nonContinuousDim.size() == 1);
+      return nonContinuousDim[0];
+    }
+    return std::nullopt;
+  }
+
   bool isEmpty() const { return getRank() == 0 && !scalar && !start && !end; }
 
   bool isMask() const { return !start && !end && !scalar && dims.size() != 0; }

include/triton-shared/Analysis/OpFoldResultUtils.h

@@ -24,6 +24,8 @@ Value materializeValue(OpBuilder &builder, Location loc, OpFoldResult ofr);
 // result of an operation too.
 std::optional<int64_t> getIntAttr(const OpFoldResult ofr);
 
+std::optional<int64_t> getConstValue(const OpFoldResult ofr);
+
 // Return if ofr contains a constant zero, either represented by an integer
 // attribute or a constant value.
 bool hasConstZero(const OpFoldResult ofr);

include/triton-shared/AnalysisStructured/PtrAnalysis.h

@@ -38,6 +38,11 @@ const extern std::string ptrAnalysisAttr;
 // shape field means the same field as tt.make_tensor_ptr; when it describes a
 // non-block pointer, shape field indicates how address wraps around (i.e.,
 // modulo); a constant 0 indicates no modulo for the dimension.
+// Multi-dimension PtrState, which has one unstructured dimension, is supported
+// for gather/scatter access. The unstructured dimension is marked by a tensor
+// type offset. The tensor offset for the unstructured dimension must be
+// expanded from a 1D tensor. The analysis will fail for multi-dimension
+// unstructured offsets.
 struct PtrState {
 
   SmallVector<OpFoldResult> offsets;
@@ -57,6 +62,10 @@ struct PtrState {
 
   bool dimHasModulo(uint32_t dim) const;
 
+  bool dimIsStructured(uint32_t dim) const;
+  int32_t getNonStructuredDim() const;
+  bool isStructured() const;
+
   bool isBlockPtr() const;
 
   void dump() const;
@@ -73,8 +82,8 @@ struct PtrState {
   LogicalResult mulState(const PtrState &lhsState, const PtrState &rhsState,
                          Operation *op, OpBuilder &builder);
 
-  tts::MakeTensorPtrOp createTTSMakeTensorPtrOp(OpBuilder &builder,
-                                                Location loc);
+  // All dim structured or only one dim unstructured
+  Operation *createTTSMakeTensorPtrOp(OpBuilder &builder, Location loc);
 };
 
 class PtrAnalysis {
@@ -162,6 +171,16 @@ class PtrAnalysis {
   LogicalResult visitOperandRem(arith::RemSIOp mulOp, PtrState &state,
                                 const Location loc, OpBuilder &builder);
 
+  // Operand is the result of arith.divsi (signed integer division).
+  // Main assumptions:
+  //  Divisor must be a constant scalar
+  //  Dividend must be 1D tensor or 2D tensor with one dimension being 1
+  //  Division is applied to offsets, strides, and scalar values
+  // Expected result:
+  //  All offset and stride values are divided by the constant divisor
+  LogicalResult visitOperandDivSI(arith::DivSIOp divOp, PtrState &state,
+                                  const Location loc, OpBuilder &builder);
+
   // Operand is the result of make_range.
   // Main assumptions:
   //  start, end, and shape are all statically known

include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td

@@ -18,12 +18,17 @@ def Triton_Structured_Dialect : Dialect {
   }];
 
   let dependentDialects = [
-    "triton::TritonDialect"
+    "triton::TritonDialect",
+    "tensor::TensorDialect"
   ];
 
   let usePropertiesForAttributes = 1;
+  let hasCanonicalizer = 1;
 }
 
+// FIXME: AnyTensor
+def TT_IndexTensorLike : AnyTypeOf<[I1Tensor, I8Tensor, I16Tensor, I32Tensor, I64Tensor, IndexTensor]>;
+
 //
 // Op Base
 //
@@ -121,6 +126,113 @@ def TTS_MakeTensorPtrOp
   //let hasCanonicalizer = 1;
 }
 
+def TTS_MakeGatherScatterTensorPtrOp
+  : TTS_Op<"make_gather_scatter_tptr", [AttrSizedOperandSegments, Pure]> {
+  // NOTE: Only support cases where the offset for each dimension is defined in a different operation.
+  //       Not support case where the offset is a tensor load from other ptr which for multiple dimension.
+  //
+  //       offset_m = tl.arange(0, M)
+  //       offset_n = tl.arange(0, N)
+  //       offset_k = tl.arange(0, K)
+  //       ld_offsets = tl.load(a_ptr + offset_m[:,None]+offsets_n[None,:])
+  //       not_support = tl.load(b_ptr + ld_offsets)
+  //       not_support2 = tl.load(b_ptr + ld_offsets * (offset_m[:,None]+offsets_n[None,:]))
+  //       not_support3 = tl.load(b_ptr + (ld_offsets * (offset_m[:,None]+offsets_n[None,:]))[:, :, None] + offset_k[None,None,:])
+  //
+  //       # Support cases where one dimension is structured while the other is not.
+  //       # For example, `offset_m[:, None] // K` is not structured, whereas `offset_n[None, :]` is structured in next line.
+  //       supported = tl.load(b_ptr + offset_m[:, None] // K + offset_n[None, :])
+  //       # `ld_offsets_1d[:, None]` is not structured, whereas `offset_n[None, :]`.
+  //       ld_offsets_1d = tl.load(a_ptr + offsets_m)
+  //       supported2 = tl.load(b_ptr + ld_offsets_1d[:, None] + offsets_n[None, :])
+
+  let summary = "create an pointer that points to a tensor in memory for gather/scatter";
+  let description = [{
+    The `tts.make_gather_scatter_tptr` operation is similar to `tts.make_tptr`.
+    The key difference is that `make_gather_scatter_tptr` accesses the tensor non-continuously.
+    Currently, only one dimension is allowed to be non-continuous.
+    This dimension is saved in `gather_scatter_dim`, and the offset for that dimension is saved in `gather_scatter_offset`.
+    Each contiguous load will load from this offset.
+    Cases with more than one non-continuous dimension are not supported.
+  }];
+
+  // base:    Base pointer used to contruct the tensor of pointers or pointer to tensor.
+  // gather_scatter_dim: The dimension for gather_scatter_offset.
+  // sizes:   Size of the data being loaded or stored.
+  // strides: The strides of the parent tensor, which means how much to increase the pointer
+  //          by when moving by 1 element in a specific axis.
+  // offsets: Offset of the block along each dimension from base.
+  // result:  A tensor of pointers.
+
+  let arguments = (ins TT_Ptr:$base,
+                       TT_IndexTensorLike:$gather_scatter_offset,
+                       I32Attr:$gather_scatter_dim,
+                       DenseI64ArrayAttr:$sizes,
+                       Variadic<Index>:$strides,
+                       Variadic<Index>:$offsets,
+                       DenseI64ArrayAttr:$static_strides,
+                       DenseI64ArrayAttr:$static_offsets,
+                       Optional<TT_BoolLike>:$gather_scatter_mask);
+
+  let results = (outs TT_PtrLike:$result);
+
+  let assemblyFormat = [{
+    $base `to` `sizes` `` `:` $sizes
+    `gather_scatter_dim` `` `:` $gather_scatter_dim
+    `gather_scatter_offset` `` `:` $gather_scatter_offset
+    (`gather_scatter_mask` `` `:` $gather_scatter_mask^)?
+    `` `,` `strides` `` `:`
+    custom<DynamicIndexList>($strides, $static_strides)
+    `` `,` `offsets` `` `:`
+    custom<DynamicIndexList>($offsets, $static_offsets)
+    attr-dict `:`  type($gather_scatter_offset) type($gather_scatter_mask) type($base) `to` type($result)
+  }];
+
+
+  let builders = [
+    // Build with mixed static and dynamic entries.
+    OpBuilder<(ins
+      "Value":$base,
+      "Value":$gather_scatter_offset,
+      "int":$gather_scatter_dim,
+      "ArrayRef<int64_t>":$sizes,
+      "ArrayRef<OpFoldResult>":$strides,
+      "ArrayRef<OpFoldResult>":$offsets)>,
+
+    OpBuilder<(ins
+      "Value":$base,
+      "Value":$gather_scatter_offset,
+      "Value":$gather_scatter_mask,
+      "int":$gather_scatter_dim,
+      "ArrayRef<int64_t>":$sizes,
+      "ArrayRef<OpFoldResult>":$strides,
+      "ArrayRef<OpFoldResult>":$offsets)>,
+  ];
+
+  let extraClassDeclaration = [{
+    /// Return a vector of all the static or dynamic fields
+    SmallVector<OpFoldResult> getMixedSizes() {
+      Builder b(getContext());
+      SmallVector<Value> dynSizes; // sizes are always static
+      return ::mlir::getMixedValues(getSizes(), dynSizes, b);
+    }
+    /// Return a vector of all the static or dynamic fields
+    SmallVector<OpFoldResult> getMixedStrides() {
+      Builder b(getContext());
+      return ::mlir::getMixedValues(getStaticStrides(), getStrides(), b);
+    }
+    SmallVector<OpFoldResult> getMixedOffsets() {
+      Builder b(getContext());
+      return ::mlir::getMixedValues(getStaticOffsets(), getOffsets(), b);
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCanonicalizer = 0;
+}
+
+// SameVariadicResultSize
+// AttrSizedResultSegments
 def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSegments, Pure]> {
   let summary = "Placeholder for the structured pointer states computed during PtrAnalysis.";
   let description = "Used to pass the offsets and strides to scf.for op to simplify IR rewrites.";

lib/Analysis/CMakeLists.txt

@@ -11,4 +11,5 @@ add_triton_library(TritonSharedAnalysis
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
+  Common
 )