[mlir][shape] Add a pattern to rewrite `shape.reduce` as `scf.for`.

[pifon2a]

@jpienaar Jacques, I put this pattern to Transform directory, because I am not really converting Shape Ops inside the reduction operator. The only thing that changes here is just the reduce op "expanded" as a loop.
I am also not completely happy with how I named the pass. So, suggestions are very-very welcome.

[silvas]

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

19

All dialects share the same "namespace" of pass names, so I recommend calling it shape-* to avoid naming conflicts with other dialect passes.

Also, lower-reduce-to-scf and lower-shape-to-shape would be my name preferences so it is consistent at a glance what are lowering passes.

mlir/lib/Dialect/Shape/Transforms/ExpandReduce.cpp

41 ↗

(On Diff #270237)

llvm naming convention is extentTensor

52 ↗

(On Diff #270237)

should this use b instead of rewriter?

57 ↗

(On Diff #270237)

can you use mlir::inlineRegion for this?

[pifon2a]

Address the comments.

[pifon2a]

mlir/lib/Dialect/Shape/Transforms/ExpandReduce.cpp

57 ↗

(On Diff #270237)

I am not sure that would look nicer since I need to convert shape.yield to scf.yield after inlining. Also blocks have different signatures. I would prefer to leave it as it is.

[herhut]

Nice to see this become a thing!

mlir/lib/Dialect/Shape/Transforms/ShapeToSCFLowering.cpp

1

Should this live in Conversion/ShapeToSCF? Only transformations within a dialect are in transforms IIUC.

46

This should be dim %extent, 0 and not rank. The rank of the extent tensor is always 1 (as it is a vector).

62

clone with a mapping already does this for you.

65

Can this be done with makeArrayRef?

[ftynse]

mlir/lib/Dialect/Shape/Transforms/ShapeToSCFLowering.cpp

1

Yes, and it should not be called "lowering". That name was deprecated long time ago.

[silvas]

mlir/lib/Dialect/Shape/Transforms/ShapeToSCFLowering.cpp

1

Can you update the docs https://mlir.llvm.org/getting_started/Glossary/#lowering ? I was not aware of that deprecation.

[pifon2a]

mlir/lib/Dialect/Shape/Transforms/ShapeToSCFLowering.cpp

65

if we could lookup arrays of values, then yes.

[pifon2a]

Address the comments.

[frgossen]

mlir/include/mlir/Conversion/ShapeToSCF/ShapeToSCF.h

14 ↗

(On Diff #270754)

If I remember correctly, the using namespace mlir variant is preferred in LLVM.

[herhut]

Neat!

mlir/include/mlir/Conversion/ShapeToSCF/ShapeToSCF.h

14 ↗

(On Diff #270754)

Only for uses, not definitions.

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

 
 class Pass;
 
 /// Creates an instance of the ShapeToShapeLowering pass that legalizes Shape
 /// dialect to be convertible to Standard. For example, `shape.num_elements` get
 /// transformed to `shape.reduce`, which can be lowered to SCF and Standard.
 std::unique_ptr<Pass> createShapeToShapeLowering();
 
+/// Creates an instance of the ShapeToSCFLowering pass that converts Shape ops
+/// to SCF dialect. For example, `shape.reduce` gets rewritten as an `scf.for`.
+/// The operations within the region of ReduceOp are cloned to the ForOp body.
+std::unique_ptr<Pass> createShapeToSCFLowering();
+
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_

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

 //===-- Passes.td - ShapeOps pass definition file ----------*- tablegen -*-===//
 //
 // 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 MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES
 #define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
-def ShapeToShapeLowering : FunctionPass<"shape-to-shape-lowering"> {
+def ShapeToShapeLowering : FunctionPass<"shape-to-shape"> {
   let summary = "Legalize Shape dialect to be convertible to Standard";
   let constructor = "mlir::createShapeToShapeLowering()";
 }
 
+def ShapeToSCFLowering : FunctionPass<"shape-to-scf"> {
+  let summary = "Convert Shape ops to SCF dialect";
+  let constructor = "mlir::createShapeToSCFLowering()";
+}
+
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES

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

 add_mlir_dialect_library(MLIRShapeOpsTransforms
   ShapeToShapeLowering.cpp
+  ShapeToSCFLowering.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ShapeOps/Transforms
 
   DEPENDS
   MLIRShapeTransformsIncGen
   )
 
 target_link_libraries(MLIRShapeOpsTransforms
   PUBLIC
   MLIRIR
   MLIRPass
   MLIRShape
   MLIRSupport
   MLIRTransforms
   )

mlir/lib/Dialect/Shape/Transforms/ShapeToSCFLowering.cpp

+//===- ShapeToSCFLowering.cpp - Expand to SCF ForOp -----------------------===//
+//
+// 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 "PassDetail.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+using namespace mlir::shape;
+
+namespace {
+/// Converts `shape.reduce` to `scf.for`.
+struct ReduceOpConverter : public OpRewritePattern<ReduceOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ReduceOp op,
+                                PatternRewriter &rewriter) const final;
+};
+} // namespace
+
+LogicalResult
+ReduceOpConverter::matchAndRewrite(ReduceOp reduceOp,
+                                   PatternRewriter &rewriter) const {
+  auto loc = reduceOp.getLoc();
+
+  Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+  Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+  Value extentTensor = rewriter.create<ToExtentTensorOp>(
+      loc,
+      RankedTensorType::get({ShapedType::kDynamicSize},
+                            rewriter.getIndexType()),
+      reduceOp.shape());
+  Value rank = rewriter.create<RankOp>(loc, extentTensor);
+
+  auto loop = rewriter.create<scf::ForOp>(
+      loc, zero, rank, one, reduceOp.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());
+
+        BlockAndValueMapping mapping;
+        Block *loopBody = reduceOp.getBody();
+        mapping.map(loopBody->getArguments(), mapped_values);
+        for (auto &nested : loopBody->without_terminator()) {
+          Operation *clone = b.clone(nested, mapping);
+          mapping.map(nested.getResults(), clone->getResults());
+        }
+
+        SmallVector<Value, 2> mappedResults;
+        for (auto result : loopBody->getTerminator()->getOperands())
+          mappedResults.push_back(mapping.lookup(result));
+        b.create<scf::YieldOp>(loc, mappedResults);
+      });
+
+  rewriter.replaceOp(reduceOp, loop.getResults());
+  return success();
+}
+
+namespace {
+struct ShapeToSCFLowering : public ShapeToSCFLoweringBase<ShapeToSCFLowering> {
+  void runOnFunction() override;
+};
+} // namespace
+
+void ShapeToSCFLowering::runOnFunction() {
+  OwningRewritePatternList patterns;
+  patterns.insert<ReduceOpConverter>(&getContext());
+
+  ConversionTarget target(getContext());
+  target.addLegalDialect<ShapeDialect, scf::SCFDialect, StandardOpsDialect>();
+  target.addIllegalOp<ReduceOp>();
+  if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+    signalPassFailure();
+}
+
+std::unique_ptr<Pass> mlir::createShapeToSCFLowering() {
+  return std::make_unique<ShapeToSCFLowering>();
+}

mlir/test/Dialect/Shape/shape-to-scf.mlir

+// RUN: mlir-opt -shape-to-scf -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: shape_reduce
+// CHECK-SAME:   [[SHAPE:%.*]]: !shape.shape) -> !shape.size {
+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: [[RANK:%.*]] = rank [[EXTENTS]] : tensor<?xindex>
+
+// CHECK-NEXT: [[RESULT:%.*]] = scf.for [[I:%.*]] = [[C0]] to [[RANK]]
+// 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: }
+// CHECK-NEXT: return [[RESULT]] : !shape.size

mlir/test/Dialect/Shape/shape-to-shape.mlir

-// RUN: mlir-opt -shape-to-shape-lowering -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -shape-to-shape -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: func @num_elements_to_reduce(
 // CHECK-SAME:    [[ARG:%.*]]: !shape.shape) -> [[SIZE_TY:!.*]] {
 func @num_elements_to_reduce(%shape : !shape.shape) -> !shape.size {
   %num_elements = shape.num_elements %shape
   return %num_elements : !shape.size
 }
 // CHECK: [[C1:%.*]] = shape.const_size 1
 // CHECK: [[NUM_ELEMENTS:%.*]] = shape.reduce([[ARG]], [[C1]])  -> [[SIZE_TY]]
 // CHECK: ^bb0({{.*}}: index, [[DIM:%.*]]: [[SIZE_TY]], [[ACC:%.*]]: [[SIZE_TY]]
 // CHECK:   [[NEW_ACC:%.*]] = shape.mul [[DIM]], [[ACC]]
 // CHECK:   shape.yield [[NEW_ACC]] : [[SIZE_TY]]
 // CHECK: }
 // CHECK: return [[NUM_ELEMENTS]] : [[SIZE_TY]]