Add lowering of tensor.insert op.

[pashu123]

Summary

The lowering of tensor.insert op has been added to the
tensor-bufferize pass.

Signed-off-by: Prashant Kumar <[e-mail]>

[silvas]

mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp

67

due to the way these work, you need to create a copy of the bufferized input memref before storing to it.

67

and yes, that's thoroughly non-optimal... there's ongoing discussions re bufferization.

[pashu123]

mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp

67

Is there any LLVM-discourse link for the discussion?

mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp

struct BufferizeExtractOp : public OpConversionPattern<tensor::ExtractOp> {
   matchAndRewrite(tensor::ExtractOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<memref::LoadOp>(op, adaptor.tensor(),
                                                 adaptor.indices());
     return success();
   }
 };
 
+struct BufferizeInsertOp : public OpConversionPattern<tensor::InsertOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(tensor::InsertOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    rewriter.create<memref::StoreOp>(op.getLoc(), adaptor.scalar(),
+                                     adaptor.dest(), adaptor.indices());
+    rewriter.replaceOp(op, {adaptor.dest()});
+    return success();
+  }
+};
+
 struct BufferizeFromElementsOp
     : public OpConversionPattern<tensor::FromElementsOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::FromElementsOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
void runOnFunction() override {
     bufferization::BufferizeTypeConverter typeConverter;
 
     ConversionTarget target(*context);
     target.addLegalDialect<scf::SCFDialect, memref::MemRefDialect>();
     target.addDynamicallyLegalDialect<arith::ArithmeticDialect,
                                       StandardOpsDialect>(
         [&](Operation *op) { return typeConverter.isLegal(op); });
     target.addLegalOp<CallOp, ReturnOp>();
-    target.addIllegalOp<tensor::CastOp, tensor::ExtractOp,
+    target.addIllegalOp<tensor::CastOp, tensor::ExtractOp, tensor::InsertOp,
                         tensor::FromElementsOp, tensor::GenerateOp>();
     bufferization::populateBufferizeMaterializationLegality(target);
 
     RewritePatternSet patterns(context);
     populateTensorBufferizePatterns(typeConverter, patterns);
     if (failed(
             applyPartialConversion(getFunction(), target, std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 void mlir::populateTensorBufferizePatterns(
     bufferization::BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns) {
   patterns.add<BufferizeCastOp, BufferizeDimOp, BufferizeExtractOp,
-               BufferizeFromElementsOp, BufferizeGenerateOp, BufferizeRankOp>(
-      typeConverter, patterns.getContext());
+               BufferizeInsertOp, BufferizeFromElementsOp, BufferizeGenerateOp,
+               BufferizeRankOp>(typeConverter, patterns.getContext());
 }
 
 std::unique_ptr<Pass> mlir::createTensorBufferizePass() {
   return std::make_unique<TensorBufferizePass>();
 }

mlir/test/Dialect/Tensor/bufferize.mlir

func @tensor.generate_unknown_ops_in_body(%arg0: index) -> tensor<?xindex> {
   %tensor = tensor.generate %arg0 {
   ^bb0(%iv: index):
     // CHECK: test.source
     %0 = "test.source"() : () -> index
     tensor.yield %0 : index
   } : tensor<?xindex>
   return %tensor : tensor<?xindex>
 }
+
+// CHECK-LABEL:  func @tensor.insert(
+// CHECK-SAME:              %[[TENSOR:.*]]: tensor<?xf32>,
+// CHECK-SAME:              %[[ELEM_VAL:.*]]: f32,
+// CHECK-SAME:              %[[IDX:.*]]: index) -> tensor<?xf32> {
+// CHECK:       %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<?xf32>
+// CHECK:       memref.store %[[ELEM_VAL]], %[[MEMREF]][%[[IDX]]] : memref<?xf32>
+// CHECK:       %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<?xf32>
+// CHECK:       return %[[RET]] : tensor<?xf32>
+// CHECK:  }
+func @tensor.insert(%arg0: tensor<?xf32>, %arg1: f32, %arg2: index) -> tensor<?xf32> {
+  %0 = tensor.insert %arg1 into %arg0[%arg2] : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}