[mlir] Add std.dynamic_tensor_from_elements bufferization.

[silvas]

Summary

It's unfortunate that this requires adding a dependency on scf dialect
to std bufferization (and hence all of std transforms). This is a bit
perilous. We might want a lib/Transforms/Bufferize/ with a separate
bufferization library per dialect?

[herhut]

Regarding the dependency: Is bufferize special enough to require its own directory or is this just a symptom of the too coarse modelling of transform dependencies in general?

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

58

The DynamicTensorFromElementsOp does not have a guaranteed iteration order, it is conceptually a map over the index space.

[silvas]

Regarding the dependency: Is bufferize special enough to require its own directory or is this just a symptom of the too coarse modelling of transform dependencies in general?

Good point. I think it's the too coarse modeling of dialect transforms.

[silvas]

update

[silvas]

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

58

Updated the op description as well.

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

def DynamicTensorFromElementsOp : Std_Op<"dynamic_tensor_from_elements",
     [RecursiveSideEffects, SingleBlockImplicitTerminator<"YieldOp">]> {
   string summary = "Creates a dynamically sized tensor from elements";
   string description = [{
     This operation creates a dynamically sized tensor with elements of any type.
     It expects one index operand per dynamic extent of the result tensor.
 
     The body region defines the tensor's elements. It takes index operands as
     its region arguments that span the index space. The element at the given
-    position is yielded with the `yield` operation (see `YieldOp`).
+    position is yielded with the `yield` operation (see `YieldOp`). There is
+    no defined ordering to the invocations of the body. It is conceptually
+    a "parallel map" operation.
 
     Example:
 
     ```mlir
       %tnsr = dynamic_tensor_from_elements %m, %n {
       ^bb0(%i : index, %j : index, %k : index):
         ...
         yield %elem : f32

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

 def ExpandAtomic : FunctionPass<"expand-atomic"> {
   let summary = "Expands AtomicRMWOp into GenericAtomicRMWOp.";
   let constructor = "mlir::createExpandAtomicPass()";
 }
 
 def StdBufferize : FunctionPass<"std-bufferize"> {
   let summary = "Bufferize the std dialect";
   let constructor = "mlir::createStdBufferizePass()";
+  let dependentDialects = ["scf::SCFDialect"];
 }
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

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

 //===- Bufferize.cpp - Bufferization for std ops --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements bufferization of std ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/Bufferize.h"
 #include "PassDetail.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
+class BufferizeDynamicTensorFromElementsOp
+    : public OpConversionPattern<DynamicTensorFromElementsOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(DynamicTensorFromElementsOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    // Allocate memory.
+    Location loc = op.getLoc();
+    DynamicTensorFromElementsOp::Adaptor transformed(operands);
+    RankedTensorType tensorType = op.getType().cast<RankedTensorType>();
+    MemRefType memrefType =
+        MemRefType::get(tensorType.getShape(), tensorType.getElementType());
+    Value result =
+        rewriter.create<AllocOp>(loc, memrefType, transformed.dynamicExtents());
+
+    // Collect loop bounds.
+    int64_t rank = tensorType.getRank();
+    Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+    Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+    SmallVector<Value, 4> lowerBounds(rank, zero);
+    SmallVector<Value, 4> steps(rank, one);
+    SmallVector<Value, 4> upperBounds;
+    int nextDynamicIndex = 0;
+    for (int i = 0; i < rank; i++) {
+      Value upperBound =
+          tensorType.isDynamicDim(i)
+              ? transformed.dynamicExtents()[nextDynamicIndex++]
+              : rewriter.create<ConstantIndexOp>(loc, memrefType.getDimSize(i));
+      upperBounds.push_back(upperBound);
+    }
+
+    // Generate tensor elements with a parallel loop.
+    rewriter.create<scf::ParallelOp>(
+        loc, lowerBounds, upperBounds, steps,
+        [&](OpBuilder &b, Location loc, ValueRange ivs) {
+          BlockAndValueMapping mapping;
+          mapping.map(op.body().getArguments(), ivs);
+          for (auto &nestedOp : op.getBody()->without_terminator())
+            b.clone(nestedOp, mapping);
+          auto yieldOp = cast<YieldOp>(op.getBody()->getTerminator());
+          b.create<StoreOp>(loc, mapping.lookup(yieldOp.value()), result, ivs);
+          b.create<scf::YieldOp>(loc);
+        });
+
+    rewriter.replaceOp(op, {result});
+    return success();
+  }
+};
+} // namespace
+
+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(),
matchAndRewrite(TensorFromElementsOp op, ArrayRef<Value> operands,
     return success();
   }
 };
 } // namespace
 
 void mlir::populateStdBufferizePatterns(MLIRContext *context,
                                         BufferizeTypeConverter &typeConverter,
                                         OwningRewritePatternList &patterns) {
-  patterns.insert<BufferizeExtractElementOp, BufferizeTensorCastOp,
-                  BufferizeTensorFromElementsOp>(typeConverter, context);
+  patterns
+      .insert<BufferizeDynamicTensorFromElementsOp, 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>();
+    target.addLegalDialect<scf::SCFDialect>();
 
     populateStdBufferizePatterns(context, typeConverter, patterns);
-    target.addIllegalOp<ExtractElementOp, TensorCastOp, TensorFromElementsOp>();
+    target.addIllegalOp<DynamicTensorFromElementsOp, ExtractElementOp,
+                        TensorCastOp, TensorFromElementsOp>();
 
     if (failed(applyPartialConversion(getFunction(), target, patterns)))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createStdBufferizePass() {
   return std::make_unique<StdBufferizePass>();
 }

mlir/lib/Dialect/StandardOps/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
   ExpandAtomic.cpp
   ExpandTanh.cpp
   FuncConversions.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS
   MLIRStandardTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
+  MLIRSCF
   MLIRStandard
   MLIRTransforms
   )

mlir/lib/Dialect/StandardOps/Transforms/PassDetail.h

 //===- PassDetail.h - GPU Pass class details --------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_STANDARD_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_STANDARD_TRANSFORMS_PASSDETAIL_H_
 
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class AtomicRMWOp;
 
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // DIALECT_STANDARD_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/Standard/bufferize.mlir

 // RUN: mlir-opt %s -std-bufferize | FileCheck %s
 
+// CHECK-LABEL:   func @dynamic_tensor_from_elements(
+// CHECK-SAME:                                       %[[ARG:.*]]: tensor<*xf32>,
+// CHECK-SAME:                                       %[[DYNAMIC_EXTENT:.*]]: index) -> tensor<?xindex> {
+// CHECK:           %[[MEMREF:.*]] = alloc(%[[DYNAMIC_EXTENT]]) : memref<?xindex>
+// CHECK:           %[[C0:.*]] = constant 0 : index
+// CHECK:           %[[C1:.*]] = constant 1 : index
+// CHECK:           scf.parallel (%[[I:.*]]) = (%[[C0]]) to (%[[DYNAMIC_EXTENT]]) step (%[[C1]]) {
+// CHECK:             %[[ELEM:.*]] = dim %[[ARG]], %[[I]] : tensor<*xf32>
+// CHECK:             store %[[ELEM]], %[[MEMREF]][%[[I]]] : memref<?xindex>
+// CHECK:             scf.yield
+// CHECK:           }
+// CHECK:           %[[RET:.*]] = tensor_load %[[MEMREF]] : memref<?xindex>
+// CHECK:           return %[[RET]] : tensor<?xindex>
+// CHECK:         }
+func @dynamic_tensor_from_elements(%arg: tensor<*xf32>, %rank: index) -> tensor<?xindex> {
+  %result = dynamic_tensor_from_elements %rank {
+  ^bb0(%i : index):
+    %elem = dim %arg, %i : tensor<*xf32>
+    yield %elem : index
+  } : tensor<?xindex>
+  return %result : tensor<?xindex>
+}
+
+// Additional test that checks the logic for intermixed static and dynamic
+// extents.
+//
+// CHECK-LABEL:   func @dynamic_tensor_from_elements_static_and_dynamic(
+// CHECK-SAME:                                                          %[[DYNAMIC_EXTENT:.*]]: index) -> tensor<16x?xindex> {
+// CHECK:           %[[MEMREF:.*]] = alloc(%[[DYNAMIC_EXTENT]]) : memref<16x?xindex>
+// CHECK:           %[[C0:.*]] = constant 0 : index
+// CHECK:           %[[C1:.*]] = constant 1 : index
+// CHECK:           %[[C16:.*]] = constant 16 : index
+// CHECK:           scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]]) to (%[[C16]], %[[DYNAMIC_EXTENT]]) step (%[[C1]], %[[C1]]) {
+// CHECK:             %[[VAL_7:.*]] = addi %[[I]], %[[J]] : index
+// CHECK:             store %[[VAL_7]], %[[MEMREF]][%[[I]], %[[J]]] : memref<16x?xindex>
+// CHECK:             scf.yield
+// CHECK:           }
+// CHECK:           %[[RET:.*]] = tensor_load %[[MEMREF]] : memref<16x?xindex>
+// CHECK:           return %[[RET]] : tensor<16x?xindex>
+// CHECK:         }
+func @dynamic_tensor_from_elements_static_and_dynamic(%arg0: index) -> tensor<16x?xindex> {
+  %result = dynamic_tensor_from_elements %arg0 {
+  ^bb0(%i: index, %j: index):
+    %sum = addi %i, %j : index
+    yield %sum : index
+  } : tensor<16x?xindex>
+  return %result : tensor<16x?xindex>
+}
+
 // 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 {