[mlir] Fix indexing of first offset operand in ops that implement OffsetSizeAndStrideOpInterface

[nicolasvasilache]

Summary

OffsetSizeAndStrideOpInterface ops may have a varying number of operands before the first
offset operand. This revision adds a method that such ops much implement to properly specify
the position of the first offset operand.

[bondhugula]

Missing a test case?

mlir/include/mlir/Interfaces/ViewLikeInterface.td

44

Nit: missing period at the end.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     }
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
     std::array<unsigned, 3> getArrayAttrRanks() {
       unsigned resultRank = getResult().getType().cast<ShapedType>().getRank();
       return {1, resultRank, resultRank};
     }
+
+    /// Return the number of leading operands before the `offsets`, `sizes` and
+    /// and `strides` operands.
+    static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefReshapeOp
 //===----------------------------------------------------------------------===//
 
 def MemRefReshapeOp: Std_Op<"memref_reshape", [
def SubIOp : IntArithmeticOp<"subi"> {
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SubViewOp
 //===----------------------------------------------------------------------===//
 
 def SubViewOp : BaseOpWithOffsetSizesAndStrides<
-    "subview", [DeclareOpInterfaceMethods<ViewLikeOpInterface>, OffsetSizeAndStrideOpInterface] >  {
+    "subview", [DeclareOpInterfaceMethods<ViewLikeOpInterface>,
+                OffsetSizeAndStrideOpInterface] >  {
   let summary = "memref subview operation";
   let description = [{
     The "subview" operation converts a memref type to another memref type
     which represents a reduced-size view of the original memref as specified by
     the operation's offsets, sizes and strides arguments.
 
     The SubView operation supports the following arguments:
 
static Type inferResultType(MemRefType sourceMemRefType,
                                 ArrayRef<int64_t> staticStrides);
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
     std::array<unsigned, 3> getArrayAttrRanks() {
       unsigned rank = getSourceType().getRank();
       return {rank, rank, rank};
     }
+
+    /// Return the number of leading operands before the `offsets`, `sizes` and
+    /// and `strides` operands.
+    static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SubTensorOp
 //===----------------------------------------------------------------------===//
 
-def SubTensorOp : BaseOpWithOffsetSizesAndStrides<"subtensor", [OffsetSizeAndStrideOpInterface]> {
+def SubTensorOp : BaseOpWithOffsetSizesAndStrides<
+    "subtensor", [OffsetSizeAndStrideOpInterface]> {
   let summary = "subtensor operation";
   let description = [{
     The "subtensor" operation extract a tensor from another tensor as
     specified by the operation's offsets, sizes and strides arguments.
 
     The subtensor operation supports the following arguments:
 
     * tensor: the "base" tensor from which to extract a subtensor.
let arguments = (ins
     Variadic<Index>:$strides,
     I64ArrayAttr:$static_offsets,
     I64ArrayAttr:$static_sizes,
     I64ArrayAttr:$static_strides
   );
   let results = (outs AnyRankedTensor:$result);
 
   let builders = [
-    // Build a SubViewOp with mixed static and dynamic entries.
+    // Build a SubTensorOp with mixed static and dynamic entries.
     OpBuilderDAG<(ins "Value":$source, "ArrayRef<int64_t>":$staticOffsets,
       "ArrayRef<int64_t>":$staticSizes, "ArrayRef<int64_t>":$staticStrides,
       "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
-    // Build a SubViewOp with all dynamic entries.
+    // Build a SubTensorOp with all dynamic entries.
     OpBuilderDAG<(ins "Value":$source, "ValueRange":$offsets,
       "ValueRange":$sizes, "ValueRange":$strides,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let extraClassDeclaration = extraBaseClassDeclaration # [{
     /// Returns the type of the base tensor operand.
     RankedTensorType getSourceType() {
static Type inferResultType(RankedTensorType sourceRankedTensorType,
                                 ArrayRef<int64_t> staticStrides);
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
     std::array<unsigned, 3> getArrayAttrRanks() {
       unsigned rank = getSourceType().getRank();
       return {rank, rank, rank};
     }
+
+    /// Return the number of leading operands before the `offsets`, `sizes` and
+    /// and `strides` operands.
+    static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
   }];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SubTensorInsertOp
 //===----------------------------------------------------------------------===//
 
-def SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides<"subtensor_insert", [OffsetSizeAndStrideOpInterface]> {
+def SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides<
+    "subtensor_insert", [OffsetSizeAndStrideOpInterface]> {
   let summary = "subtensor_insert operation";
   let description = [{
     The "subtensor_insert" operation insert a tensor `source` into another
     tensor `dest` as specified by the operation's offsets, sizes and strides
     arguments.
 
     It returns a copy of `dest` with the proper subtensor updated with the value
     of `source`.
let arguments = (ins
     Variadic<Index>:$strides,
     I64ArrayAttr:$static_offsets,
     I64ArrayAttr:$static_sizes,
     I64ArrayAttr:$static_strides
   );
   let results = (outs AnyRankedTensor:$result);
 
   let builders = [
-    // Build a SubViewOp with mixed static and dynamic entries.
+    // Build a SubTensorInsertOp with mixed static and dynamic entries.
     OpBuilderDAG<(ins "Value":$source, "Value":$dest,
       "ArrayRef<int64_t>":$staticOffsets, "ArrayRef<int64_t>":$staticSizes,
       "ArrayRef<int64_t>":$staticStrides, "ValueRange":$offsets,
       "ValueRange":$sizes, "ValueRange":$strides,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
-    // Build a SubViewOp with all dynamic entries.
+    // Build a SubTensorInsertOp with all dynamic entries.
     OpBuilderDAG<(ins "Value":$source, "Value":$dest, "ValueRange":$offsets,
       "ValueRange":$sizes, "ValueRange":$strides,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let extraClassDeclaration = extraBaseClassDeclaration # [{
     /// Returns the type of the base tensor operand.
     RankedTensorType getSourceType() {
       return source().getType().cast<RankedTensorType>();
     }
 
     /// The result of a subtensor is always a tensor.
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
     std::array<unsigned, 3> getArrayAttrRanks() {
       unsigned rank = getSourceType().getRank();
       return {rank, rank, rank};
     }
+
+    /// Return the number of leading operands before the `offsets`, `sizes` and
+    /// and `strides` operands.
+    static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TanhOp
 //===----------------------------------------------------------------------===//
 
 def TanhOp : FloatUnaryOp<"tanh"> {

mlir/include/mlir/Interfaces/ViewLikeInterface.td

let description = [{
     Common interface for ops that allow specifying mixed dynamic and static
     offsets, sizes and strides variadic operands.
     Ops that implement this interface need to expose the following methods:
       1. `getArrayAttrRanks` to specify the length of static integer
           attributes.
       2. `offsets`, `sizes` and `strides` variadic operands.
       3. `static_offsets`, resp. `static_sizes` and `static_strides` integer
           array attributes.
+      4. `getOffsetSizeAndStrideStartOperandIndex` method that specifies the
+         starting index of the OffsetSizeAndStrideOpInterface operands
 
     The invariants of this interface are:
       1. `static_offsets`, `static_sizes` and `static_strides` have length
           exactly `getArrayAttrRanks()`[0] (resp. [1], [2]).
       2. `offsets`, `sizes` and `strides` have each length at most
          `getArrayAttrRanks()`[0] (resp. [1], [2]).
       3. if an entry of `static_offsets` (resp. `static_sizes`,
          `static_strides`) is equal to a special sentinel value, namely
          `ShapedType::kDynamicStrideOrOffset` (resp. `ShapedType::kDynamicSize`,
          `ShapedType::kDynamicStrideOrOffset`), then the corresponding entry is
          a dynamic offset (resp. size, stride).
       4. a variadic `offset` (resp. `sizes`, `strides`) operand  must be present
          for each dynamic offset (resp. size, stride).
+      5. `offsets`, `sizes` and `strides` operands are specified in this order
+         at operand index starting at `getOffsetSizeAndStrideStartOperandIndex`.
 
     This interface is useful to factor out common behavior and provide support
     for carrying or injecting static behavior through the use of the static
     attributes.
   }];
 
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<
       /*desc=*/[{
+        Return the number of leading operands before the `offsets`, `sizes` and
+        and `strides` operands.
+      }],
+      /*retTy=*/"unsigned",
+      /*methodName=*/"getOffsetSizeAndStrideStartOperandIndex",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the expected rank of each of the`static_offsets`, `static_sizes`
         and `static_strides` attributes.
       }],
       /*retTy=*/"std::array<unsigned, 3>",
       /*methodName=*/"getArrayAttrRanks",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        return $_op.offsets();
-      }]
+      /*args=*/(ins)
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the dynamic offset operands.
       }],
       /*retTy=*/"OperandRange",
       /*methodName=*/"offsets",
       /*args=*/(ins),
InterfaceMethod<
       /*retTy=*/"unsigned",
       /*methodName=*/"getIndexOfDynamicOffset",
       /*args=*/(ins "unsigned":$idx),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert($_op.isDynamicOffset(idx) && "expected dynamic offset");
         auto numDynamic = getNumDynamicEntriesUpToIdx(
           static_offsets().template cast<ArrayAttr>(),
-          ShapedType::isDynamicStrideOrOffset, 
+          ShapedType::isDynamicStrideOrOffset,
           idx);
-        return 1 + numDynamic;
+        return $_op.getOffsetSizeAndStrideStartOperandIndex() + numDynamic;
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Assert the size `idx` is dynamic and return the position of the
         corresponding operand.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getIndexOfDynamicSize",
       /*args=*/(ins "unsigned":$idx),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert($_op.isDynamicSize(idx) && "expected dynamic size");
         auto numDynamic = getNumDynamicEntriesUpToIdx(
           static_sizes().template cast<ArrayAttr>(), ShapedType::isDynamic, idx);
-        return 1 + offsets().size() + numDynamic;
+        return $_op.getOffsetSizeAndStrideStartOperandIndex() +
+          offsets().size() + numDynamic;
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Assert the stride `idx` is dynamic and return the position of the
         corresponding operand.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getIndexOfDynamicStride",
       /*args=*/(ins "unsigned":$idx),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert($_op.isDynamicStride(idx) && "expected dynamic stride");
         auto numDynamic = getNumDynamicEntriesUpToIdx(
           static_strides().template cast<ArrayAttr>(),
-          ShapedType::isDynamicStrideOrOffset, 
+          ShapedType::isDynamicStrideOrOffset,
           idx);
-        return 1 + offsets().size() + sizes().size() + numDynamic;
+        return $_op.getOffsetSizeAndStrideStartOperandIndex() +
+          offsets().size() + sizes().size() + numDynamic;
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Helper method to compute the number of dynamic entries of `attr`, up to
         `idx` using `isDynamic` to determine whether an entry is dynamic.
       }],
       /*retTy=*/"unsigned",