[mlir] Add canonicaliations for subtensor_insert operation.

[mravishankar]

Summary

Add canonicalizers to subtensor_insert operations need canonicalizers
that propagate the constant arguments within offsets, sizes and
strides. Also add pattern to propogate tensor_cast operations.

[rriddle]

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

3310

Why do you need this method?

[mravishankar]

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

3310

It is used in SubTensorInsertOpCastFolder below. When you output type changes, this adds the cast so that replacement works fine. The cast is canonicalized away later.

[rriddle]

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

3310

Why do you need a new function for that? That is what I meant. This looks like something that would just be inlined into the pattern.

[mravishankar]

Address comments

[mravishankar]

Remove unneeded method.

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

3310

Fair enough. I just copied what was done for subtensor and subview above. But that was done since those two ops use the same template method. Removed this.

[mravishankar]

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

3310

At least I thought I did. Let me reupload.

[mravishankar]

Reupload

[rriddle]

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

3805

copy/paste? : subview

3813

3833–3836

3841

nit: Drop trivial braces.

[mravishankar]

Rebase and fix stale comments.

[mravishankar]

Addressing missed comments.

[mravishankar]

(push)

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

std::array<unsigned, 3> getArrayAttrMaxRanks() {
       return {rank, rank, rank};
     }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; }
   }];
 
+  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TensorLoadOp
 //===----------------------------------------------------------------------===//
 
 def TensorLoadOp : Std_Op<"tensor_load",

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

static void replaceWithNewOp(PatternRewriter &rewriter, SubViewOp op,
   rewriter.replaceOpWithNewOp<MemRefCastOp>(op, newOp, op.getType());
 }
 
 static void replaceWithNewOp(PatternRewriter &rewriter, SubTensorOp op,
                              SubTensorOp newOp) {
   rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
 }
 
+static void replaceWithNewOp(PatternRewriter &rewriter, SubTensorInsertOp op,
+                             SubTensorInsertOp newOp) {
+  rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
+}
+
 /// Pattern to rewrite a subview op with constant arguments.
 template <typename OpType>
 class OpWithOffsetSizesAndStridesConstantArgumentFolder final
     : public OpRewritePattern<OpType> {
 public:
   using OpRewritePattern<OpType>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpType op,
OpFoldResult SubTensorInsertOp::fold(ArrayRef<Attribute>) {
   if (getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->source();
   if (succeeded(tensor::foldTensorCast(*this)))
     return this->source();
   return OpFoldResult();
 }
 
+namespace {
+
+/// Pattern to rewrite a subview op with constant arguments.
+class SubTensorInsertOpConstantArgumentFolder final
+    : public OpRewritePattern<SubTensorInsertOp> {
+public:
+  using OpRewritePattern<SubTensorInsertOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SubTensorInsertOp subTensorInsertOp,
+                                PatternRewriter &rewriter) const override {
+    // No constant operand, just return;
+    if (llvm::none_of(subTensorInsertOp.getOperands(), [](Value operand) {
+          return matchPattern(operand, m_ConstantIndex());
+        }))
+      return failure();
+
+    // At least one of offsets/sizes/strides is a new constant.
+    // Form the new list of operands and constant attributes from the existing.
+    SmallVector<OpFoldResult> mixedOffsets(subTensorInsertOp.getMixedOffsets());
+    SmallVector<OpFoldResult> mixedSizes(subTensorInsertOp.getMixedSizes());
+    SmallVector<OpFoldResult> mixedStrides(subTensorInsertOp.getMixedStrides());
+    canonicalizeSubViewPart(mixedOffsets, ShapedType::isDynamicStrideOrOffset);
+    canonicalizeSubViewPart(mixedSizes, ShapedType::isDynamic);
+    canonicalizeSubViewPart(mixedStrides, ShapedType::isDynamicStrideOrOffset);
+
+    // Create the new op in canonical form.
+    Value source = subTensorInsertOp.source();
+    RankedTensorType sourceType = source.getType().cast<RankedTensorType>();
+    SmallVector<int64_t, 4> shape = llvm::to_vector<4>(
+        llvm::map_range(mixedSizes, [](OpFoldResult valueOrAttr) -> int64_t {
+          auto attr = valueOrAttr.dyn_cast<Attribute>();
+          if (!attr)
+            return ShapedType::kDynamicSize;
+          return attr.cast<IntegerAttr>().getInt();
+        }));
+    RankedTensorType newSourceType =
+        RankedTensorType::get(shape, sourceType.getElementType());
+    Location loc = subTensorInsertOp.getLoc();
+    if (sourceType != newSourceType) {
+      source = rewriter.create<tensor::CastOp>(loc, newSourceType, source);
+    }
+    rewriter.replaceOpWithNewOp<SubTensorInsertOp>(
+        subTensorInsertOp, source, subTensorInsertOp.dest(), mixedOffsets,
+        mixedSizes, mixedStrides);
+    return success();
+  }
+};
+
+/// Fold tensor_casts with subtensor_insert operations.
+struct SubTensorInsertOpCastFolder final
+    : public OpRewritePattern<SubTensorInsertOp> {
+  using OpRewritePattern<SubTensorInsertOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SubTensorInsertOp subTensorOp,
+                                PatternRewriter &rewriter) const override {
+    if (llvm::any_of(subTensorOp.getOperands(), [](Value operand) {
+          return matchPattern(operand, m_ConstantIndex());
+        }))
+      return failure();
+
+    auto getSourceOfCastOp = [](Value v) -> Optional<Value> {
+      auto castOp = v.getDefiningOp<tensor::CastOp>();
+      if (!castOp || !canFoldIntoConsumerOp(castOp))
+        return llvm::None;
+      return castOp.source();
+    };
+    Optional<Value> sourceCastSource = getSourceOfCastOp(subTensorOp.source());
+    Optional<Value> destCastSource = getSourceOfCastOp(subTensorOp.dest());
+    if (!sourceCastSource && !destCastSource &&
+        subTensorOp.dest().getType() == subTensorOp.getResult().getType())
+      return failure();
+
+    auto newOp = rewriter.create<SubTensorInsertOp>(
+        subTensorOp.getLoc(),
+        (sourceCastSource ? *sourceCastSource : subTensorOp.source()),
+        (destCastSource ? *destCastSource : subTensorOp.dest()),
+        subTensorOp.getMixedOffsets(), subTensorOp.getMixedSizes(),
+        subTensorOp.getMixedStrides());
+
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(subTensorOp,
+                                                subTensorOp.getType(), newOp);
+    return success();
+  }
+};
+} // namespace
+
+void SubTensorInsertOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<SubTensorInsertOpConstantArgumentFolder,
+                 SubTensorInsertOpCastFolder>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // TensorLoadOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult TensorLoadOp::fold(ArrayRef<Attribute>) {
   if (auto tensorToMemref = memref().getDefiningOp<TensorToMemrefOp>())
     return tensorToMemref.tensor();
   return {};

