Support index_select style gather/scatter access.

Save stride for unstructured dimension along chain of muls.
Support i64 index for TTS_MakeGatherScatterTensorPtrOp.
Fix the stride issue caused by the offset and strides of subview are relative to the strided memref of the input memref while the offset and stride of reinterpret_cast are relative to the base underlying memory of the memref.

Author: Xiang Li

include/triton-shared/AnalysisStructured/PtrAnalysis.h

@@ -45,8 +45,10 @@ const extern std::string ptrAnalysisAttr;
 // unstructured offsets. Later, when using the tensor offset to calculate the
 // address, it will be collapsed to 1D. To support gather/scatter access, treat
 // the unstructured offset as a whole offset instead of decoding the pointer
-// arithmetic on it. The stride is set to 1 so it still matches the offset *
-// stride formula
+// arithmetic on it except scalar mul.
+// The stride is set to 1 when there's no scalar mul so it still matches the offset *
+// stride formula. When there're scalar muls, the stride is set to the multiplication
+// of all the scalar strides.
 struct PtrState {
   SmallVector<OpFoldResult> offsets;
   SmallVector<OpFoldResult> sizes;

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

@@ -120,6 +120,8 @@ def TTS_MakeTensorPtrOp
   //let hasCanonicalizer = 1;
 }
 
+def TT_IndexTensorLike : AnyTypeOf<[I32Tensor, I64Tensor]>;
+
 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.
@@ -160,7 +162,7 @@ def TTS_MakeGatherScatterTensorPtrOp
   // result:  A tensor of pointers.
 
   let arguments = (ins TT_Ptr:$base,
-                       I32Tensor:$gather_scatter_offset,
+                       TT_IndexTensorLike:$gather_scatter_offset,
                        I32Attr:$gather_scatter_dim,
                        DenseI64ArrayAttr:$sizes,
                        Variadic<Index>:$strides,

lib/AnalysisStructured/PtrAnalysis.cpp

@@ -398,8 +398,10 @@ LogicalResult PtrState::mulState(const PtrState &lhsState,
       OpFoldResult newOffset =
           mulOFRs(lhs->offsets[i], rhsStride, loc, builder);
       offsets.push_back(newOffset);
-      // Set stride to 1 when not continuous.
-      strides.push_back(builder.getIndexAttr(1));
+      // Mul the scalart to stride.
+      OpFoldResult newStride =
+          mulOFRs(lhs->strides[i], rhs->scalar, loc, builder);
+      strides.push_back(newStride);
     }
     OpFoldResult newShape =
         mulOFRs(lhs->shape[i], rhs->scalar, loc, builder);

lib/Conversion/StructuredToMemref/StructuredToMemref.cpp

@@ -865,6 +865,9 @@ struct LoadConverter : public OpConversionPattern<tts::LoadOp> {
                                                   gatherOffsetElt.getResult(),
                                                   gatherDim, rewriter);
     unsigned rank = ptr.getSizes().size();
+    // Set strides to 1 for subview multiplies the existing strides with the
+    // stride of the subview.
+    SmallVector<OpFoldResult> oneStrides(rank, OpFoldResult(step));
     // subview from srcPtr for mask.
     // With offsets[gatherDim] set to 0 since the offset already in
     // reinterpret_cast. With sizes[gatherDim] set to 1 since we are load one
@@ -875,24 +878,24 @@ struct LoadConverter : public OpConversionPattern<tts::LoadOp> {
       sizes = mixedDims;
       // maskOffsets should be all zero, since srcPtr already has the offsets.
       SmallVector<OpFoldResult> maskOffsets(rank, OpFoldResult(lowerBound));
-      // Use allocStrides for subview.
+      // Use oneStrides for subview.
       auto dstSubViewType = memref::SubViewOp::inferResultType(
-          cast<MemRefType>(srcPtr.getType()), maskOffsets, sizes, allocStrides);
+          cast<MemRefType>(srcPtr.getType()), maskOffsets, sizes, oneStrides);
       srcPtr =
           rewriter
               .create<memref::SubViewOp>(loc, cast<MemRefType>(dstSubViewType),
-                                         srcPtr, maskOffsets, sizes, allocStrides)
+                                         srcPtr, maskOffsets, sizes, oneStrides)
               .getResult();
     }
 
     // alloc[inductionVar]
     SmallVector<OpFoldResult> allocOffsets(rank, OpFoldResult(lowerBound));
     allocOffsets[gatherDim] = inductionVar;
     auto dstAllocType = memref::SubViewOp::inferResultType(
-        allocType, allocOffsets, sizes, allocStrides);
+        allocType, allocOffsets, sizes, oneStrides);
     auto dstSubview = rewriter.create<memref::SubViewOp>(
         loc, cast<MemRefType>(dstAllocType), alloc, allocOffsets, sizes,
-        allocStrides);
+        oneStrides);
     // Copy srcPtr to alloc[inductionVar].
     rewriter.create<memref::CopyOp>(loc, srcPtr, dstSubview);
 
@@ -990,8 +993,11 @@ struct StoreConverter : public OpConversionPattern<tts::StoreOp> {
       mixedDims[gatherDim] = sizes[gatherDim];
       sizes = mixedDims;
     }
+    // Set strides to 1 for subview/extract_slice multiplies the existing strides with the
+    // stride of the subview.
+    SmallVector<OpFoldResult> oneStrides(rank, OpFoldResult(step));
     auto slice = rewriter.create<tensor::ExtractSliceOp>(
-        loc, stVal, stValOffsets, sizes, strides);
+        loc, stVal, stValOffsets, sizes, oneStrides);
 
     // reinterpret_cast to current row as memRefPtr[gatherOffsetElt].
     Value dstPtr = rewriteGatherScatterPtrElement(staticSizes, ptr, memRefPtr,
@@ -1003,15 +1009,14 @@ struct StoreConverter : public OpConversionPattern<tts::StoreOp> {
       // maskOffsets should be all zero, since srcPtr already has the offsets.
       SmallVector<OpFoldResult> maskOffsets(rank, OpFoldResult(lowerBound));
       auto dstType = memref::SubViewOp::inferResultType(
-          cast<MemRefType>(dstPtr.getType()), maskOffsets, sizes, strides);
+          cast<MemRefType>(dstPtr.getType()), maskOffsets, sizes, oneStrides);
 
       dstPtr =
           rewriter
               .create<memref::SubViewOp>(loc, cast<MemRefType>(dstType), dstPtr,
-                                         maskOffsets, sizes, strides)
+                                         maskOffsets, sizes, oneStrides)
               .getResult();
     }
-
     // store slice to dstPtr.
     auto storeOp = rewriter.create<bufferization::MaterializeInDestinationOp>(
         loc, slice, dstPtr);