MLIR][Shape] Add function signature conversion to `shape` to `scf` lowering.

[frgossen]

Summary

The pass lowers fully to scf (and std where unavoidable).
It includes patterns from the shape to std lowering pass for testing
purposes.
This change prepares for the lowering of shape.shape_eq to scf.

Depends On D82528

[frgossen]

Remove CHECK-DAGs

[rriddle]

mlir/lib/IR/Builders.cpp

169 ↗

(On Diff #273388)

This method doesn't look like it is used in this revision.

[frgossen]

Rebase

[frgossen]

No longer relevant for the new lowering strategy

mlir/include/mlir/Dialect/SCF/SCFOps.td

let description = [{
     If the parent operation defines no values, then the "scf.yield" may be
     left out in the custom syntax and the builders will insert one implicitly.
     Otherwise, it has to be present in the syntax to indicate which values are
     yielded.
   }];
 
   let arguments = (ins Variadic<AnyType>:$results);
   let builders = [
-    OpBuilder<"OpBuilder &builder, OperationState &result",
-              [{ /* nothing to do */ }]>
+    OpBuilder<"OpBuilder &builder, OperationState &state",
+              [{ /* nothing to do */ }]>,
+    OpBuilder<"OpBuilder &builder, OperationState &state, Value result",
+              [{ build(builder, state, ValueRange({result})); }]>
   ];
 }
 #endif // MLIR_DIALECT_SCF_SCFOPS

mlir/lib/Conversion/ShapeToSCF/CMakeLists.txt

add_mlir_conversion_library(MLIRShapeToSCF
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRShape
   MLIRPass
   MLIRSCF
+  MLIRShapeToStandard
   MLIRTransforms
   )

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

 //===- ShapeToSCF.cpp - conversion from Shape to SCF dialect --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/ShapeToSCF/ShapeToSCF.h"
 
 #include "../PassDetail.h"
+#include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::shape;
+using namespace mlir::scf;
 
+/// Conversion patterns.
 namespace {
 /// Converts `shape.reduce` to `scf.for`.
-struct ReduceOpConverter : public OpRewritePattern<ReduceOp> {
+struct ReduceOpConverter : public OpConversionPattern<shape::ReduceOp> {
 public:
-  using OpRewritePattern::OpRewritePattern;
+  using OpConversionPattern<shape::ReduceOp>::OpConversionPattern;
 
-  LogicalResult matchAndRewrite(ReduceOp op,
-                                PatternRewriter &rewriter) const final;
+  LogicalResult
+  matchAndRewrite(shape::ReduceOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 LogicalResult
-ReduceOpConverter::matchAndRewrite(ReduceOp reduceOp,
-                                   PatternRewriter &rewriter) const {
+ReduceOpConverter::matchAndRewrite(shape::ReduceOp reduceOp,
+                                   ArrayRef<Value> operands,
+                                   ConversionPatternRewriter &rewriter) const {
+  shape::ReduceOp::Adaptor transformed(operands);
   auto loc = reduceOp.getLoc();
 
+  Type indexTy = rewriter.getIndexType();
   Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+  Value size = rewriter.create<DimOp>(loc, indexTy, transformed.shape(), zero);
   Value one = rewriter.create<ConstantIndexOp>(loc, 1);
-  Value extentTensor = rewriter.create<ToExtentTensorOp>(
-      loc,
-      RankedTensorType::get({ShapedType::kDynamicSize},
-                            rewriter.getIndexType()),
-      reduceOp.shape());
-  Value size =
-      rewriter.create<DimOp>(loc, rewriter.getIndexType(), extentTensor, zero);
-
   auto loop = rewriter.create<scf::ForOp>(
-      loc, zero, size, one, reduceOp.initVals(),
+      loc, zero, size, one, transformed.initVals(),
       [&](OpBuilder &b, Location nestedLoc, Value iv, ValueRange args) {
-        Value indexExtent = b.create<ExtractElementOp>(loc, extentTensor, iv);
-        Value sizeExtent = b.create<IndexToSizeOp>(loc, indexExtent);
-
-        SmallVector<Value, 2> mapped_values{iv, sizeExtent};
-        mapped_values.append(args.begin(), args.end());
-
+        // TODO: Consider to determine the extent with a `std.dim` operation if
+        // the shape originates in a tensor.
+        Value extent = b.create<ExtractElementOp>(loc, transformed.shape(), iv);
+
+        // Copy `shape.reduce` body without terminator and substitute values.
+        SmallVector<Value, 2> mappedValues{iv, extent};
+        mappedValues.append(args.begin(), args.end());
         BlockAndValueMapping mapping;
         Block *reduceBody = reduceOp.getBody();
-        mapping.map(reduceBody->getArguments(), mapped_values);
+        mapping.map(reduceBody->getArguments(), mappedValues);
         for (auto &nested : reduceBody->without_terminator())
           b.clone(nested, mapping);
 
+        // Copy terminator and substitute values.
         SmallVector<Value, 2> mappedResults;
         for (auto result : reduceBody->getTerminator()->getOperands())
           mappedResults.push_back(mapping.lookup(result));
         b.create<scf::YieldOp>(loc, mappedResults);
       });
 
   rewriter.replaceOp(reduceOp, loop.getResults());
   return success();
ShapeOfOpConverter::matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
   auto shapeIntVal = rewriter.create<TensorLoadOp>(loc, memVal);
   auto indexTy = rewriter.getIndexType();
   auto shapeTy = RankedTensorType::get({ShapedType::kDynamicSize}, indexTy);
   rewriter.replaceOpWithNewOp<IndexCastOp>(op.getOperation(), shapeIntVal,
                                            shapeTy);
   return success();
 }
 
+/// Type conversions.
+class ShapeTypeConverter : public TypeConverter {
+public:
+  using TypeConverter::convertType;
+
+  ShapeTypeConverter(MLIRContext *ctx) {
+    // Add default pass-through conversion.
+    addConversion([&](Type type) { return type; });
+
+    addConversion([ctx](SizeType type) { return IndexType::get(ctx); });
+    addConversion([ctx](ShapeType type) {
+      return RankedTensorType::get({ShapedType::kDynamicSize},
+                                   IndexType::get(ctx));
+    });
+  }
+};
+
+/// Conversion pass.
 namespace {
 struct ConvertShapeToSCFPass
     : public ConvertShapeToSCFBase<ConvertShapeToSCFPass> {
   void runOnFunction() override;
 };
 } // namespace
 
 void ConvertShapeToSCFPass::runOnFunction() {
+  // Setup type conversion.
   MLIRContext &ctx = getContext();
+  ShapeTypeConverter typeConverter(&ctx);
 
   // Populate conversion patterns.
   OwningRewritePatternList patterns;
   populateShapeToSCFConversionPatterns(patterns, &ctx);
+  populateShapeToStandardConversionPatterns(patterns, &ctx);
+  populateFuncOpTypeConversionPattern(patterns, &ctx, typeConverter);
 
   // Setup target legality.
   ConversionTarget target(getContext());
-  target.addLegalDialect<ShapeDialect, scf::SCFDialect, StandardOpsDialect>();
-  target.addIllegalOp<ReduceOp, ShapeOfOp>();
-
-  // Apply conversion.
+  target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
+  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+    return typeConverter.isSignatureLegal(op.getType()) &&
+           typeConverter.isLegal(&op.getBody());
+  });
   if (failed(applyPartialConversion(getFunction(), target, patterns)))
     signalPassFailure();
 }
 
 void mlir::populateShapeToSCFConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   patterns.insert<ReduceOpConverter, ShapeOfOpConverter>(ctx);
 }
 
 std::unique_ptr<FunctionPass> mlir::createConvertShapeToSCFPass() {
   return std::make_unique<ConvertShapeToSCFPass>();
 }

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

 // RUN: mlir-opt -convert-shape-to-scf -split-input-file %s | FileCheck %s
 
-// CHECK-LABEL: shape_reduce
-// CHECK-SAME:   [[SHAPE:%.*]]: !shape.shape) -> !shape.size {
+// CHECK-LABEL: @shape_reduce
+// CHECK-SAME:  (%[[SHAPE:.*]]: tensor<?xindex>) -> index
 func @shape_reduce(%shape : !shape.shape) -> !shape.size {
   %init = shape.const_size 1
   %num_elements = shape.reduce(%shape, %init) -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %acc: !shape.size):
       %new_acc = shape.mul %acc, %dim
       shape.yield %new_acc : !shape.size
   }
   return %num_elements : !shape.size
 }
