diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td --- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td +++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td @@ -2177,6 +2177,10 @@ unsigned resultRank = getResult().getType().cast().getRank(); return {1, resultRank, resultRank}; } + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } }]; } @@ -3031,7 +3035,8 @@ //===----------------------------------------------------------------------===// def SubViewOp : BaseOpWithOffsetSizesAndStrides< - "subview", [DeclareOpInterfaceMethods, OffsetSizeAndStrideOpInterface] > { + "subview", [DeclareOpInterfaceMethods, + OffsetSizeAndStrideOpInterface] > { let summary = "memref subview operation"; let description = [{ The "subview" operation converts a memref type to another memref type @@ -3217,6 +3222,10 @@ 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; @@ -3227,7 +3236,8 @@ // 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 @@ -3279,12 +3289,12 @@ 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":$staticOffsets, "ArrayRef":$staticSizes, "ArrayRef":$staticStrides, "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, CArg<"ArrayRef", "{}">:$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", "{}">:$attrs)> @@ -3315,6 +3325,10 @@ 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; @@ -3324,7 +3338,8 @@ // 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 @@ -3369,13 +3384,13 @@ 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":$staticOffsets, "ArrayRef":$staticSizes, "ArrayRef":$staticStrides, "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, CArg<"ArrayRef", "{}">:$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", "{}">:$attrs)> @@ -3398,6 +3413,10 @@ 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; } }]; } diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.td b/mlir/include/mlir/Interfaces/ViewLikeInterface.td --- a/mlir/include/mlir/Interfaces/ViewLikeInterface.td +++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.td @@ -40,6 +40,8 @@ 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 @@ -53,6 +55,8 @@ 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 @@ -62,6 +66,15 @@ 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` @@ -69,11 +82,7 @@ }], /*retTy=*/"std::array", /*methodName=*/"getArrayAttrRanks", - /*args=*/(ins), - /*methodBody=*/"", - /*defaultImplementation=*/[{ - return $_op.offsets(); - }] + /*args=*/(ins) >, InterfaceMethod< /*desc=*/[{ @@ -250,9 +259,9 @@ assert($_op.isDynamicOffset(idx) && "expected dynamic offset"); auto numDynamic = getNumDynamicEntriesUpToIdx( static_offsets().template cast(), - ShapedType::isDynamicStrideOrOffset, + ShapedType::isDynamicStrideOrOffset, idx); - return 1 + numDynamic; + return $_op.getOffsetSizeAndStrideStartOperandIndex() + numDynamic; }] >, InterfaceMethod< @@ -268,7 +277,8 @@ assert($_op.isDynamicSize(idx) && "expected dynamic size"); auto numDynamic = getNumDynamicEntriesUpToIdx( static_sizes().template cast(), ShapedType::isDynamic, idx); - return 1 + offsets().size() + numDynamic; + return $_op.getOffsetSizeAndStrideStartOperandIndex() + + offsets().size() + numDynamic; }] >, InterfaceMethod< @@ -284,9 +294,10 @@ assert($_op.isDynamicStride(idx) && "expected dynamic stride"); auto numDynamic = getNumDynamicEntriesUpToIdx( static_strides().template cast(), - ShapedType::isDynamicStrideOrOffset, + ShapedType::isDynamicStrideOrOffset, idx); - return 1 + offsets().size() + sizes().size() + numDynamic; + return $_op.getOffsetSizeAndStrideStartOperandIndex() + + offsets().size() + sizes().size() + numDynamic; }] >, InterfaceMethod<