[MLIR][Shape] Lower `shape_of` to `dynamic_tensor_from_elements`

[frgossen]

Summary

Take advantage of the new dynamic_tensor_from_elements operation in std.
Instead of stack-allocated memory, we can now lower directly to a single std
operation.

Depends On D86779

[herhut]

Mostly nits. Thanks for adding this!

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

460

not: iv or idx? Could also use dim as this is what this value represents.

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

1704

result.regions.front()?

1705

Maybe use b.createBlock and a SmallVector to represent the index types?

1714

createBlock would do this for you, so you have to move the guard in front of it.

[jpienaar]

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

1509

Nit: could you add a comment here? (just basic tblgen one as I don't think we have way to generate comments in C++ from here yet)

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

460

Could args.front() be used?

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

1708

Is this verified so that we can just cast here?

[frgossen]

Address comments

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

1708

Yes, result type is AnyRankedTensor.

[herhut]

Thanks!

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

```mlir
         yield %elem : f32
       } : tensor<?x3x?f32>
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamicExtents);
   let results = (outs AnyRankedTensor:$result);
   let regions = (region SizedRegion<1>:$body);
+
+  let builders = [
+    // Build op and populate its body per callback function.
+    OpBuilder<"OpBuilder &b, OperationState &result, Type resultTy, "
+              "ValueRange dynamicExtents, "
+              "function_ref<void(OpBuilder &, Location, ValueRange)>">,
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // ExpOp
 //===----------------------------------------------------------------------===//
 
 def ExpOp : FloatUnaryOp<"exp"> {
   let summary = "base-e exponential of the specified value";

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

LogicalResult ShapeOfOpConversion::matchAndRewrite(
     ShapeOfOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
 
   // For now, only error-free types are supported by this lowering.
   if (op.getType().isa<ShapeType>())
     return failure();
 
   // For ranked tensor arguments, lower to `tensor_from_elements`.
+  auto loc = op.getLoc();
   ShapeOfOp::Adaptor transformed(operands);
   Value tensor = transformed.arg();
   Type tensorTy = tensor.getType();
   if (tensorTy.isa<RankedTensorType>()) {
 
     // Build values for individual extents.
     SmallVector<Value, 8> extentValues;
     RankedTensorType rankedTensorTy = tensorTy.cast<RankedTensorType>();
     int64_t rank = rankedTensorTy.getRank();
-    auto loc = op.getLoc();
     for (int64_t i = 0; i < rank; i++) {
       if (rankedTensorTy.isDynamicDim(i)) {
         Value extent = rewriter.create<DimOp>(loc, tensor, i);
         extentValues.push_back(extent);
       } else {
         Value extent =
             rewriter.create<ConstantIndexOp>(loc, rankedTensorTy.getDimSize(i));
         extentValues.push_back(extent);
       }
     }
 
     // Materialize extent tensor.
     Value staticExtentTensor =
         rewriter.create<TensorFromElementsOp>(loc, extentValues);
     rewriter.replaceOpWithNewOp<TensorCastOp>(op, staticExtentTensor,
                                               op.getType());
     return success();
   }
 
-  // Allocate stack memory.
-  auto loc = op.getLoc();
+  // Lower to `dynamic_tensor_from_elements` otherwise.
+  auto *ctx = rewriter.getContext();
   Value rank = rewriter.create<mlir::RankOp>(loc, tensor);
-  Type indexTy = rewriter.getIndexType();
-  Type memTy = MemRefType::get({ShapedType::kDynamicSize}, indexTy);
-  Value mem = rewriter.create<AllocaOp>(loc, memTy, ValueRange{rank});
-
-  // Copy shape extents to stack-allocated memory.
-  Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
-  Value one = rewriter.create<ConstantIndexOp>(loc, 1);
-  rewriter.create<scf::ForOp>(
-      loc, zero, rank, one, llvm::None,
-      [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {
-        Value dim = rewriter.create<DimOp>(loc, tensor, iv);
-        rewriter.create<StoreOp>(loc, dim, mem, ValueRange{iv});
-        rewriter.create<scf::YieldOp>(loc);
+  rewriter.replaceOpWithNewOp<DynamicTensorFromElementsOp>(
+      op, getExtentTensorType(ctx), ValueRange{rank},
+      [&](OpBuilder &b, Location loc, ValueRange args) {
+        Value dim = args.front();
+        Value extent = b.create<DimOp>(loc, tensor, dim);
+        b.create<mlir::YieldOp>(loc, extent);
       });
 
-  // Load extents to tensor value.
-  rewriter.replaceOpWithNewOp<TensorLoadOp>(op.getOperation(), mem);
   return success();
 }
 
 namespace {
 class ToExtentTensorOpConversion
     : public OpConversionPattern<ToExtentTensorOp> {
 public:
   using OpConversionPattern<ToExtentTensorOp>::OpConversionPattern;

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

static LogicalResult verify(DynamicTensorFromElementsOp op) {
   if (yieldOp.value().getType() != resultTy.getElementType())
     return op.emitOpError(
         "body must be terminated with a `yield` operation of the tensor "
         "element type");
 
   return success();
 }
 
+void DynamicTensorFromElementsOp::build(
+    OpBuilder &b, OperationState &result, Type resultTy,
+    ValueRange dynamicExtents,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
+  build(b, result, resultTy, dynamicExtents);
+
+  // Build and populate body.
+  OpBuilder::InsertionGuard guard(b);
+  Region *bodyRegion = result.regions.front().get();
+  auto rank = resultTy.cast<RankedTensorType>().getRank();
+  SmallVector<Type, 2> argumentTypes(rank, b.getIndexType());
+  Block *bodyBlock =
+      b.createBlock(bodyRegion, bodyRegion->end(), argumentTypes);
+  bodyBuilder(b, result.location, bodyBlock->getArguments());
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ExtractElementOp op) {
   // Verify the # indices match if we have a ranked type.
   auto aggregateType = op.getAggregate().getType().cast<ShapedType>();
   if (aggregateType.hasRank() &&
       aggregateType.getRank() != op.getNumOperands() - 1)
     return op.emitOpError("incorrect number of indices for extract_element");

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 
 // -----
 
 // Lower `shape_of` for unranked tensors.
 // CHECK-LABEL: @shape_of_unranked
 // CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
 func @shape_of_unranked(%arg : tensor<*xf32>) {
   // CHECK: %[[RANK:.*]] = rank %[[ARG]] : tensor<*xf32>
-  // CHECK: %[[SHAPE_MEM:.*]] = alloca(%[[RANK]]) : memref<?xindex>
-  // CHECK: %[[C0:.*]] = constant 0 : index
-  // CHECK: %[[C1:.*]] = constant 1 : index
-  // CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[RANK]] step %[[C1]] {
-  // CHECK:   %[[DIM:.]] = dim %[[ARG]], %[[I]] : tensor<*xf32>
-  // CHECK:   store %[[DIM]], %[[SHAPE_MEM]][%[[I]]] : memref<?xindex>
-  // CHECK: }
-  // CHECK: %[[SHAPE:.*]] = tensor_load %[[SHAPE_MEM]] : memref<?xindex>
+  // CHECK: %[[SHAPE:.*]] = dynamic_tensor_from_elements %[[RANK]] {
+  // CHECK: ^bb0(%[[I:.*]]: index):
+  // CHECK:   %[[EXTENT:.*]] = dim %[[ARG]], %[[I]] : tensor<*xf32>
+  // CHECK:   yield %[[EXTENT]] : index
+  // CHECK: } : tensor<?xindex>
   %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
   return
 }
 
 // -----
 
 // Don't lower `shape_of` with `shape.shape` type.
 // CHECK-LABEL: @shape_of