-// CHECK-NEXT: [[SHAPE_C1:%.*]] = shape.const_size 1
-// CHECK-NEXT: [[C0:%.*]] = constant 0 : index
-// CHECK-NEXT: [[C1:%.*]] = constant 1 : index
-
-// CHECK-NEXT: [[EXTENTS:%.*]] = shape.to_extent_tensor [[SHAPE]]
-// CHECK-NEXT: [[SIZE:%.*]] = dim [[EXTENTS]], [[C0]] : tensor<?xindex>
-
-// CHECK-NEXT: [[RESULT:%.*]] = scf.for [[I:%.*]] = [[C0]] to [[SIZE]]
-// CHECK-SAME:       step [[C1]] iter_args([[ACC:%.*]] = [[SHAPE_C1]])
-// CHECK-NEXT:   [[EXTENT_INDEX:%.*]] = extract_element [[EXTENTS]]{{\[}}[[I]]]
-// CHECK-NEXT:   [[EXTENT:%.*]] = shape.index_to_size [[EXTENT_INDEX]]
-// CHECK-NEXT:   [[NEW_ACC:%.*]] = shape.mul [[ACC]], [[EXTENT]]
-// CHECK-NEXT:   scf.yield [[NEW_ACC]] : !shape.size
+// CHECK-NEXT: %[[INIT:.*]] = constant 1 : index
+// CHECK-NEXT: %[[C0:.*]] = constant 0 : index
+// CHECK-NEXT: %[[SIZE:.*]] = dim %[[SHAPE]], %[[C0]] : tensor<?xindex>
+// CHECK-NEXT: %[[C1:.*]] = constant 1 : index
+// CHECK-NEXT: %[[RESULT:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[SIZE]]
+// CHECK-SAME:     step %[[C1]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (index) {
+// CHECK-NEXT:   %[[EXTENT:.*]] = extract_element %[[SHAPE]][%[[I]]]
+// CHECK-NEXT:   %[[ACC_NEXT:.*]] = muli %[[ACC]], %[[EXTENT]]
+// CHECK-NEXT:   scf.yield %[[ACC_NEXT]] : index
 // CHECK-NEXT: }
-// CHECK-NEXT: return [[RESULT]] : !shape.size
+// CHECK-NEXT: return %[[RESULT]] : index
 
 // -----
 
 // Lower `shape_of` for unranked tensors.
 // CHECK-LABEL: @shape_of_unranked
 // CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
 func @shape_of_unranked(%arg : tensor<*xf32>) {
   // CHECK-DAG: %[[RANK:.*]] = rank %[[ARG]] : tensor<*xf32>