[mlir] Add subtensor_insert operation

[nicolasvasilache]

Summary

This revision adds the subtensor_insert operation.
This operation is discussed in https://llvm.discourse.group/t/an-update-on-linalg-on-tensors/1878
and is useful for extending linalg + transformations on tensors.

[pifon2a]

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

3105

typo: mixed

3158

shouldn't it be always tensor-rank number of static offsets with all -1s in case we have none and rank - |static offsets| number of dynamic offsets?

3167

typo: be

3168

s/and be/and /

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2687

"||"?

2716

"||"?

3253–3254

Would this also work SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamicSize); ?

3255

Here staticStridesVector == staticOffsetsVector, can we have just one vector?

3275–3280

remove ? or is it "under construction"?

mlir/test/IR/core-ops.mlir

900

please, provide the test for the mixed case.

[nicolasvasilache]

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

3158

rephrased things here and everywhere, thanks for noting!

3168

rephrased a bit differently

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2687

if parseExtraOperand is not specified we want to skip calling it.

2716

false is the "success" here so if parseExtraOperand is not specified we do not want an extra resolveOperand.

3255

No, they don't have the same sentinel.
The dynamic size sentinel was historically encoded to -1 and evolving that has been proven very intrusive ..

[nicolasvasilache]

Address review comments.

[jpienaar]