mlir/test/Dialect/Standard/canonicalize.mlir

 // Tensor cast is folded away.
 //   CHECK-NOT:   tensor.cast
 //       CHECK:   return %[[S]] : tensor<4x6x16x32xi8>
 func @rank_reducing_subtensor_insert_of_cast(%a : tensor<16x32xi8>, %b : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {
   %cast = tensor.cast %a : tensor<16x32xi8> to tensor<?x32xi8>
   %res = subtensor_insert %cast into %b[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor<?x32xi8> into tensor<4x6x16x32xi8>
   return %res : tensor<4x6x16x32xi8>
 }
+
+// -----
+
+func @subtensor_canonicalize(%arg0 : tensor<2x?xi32>, %arg1 : tensor<i32>,
+    %arg2 : index, %arg3 : index) -> tensor<?x?xi32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c8 = constant 8 : index
+  %0 = dim %arg0, %c1 : tensor<2x?xi32>
+  %1 = tensor.extract %arg1[] : tensor<i32>
+  %2 = tensor.generate %arg2, %c8 {
+  ^bb0(%arg4: index, %arg5: index):
+    tensor.yield %1 : i32
+  } : tensor<?x?xi32>
+  %3 = subtensor_insert %arg0 into %2[%c0, %arg3] [%c2, %0] [%c1, %c1] : tensor<2x?xi32> into tensor<?x?xi32>
+  return %3 : tensor<?x?xi32>
+}
+// CHECK-LABEL: func @subtensor_canonicalize
+//       CHECK:   %[[UPDATED:.+]] = subtensor_insert %{{.+}} into %{{.+}}[0, %{{.+}}] [2, %{{.+}}] [1, 1]
+//  CHECK-SAME:     tensor<2x?xi32> into tensor<?x8xi32>
+//       CHECK:   %[[CAST:.+]] = tensor.cast %[[UPDATED]]
+//       CHECK:   return %[[CAST]]
+
+// -----
+
+func @subtensor_insert_output_dest_canonicalize(%arg0 : tensor<2x3xi32>, %arg1 : tensor<i32>) -> tensor<3x9xi32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c9 = constant 9 : index
+  %c3 = constant 3 : index
+  %2 = tensor.extract %arg1[] : tensor<i32>
+  %4 = tensor.generate %c3, %c9 {
+  ^bb0(%arg2: index, %arg3: index):
+    tensor.yield %2 : i32
+  } : tensor<?x?xi32>
+  %5 = subtensor_insert %arg0 into %4[%c0, %c1] [%c2, %c3] [1, 1] : tensor<2x3xi32> into tensor<?x?xi32>
+  %6 = tensor.cast %5 : tensor<?x?xi32> to tensor<3x9xi32>
+  return %6 : tensor<3x9xi32>
+}
+// CHECK-LABEL: func @subtensor_insert_output_dest_canonicalize
+//  CHECK-SAME:   %[[ARG0:[a-zA-z0-9_]+]]: tensor<2x3xi32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<i32>
+//       CHECK:   %[[PAD:.+]] = tensor.extract %[[ARG1]]
+//       CHECK:   %[[GENERATE:.+]] = tensor.generate
+//       CHECK:   %[[RESULT:.+]] = subtensor_insert %[[ARG0]] into %[[GENERATE]]
+//       CHECK:   return %[[RESULT]]