[mlir] Add some more std bufferize patterns.

[silvas]

Summary

Add bufferizations for extract_element and tensor_from_elements.

[herhut]

Thanks!

[silvas]

update

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

 #include "PassDetail.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
+class BufferizeExtractElementOp : public OpConversionPattern<ExtractElementOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ExtractElementOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    ExtractElementOp::Adaptor adaptor(operands);
+    rewriter.replaceOpWithNewOp<LoadOp>(op, adaptor.aggregate(),
+                                        adaptor.indices());
+    return success();
+  }
+};
+} // namespace
+
+namespace {
 class BufferizeTensorCastOp : public OpConversionPattern<TensorCastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(TensorCastOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto resultType = getTypeConverter()->convertType(op.getType());
     rewriter.replaceOpWithNewOp<MemRefCastOp>(op, resultType, operands[0]);
     return success();
   }
 };
 } // namespace
 
+namespace {
+class BufferizeTensorFromElementsOp
+    : public OpConversionPattern<TensorFromElementsOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(TensorFromElementsOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    int numberOfElements = op.elements().size();
+    auto resultType = MemRefType::get(
+        {numberOfElements}, op.getType().cast<TensorType>().getElementType());
+    Value result = rewriter.create<AllocOp>(op.getLoc(), resultType);
+    for (auto element : llvm::enumerate(op.elements())) {
+      Value index =
+          rewriter.create<ConstantIndexOp>(op.getLoc(), element.index());
+      rewriter.create<StoreOp>(op.getLoc(), element.value(), result, index);
+    }
+    rewriter.replaceOp(op, {result});
+    return success();
+  }
+};
+} // namespace
+
 void mlir::populateStdBufferizePatterns(MLIRContext *context,
                                         BufferizeTypeConverter &typeConverter,
                                         OwningRewritePatternList &patterns) {
-  patterns.insert<BufferizeTensorCastOp>(typeConverter, context);
+  patterns.insert<BufferizeExtractElementOp, BufferizeTensorCastOp,
+                  BufferizeTensorFromElementsOp>(typeConverter, context);
 }
 
 namespace {
 struct StdBufferizePass : public StdBufferizeBase<StdBufferizePass> {
   void runOnFunction() override {
     auto *context = &getContext();
     BufferizeTypeConverter typeConverter;
     OwningRewritePatternList patterns;
     ConversionTarget target(*context);
 
     target.addLegalDialect<StandardOpsDialect>();
 
     populateStdBufferizePatterns(context, typeConverter, patterns);
-    target.addIllegalOp<TensorCastOp>();
+    target.addIllegalOp<ExtractElementOp, TensorCastOp, TensorFromElementsOp>();
 
-    if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+    if (failed(applyPartialConversion(getFunction(), target, patterns)))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createStdBufferizePass() {
   return std::make_unique<StdBufferizePass>();
 }

mlir/test/Dialect/Standard/bufferize.mlir

 // RUN: mlir-opt %s -std-bufferize | FileCheck %s
 
+// CHECK-LABEL:   func @extract_element(
+// CHECK-SAME:                          %[[TENSOR:.*]]: tensor<?xf32>,
+// CHECK-SAME:                          %[[IDX:.*]]: index) -> f32 {
+// CHECK:           %[[MEMREF:.*]] = tensor_to_memref %[[TENSOR]] : memref<?xf32>
+// CHECK:           %[[RET:.*]] = load %[[MEMREF]][%[[IDX]]] : memref<?xf32>
+// CHECK:           return %[[RET]] : f32
+// CHECK:         }
+func @extract_element(%arg0: tensor<?xf32>, %arg1: index) -> f32 {
+  %0 = extract_element %arg0[%arg1] : tensor<?xf32>
+  return %0 : f32
+}
+
 // CHECK-LABEL:   func @tensor_cast(
 // CHECK-SAME:                      %[[TENSOR:.*]]: tensor<?xindex>) -> tensor<2xindex> {
 // CHECK:           %[[MEMREF:.*]] = tensor_to_memref %[[TENSOR]]
 // CHECK:           %[[CASTED:.*]] = memref_cast %[[MEMREF]] : memref<?xindex> to memref<2xindex>
 // CHECK:           %[[RET:.*]] = tensor_load %[[CASTED]]
 // CHECK:           return %[[RET]] : tensor<2xindex>
 func @tensor_cast(%arg0: tensor<?xindex>) -> tensor<2xindex> {
   %0 = tensor_cast %arg0 : tensor<?xindex> to tensor<2xindex>
   return %0 : tensor<2xindex>
 }
+
+// CHECK-LABEL:   func @tensor_from_elements(
+// CHECK-SAME:                               %[[ELEM0:.*]]: index,
+// CHECK-SAME:                               %[[ELEM1:.*]]: index) -> tensor<2xindex> {
+// CHECK:           %[[MEMREF:.*]] = alloc()
+// CHECK:           %[[C0:.*]] = constant 0 : index
+// CHECK:           store %[[ELEM0]], %[[MEMREF]][%[[C0]]]
+// CHECK:           %[[C1:.*]] = constant 1 : index
+// CHECK:           store %[[ELEM1]], %[[MEMREF]][%[[C1]]]
+// CHECK:           %[[RET:.*]] = tensor_load %[[MEMREF]]
+// CHECK:           return %[[RET]] : tensor<2xindex>
+func @tensor_from_elements(%arg0: index, %arg1: index) -> tensor<2xindex> {
+  %0 = tensor_from_elements %arg0, %arg1 : tensor<2xindex>
+  return %0 : tensor<2xindex>
+}