This seems to have broken GCC 5 (https://buildkite.com/mlir/mlir-core/builds/8301#e6b6e5a7-85bd-4454-ae2b-da3b6cf84b8a). Could you update? Or should we revert?

[nicolasvasilache]

had a potential fix in flight but did not land, trying again

[bondhugula]

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

3147

insert -> inserts

3148–3149

This description would make it look like the op should be called a tensor_insert. Consider rephrasing along the lines: `... while specifying offsets, ..., and sizes for the result tensor".

3154

has the encodes -> encodes

3169

encodes that the -> specifies that the

3172

Doc description is missing an example.

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

let builders = [
     // Build a SubViewOp with all dynamic entries.
     OpBuilder<
       "Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, "
       "ArrayRef<NamedAttribute> attrs = {}">
   ];
 
   let extraClassDeclaration = extraBaseClassDeclaration # [{
     /// Returns the type of the base memref operand.
-    MemRefType getSourceMemRefType() {
+    MemRefType getSourceType() {
       return source().getType().cast<MemRefType>();
     }
 
     /// The result of a subview is always a memref.
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
 
     /// A subview result type can be fully inferred from the source type and the
     /// static representation of offsets, sizes and strides. Special sentinels
let arguments = (ins
     Variadic<Index>:$sizes,
     Variadic<Index>:$strides,
     I64ArrayAttr:$static_offsets,
     I64ArrayAttr:$static_sizes,
     I64ArrayAttr:$static_strides
   );
   let results = (outs AnyRankedTensor:$result);
 
+  let builders = [
+    // Build a SubViewOp with mized static and dynamic entries.
+    OpBuilder<
+      "Value source, ArrayRef<int64_t> staticOffsets, "
+      "ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides, "
+      "ValueRange offsets, ValueRange sizes, ValueRange strides, "
+      "ArrayRef<NamedAttribute> attrs = {}">,
+    // Build a SubViewOp with all dynamic entries.
+    OpBuilder<
+      "Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, "
+      "ArrayRef<NamedAttribute> attrs = {}">
+  ];
+
   let extraClassDeclaration = extraBaseClassDeclaration # [{
     /// Returns the type of the base tensor operand.
-    RankedTensorType getSourceRankedTensorType() {
+    RankedTensorType getSourceType() {
       return source().getType().cast<RankedTensorType>();
     }
 
     /// The result of a subtensor is always a tensor.
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
     /// A subview result type can be fully inferred from the source type and the
     /// static representation of offsets, sizes and strides. Special sentinels
     /// encode the dynamic case.
     static Type inferResultType(RankedTensorType sourceRankedTensorType,
                                 ArrayRef<int64_t> staticOffsets,
                                 ArrayRef<int64_t> staticSizes,
                                 ArrayRef<int64_t> staticStrides);
   }];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// SubTensorInsertOp
+//===----------------------------------------------------------------------===//
+
+def SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides<"subtensor_insert"> {
+  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`.
+
+    The subtensor_insert operation supports the following arguments:
+
+    * source: the tensor that is inserted.
+    * dest: the tensor into which the source tensor is inserted.
+    * offsets: tensor-rank number of dynamic offsets or static integer
+               attributes into the "base" tensor from which to extract the
+               subtensor.
+    * sizes: tensor-rank number of dynamic sizes or static integer attributes
+             which specify the sizes of the result tensor type.
+    * strides: tensor-rank number of dynamic strides or static integer
+               attributes multiplicatively to the base memref strides in each
+               dimension.
+
+    After buffer-allocation, the "subtensor_insert" op is expected to ne erased
+    and be replaced by an in-place buffer update.
+  }];
+
+  let arguments = (ins
+    AnyRankedTensor:$source,
+    AnyRankedTensor:$dest,
+    Variadic<Index>:$offsets,
+    Variadic<Index>:$sizes,
+    Variadic<Index>:$strides,
+    I64ArrayAttr:$static_offsets,
+    I64ArrayAttr:$static_sizes,
+    I64ArrayAttr:$static_strides
+  );
+  let results = (outs AnyRankedTensor:$result);
+
+  let builders = [
+    // Build a SubViewOp with mized static and dynamic entries.
+    OpBuilder<
+      "Value source, Value dest, ArrayRef<int64_t> staticOffsets, "
+      "ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides, "
+      "ValueRange offsets, ValueRange sizes, ValueRange strides, "
+      "ArrayRef<NamedAttribute> attrs = {}">,
+    // Build a SubViewOp with all dynamic entries.
+    OpBuilder<
+      "Value source, Value dest, ValueRange offsets, ValueRange sizes, "
+      "ValueRange strides, 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>();
+    }
+  }];
+
+  // let hasCanonicalizer = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // TanhOp
 //===----------------------------------------------------------------------===//
 
 def TanhOp : FloatUnaryOp<"tanh"> {
   let summary = "hyperbolic tangent of the specified value";
   let description = [{
     Syntax:
 

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/IR/OpsEnums.cpp.inc"
 
 using namespace mlir;

 }
 
 /// Print SubViewOp in the form:
 /// ```
 ///   subview ssa-name `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
 ///     `:` strided-memref-type `to` strided-memref-type
 /// ```
 template <typename OpType>
-static void printOpWithOffsetsSizesAndStrides(OpAsmPrinter &p, OpType op) {
+static void printOpWithOffsetsSizesAndStrides(
+    OpAsmPrinter &p, OpType op,
+    llvm::function_ref<void(OpAsmPrinter &p, OpType op)> printExtraOperands =
+        [](OpAsmPrinter &p, OpType op) {},
+    StringLiteral resultTypeKeyword = "to") {
   int stdDotLen = StandardOpsDialect::getDialectNamespace().size() + 1;
   p << op.getOperation()->getName().getStringRef().drop_front(stdDotLen) << ' ';
-  p << op.getOperand(0);
+  p << op.source();
+  printExtraOperands(p, op);
   printSubViewListOfOperandsOrIntegers(p, op.offsets(), op.static_offsets(),
                                        ShapedType::isDynamicStrideOrOffset);
   printSubViewListOfOperandsOrIntegers(p, op.sizes(), op.static_sizes(),
                                        ShapedType::isDynamic);
   printSubViewListOfOperandsOrIntegers(p, op.strides(), op.static_strides(),
                                        ShapedType::isDynamicStrideOrOffset);
   p.printOptionalAttrDict(op.getAttrs(),
                           /*elidedAttrs=*/{OpType::getSpecialAttrNames()});
-  p << " : " << op.getOperand(0).getType() << " to " << op.getType();
+  p << " : " << op.getSourceType() << " " << resultTypeKeyword << " "
+    << op.getType();
 }
 
 static void print(OpAsmPrinter &p, SubViewOp op) {
   return printOpWithOffsetsSizesAndStrides<SubViewOp>(p, op);
 }
 
-/// Parse SubViewOp of the form:
+/// Parse of the form:
 /// ```
-///   `name` ssa-name `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
-///     `:` strided-memref-type `to` strided-memref-type
+///   `name` ssa-name (extra-operands)?
+///       `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
+///     `:` strided-memref-type `resultTypeKeyword strided-memref-type
 /// ```
 template <typename OpType>
-static ParseResult parseOpWithOffsetsSizesAndStrides(OpAsmParser &parser,
-                                                     OperationState &result) {
-  OpAsmParser::OperandType srcInfo;
+static ParseResult parseOpWithOffsetsSizesAndStrides(
+    OpAsmParser &parser, OperationState &result,
+    std::function<ParseResult(OpAsmParser &p,
+                              OpAsmParser::OperandType &dstInfo)>
+        parseExtraOperand = nullptr,
+    StringLiteral resultTypeKeyword = "to") {
+  OpAsmParser::OperandType srcInfo, dstInfo;
   SmallVector<OpAsmParser::OperandType, 4> offsetsInfo, sizesInfo, stridesInfo;
   auto indexType = parser.getBuilder().getIndexType();
   Type srcType, dstType;
   if (parser.parseOperand(srcInfo))
     return failure();
+  if (parseExtraOperand && parseExtraOperand(parser, dstInfo))
+    return failure();
   if (parseListOfOperandsOrIntegers(
           parser, result, OpType::getStaticOffsetsAttrName(),
           ShapedType::kDynamicStrideOrOffset, offsetsInfo) ||
       parseListOfOperandsOrIntegers(parser, result,
                                     OpType::getStaticSizesAttrName(),
                                     ShapedType::kDynamicSize, sizesInfo) ||
       parseListOfOperandsOrIntegers(
           parser, result, OpType::getStaticStridesAttrName(),
           ShapedType::kDynamicStrideOrOffset, stridesInfo))
     return failure();
 
+  // Handle segment sizes.
   auto b = parser.getBuilder();
-  SmallVector<int, 4> segmentSizes{1, static_cast<int>(offsetsInfo.size()),
+  SmallVector<int, 4> segmentSizes = {1, static_cast<int>(offsetsInfo.size()),
                                       static_cast<int>(sizesInfo.size()),
                                       static_cast<int>(stridesInfo.size())};
+  // If we parse an extra operand it needs to appear in the segmentSizes
+  if (parseExtraOperand)
+    segmentSizes.insert(segmentSizes.begin(), 1);
   result.addAttribute(OpType::getOperandSegmentSizeAttr(),
                       b.getI32VectorAttr(segmentSizes));
 
   return failure(
       parser.parseOptionalAttrDict(result.attributes) ||
       parser.parseColonType(srcType) ||
+      parser.parseKeywordType(resultTypeKeyword.str().c_str(), dstType) ||
       parser.resolveOperand(srcInfo, srcType, result.operands) ||
+      (parseExtraOperand &&
+       parser.resolveOperand(dstInfo, dstType, result.operands)) ||
       parser.resolveOperands(offsetsInfo, indexType, result.operands) ||
       parser.resolveOperands(sizesInfo, indexType, result.operands) ||
       parser.resolveOperands(stridesInfo, indexType, result.operands) ||
-      parser.parseKeywordType("to", dstType) ||
       parser.addTypeToList(dstType, result.types));
 }
 
 static ParseResult parseSubViewOp(OpAsmParser &parser, OperationState &result) {
   return parseOpWithOffsetsSizesAndStrides<SubViewOp>(parser, result);
 }
 
 void mlir::SubViewOp::build(OpBuilder &b, OperationState &result, Value source,
if (failed(verifyOpWithOffsetSizesAndStridesPart(
           op, "stride", op.getStaticStridesAttrName(), op.static_strides(),
           ShapedType::isDynamicStrideOrOffset, op.strides())))
     return failure();
   return success();
 }
 
 /// Verifier for SubViewOp.
 static LogicalResult verify(SubViewOp op) {
-  MemRefType baseType = op.getSourceMemRefType();
+  MemRefType baseType = op.getSourceType();
   MemRefType subViewType = op.getType();
 
   // The base memref and the view memref should be in the same memory space.
   if (baseType.getMemorySpace() != subViewType.getMemorySpace())
     return op.emitError("different memory spaces specified for base memref "
                         "type ")
            << baseType << " and subview memref type " << subViewType;
 

 
 /// Verifier for SubViewOp.
 static LogicalResult verify(SubTensorOp op) {
   if (failed(verifyOpWithOffsetSizesAndStrides(op)))
     return failure();
 
   // Verify result type against inferred type.
   auto expectedType = SubTensorOp::inferResultType(
-      op.getSourceRankedTensorType(),
-      extractFromI64ArrayAttr(op.static_offsets()),
+      op.getSourceType(), extractFromI64ArrayAttr(op.static_offsets()),
       extractFromI64ArrayAttr(op.static_sizes()),
       extractFromI64ArrayAttr(op.static_strides()));
   if (op.getType() != expectedType)
     return op.emitError("expected result type to be ") << expectedType;
 
   return success();
 }
 
 void SubTensorOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                               MLIRContext *context) {
   results
       .insert<OpWithOffsetSizesAndStridesConstantArgumentFolder<SubTensorOp>>(
           context);
 }
 
 //===----------------------------------------------------------------------===//
+// SubTensorInsertOp
+//===----------------------------------------------------------------------===//
+
+static void print(OpAsmPrinter &p, SubTensorInsertOp op) {
+  return printOpWithOffsetsSizesAndStrides<SubTensorInsertOp>(
+      p, op,
+      [](OpAsmPrinter &p, SubTensorInsertOp op) { p << " into " << op.dest(); },
+      /*resultTypeKeyword=*/"into");
+}
+
+static ParseResult parseSubTensorInsertOp(OpAsmParser &parser,
+                                          OperationState &result) {
+  return parseOpWithOffsetsSizesAndStrides<SubTensorInsertOp>(
+      parser, result,
+      [](OpAsmParser &parser, OpAsmParser::OperandType &dstInfo) {
+        return failure(parser.parseKeyword("into") ||
+                       parser.parseOperand(dstInfo));
+      },
+      "into");
+}
+
+void mlir::SubTensorInsertOp::build(
+    OpBuilder &b, OperationState &result, Value source, Value dest,
+    ArrayRef<int64_t> staticOffsets, ArrayRef<int64_t> staticSizes,
+    ArrayRef<int64_t> staticStrides, ValueRange offsets, ValueRange sizes,
+    ValueRange strides, ArrayRef<NamedAttribute> attrs) {
+  build(b, result, dest.getType(), source, dest, offsets, sizes, strides,
+        b.getI64ArrayAttr(staticOffsets), b.getI64ArrayAttr(staticSizes),
+        b.getI64ArrayAttr(staticStrides));
+  result.addAttributes(attrs);
+}
+
+/// Build a SubViewOp with all dynamic entries: `staticOffsets`, `staticSizes`
+/// and `staticStrides` are  automatically filled with source-memref-rank
+/// sentinel values that encode dynamic entries.
+void mlir::SubTensorInsertOp::build(OpBuilder &b, OperationState &result,
+                                    Value source, Value dest,
+                                    ValueRange offsets, ValueRange sizes,
+                                    ValueRange strides,
+                                    ArrayRef<NamedAttribute> attrs) {
+  auto sourceRankedTensorType = source.getType().cast<RankedTensorType>();
+  unsigned rank = sourceRankedTensorType.getRank();
+  SmallVector<int64_t, 4> staticOffsetsVector;
+  staticOffsetsVector.assign(rank, ShapedType::kDynamicStrideOrOffset);
+  SmallVector<int64_t, 4> staticSizesVector;
+  staticSizesVector.assign(rank, ShapedType::kDynamicSize);
+  SmallVector<int64_t, 4> staticStridesVector;
+  staticStridesVector.assign(rank, ShapedType::kDynamicStrideOrOffset);
+  build(b, result, source, dest, staticOffsetsVector, staticSizesVector,
+        staticStridesVector, offsets, sizes, strides, attrs);
+}
+
+SmallVector<Range, 8> SubTensorInsertOp::getOrCreateRanges(OpBuilder &b,
+                                                           Location loc) {
+  return ::getOrCreateRangesImpl(*this, b, loc);
+}
+
+/// Verifier for SubViewOp.
+static LogicalResult verify(SubTensorInsertOp op) {
+  if (failed(verifyOpWithOffsetSizesAndStrides(op)))
+    return failure();
+  if (op.getType() != op.dest().getType())
+    return op.emitError("expected result type to be ") << op.dest().getType();
+  return success();
+}
+
+// void SubTensorInsertOp::getCanonicalizationPatterns(
+//     OwningRewritePatternList &results, MLIRContext *context) {
+//   results.insert<
+//       OpWithOffsetSizesAndStridesConstantArgumentFolder<SubTensorInsertOp>>(
+//       context);
+// }
+
+//===----------------------------------------------------------------------===//
 // TensorCastOp
 //===----------------------------------------------------------------------===//
 
 bool TensorCastOp::areCastCompatible(Type a, Type b) {
   auto aT = a.dyn_cast<TensorType>();
   auto bT = b.dyn_cast<TensorType>();
   if (!aT || !bT)
     return false;

mlir/test/IR/core-ops.mlir

 // CHECK-LABEL: func @assume_alignment
 // CHECK-SAME: %[[MEMREF:.*]]: memref<4x4xf16>
 func @assume_alignment(%0: memref<4x4xf16>) {
   // CHECK: assume_alignment %[[MEMREF]], 16 : memref<4x4xf16>
   assume_alignment %0, 16 : memref<4x4xf16>
   return
 }
 
-
 // CHECK-LABEL: func @subtensor({{.*}}) {
 func @subtensor(%t: tensor<8x16x4xf32>, %idx : index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
 
   // CHECK: subtensor
   // CHECK-SAME: tensor<8x16x4xf32> to tensor<?x?x?xf32>
   %1 = subtensor %t[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1]
     : tensor<8x16x4xf32> to tensor<?x?x?xf32>
 
   // CHECK: subtensor
   // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4x4xf32>
   %2 = subtensor %t[0, 2, 0][4, 4, 4][1, 1, 1]
     : tensor<8x16x4xf32> to tensor<4x4x4xf32>
 
   return
 }
+
+// CHECK-LABEL: func @subtensor_insert({{.*}}) {
+func @subtensor_insert(%t: tensor<8x16x4xf32>, %t2: tensor<16x32x8xf32>, %idx : index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+
+  // CHECK: subtensor_insert
+  // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32>
+  %1 = subtensor_insert %t into %t2[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1]
+    : tensor<8x16x4xf32> into tensor<16x32x8xf32>
+
+  return
+}