[mlir][tensor] Canonicalize scalar `tensor.insert_slice(tensor.insert, _)` to `tensor.insert`

[christopherbate]

Summary

Canonicalizes the pattern

%0 = tensor.insert %scalar into %t1[...] : (scalar tensor type)
%1 = tensor.insert_slice %0 into %t2[<indices>]

into

%1 = tensor.insert %scalar into %t2[<indices>]

This has a side effect on bufferization: prior to change canonicalization, the
IR below would result in two allocations (even with empty tensor elimination),
whereas afterwards it results in just the creation of two memref.store ops:

func.func @func(%arg0 : f32, %arg1: f32, %arg2: tensor<4xf32>) -> (tensor<4xf32>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c3 = arith.constant 3 : index
  %e1 = tensor.empty() : tensor<1xf32>
  %e2 = tensor.empty() : tensor<f32>
  %0 =  tensor.insert %arg0 into %e1[%c0] : tensor<1xf32>
  %1 =  tensor.insert %arg1 into %e2[] : tensor<f32>
  %2 = tensor.insert_slice %0 into %arg2[%c2][1][1] : tensor<1xf32> into tensor<4xf32>
  %3 = tensor.insert_slice %1 into %2[%c3][1][1] : tensor<f32> into tensor<4xf32>
  return %3 : tensor<4xf32>
}

[nicolasvasilache]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

2500

Seems to me you either need to check that the tensor.insert indices are all 0 or that you need to combine them together no ?

Additionally, have you looked at mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp ?
It seems you could add a new pattern there ?

Note that those patterns are not canonicalizations because they are not always desirable, so we made them opt-in.

[christopherbate]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

2500

The producer (tensor.insert) is producing a scalar. That's what line 2500 is checking. The consumer (tensor.insert_slice) is then inserting that scalar into some position. So I don't think we need to check the indices of tensor.insert? Or did I get something wrong here?

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Value cast = rewriter.create<tensor::CastOp>(
         insertSliceOp.getLoc(), newSrcType, insertSliceOp.getSource());
     rewriter.replaceOpWithNewOp<InsertOpTy>(
         insertSliceOp, cast, insertSliceOp.getDest(),
         insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
         insertSliceOp.getMixedStrides());
     return success();
   }
 };
+
+/// Canonicalizes the pattern
+/// ```
+/// %0 = tensor.insert %scalar into %t1[...] : (scalar tensor type)
+/// %1 = tensor.insert_slice %0 into %t2[<indices>]
+/// ```
+/// into
+/// ```
+/// %1 = tensor.insert %scalar into %t2[<indices>]
+/// ```
+struct InsertSliceToInsertRewriter : public OpRewritePattern<InsertSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::InsertSliceOp op,
+                                PatternRewriter &rewriter) const override {
+
+    RankedTensorType sourceType = op.getSourceType();
+    // The `tensor.insert` result (`insert_slice` source) should be a scalar
+    // so that we know forwarding makes sense here.
+    if (!sourceType.hasStaticShape() || sourceType.getNumElements() != 1)
+      return failure();
+
+    auto insertOp = op.getSource().getDefiningOp<tensor::InsertOp>();
+    if (!insertOp)
+      return failure();
+
+    SmallVector<Value> indices = mlir::getValueOrCreateConstantIndexOp(
+        rewriter, op.getLoc(), op.getMixedOffsets());
+    rewriter.replaceOpWithNewOp<tensor::InsertOp>(op, insertOp.getScalar(),
+                                                  op.getDest(), indices);
+    return success();
+  }
+};
+
 } // namespace
 
 llvm::SmallBitVector InsertSliceOp::getDroppedDims() {
   return ::getDroppedDims(getSourceType().getShape(), getMixedSizes());
 }
 
 void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<InsertSliceOpConstantArgumentFolder<InsertSliceOp>,
               InsertSliceOpCastFolder<InsertSliceOp>,
-              InsertSliceOpSourceCastInserter<InsertSliceOp>>(context);
+              InsertSliceOpSourceCastInserter<InsertSliceOp>,
+              InsertSliceToInsertRewriter>(context);
 }
 
 Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder &b,
                                                              Location loc,
                                                              Value tensor,
                                                              Value dest) {
   auto rankedTensorType = llvm::cast<RankedTensorType>(dest.getType());
   unsigned rank = rankedTensorType.getRank();

mlir/test/Dialect/Tensor/canonicalize.mlir

func.func @fold_insert(%arg0 : index) -> (tensor<4xf32>) {
   %1 = arith.constant 4.0 : f32
   %ins_1 = tensor.insert %1 into %0[%arg0] : tensor<4xf32>
   // CHECK-NEXT: return %[[C4]]
   return %ins_1 : tensor<4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @scalar_insert_slice_to_insert
+//  CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: tensor<4xf32>
+func.func @scalar_insert_slice_to_insert(%arg0 : f32, %arg1: f32, %arg2: tensor<4xf32>) -> (tensor<4xf32>) {
+  // Canonicalize an insert_slice into an insert.
+  // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %e1 = tensor.empty() : tensor<1xf32>
+  %e2 = tensor.empty() : tensor<f32>
+  %0 =  tensor.insert %arg0 into %e1[%c0] : tensor<1xf32>
+  %1 =  tensor.insert %arg1 into %e2[] : tensor<f32>
+  // CHECK: %[[INS:.+]] = tensor.insert %[[ARG0]] into %[[ARG2]][%[[C2]]]
+  // CHECK: %[[INS1:.+]] = tensor.insert %[[ARG1]] into %[[INS]][%[[C3]]]
+  // CHECK: return %[[INS1]]
+  %2 = tensor.insert_slice %0 into %arg2[%c2][1][1] : tensor<1xf32> into tensor<4xf32>
+  %3 = tensor.insert_slice %1 into %2[%c3][1][1] : tensor<f32> into tensor<4xf32>
+  return %3 : tensor<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @extract_from_tensor.cast
 // CHECK-SAME: %[[TENSOR:.*]]: tensor<9xf32>
 func.func @extract_from_tensor.cast(%tensor: tensor<9xf32>) -> f32 {
   // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
   %c0 = arith.constant 0 : index
   // CHECK-NOT: tensor.cast
   %casted = tensor.cast %tensor : tensor<9xf32> to tensor<?xf32>
   // CHECK-NEXT: tensor.extract %[[TENSOR]][%[[C0]]]