[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 ↗

(On Diff #270237)

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 ↗

(On Diff #270365)

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

46 ↗

(On Diff #270365)

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

62 ↗

(On Diff #270365)

clone with a mapping already does this for you.

65 ↗

(On Diff #270365)

Can this be done with makeArrayRef?

[ftynse]

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

1 ↗

(On Diff #270365)

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

[silvas]

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

1 ↗

(On Diff #270365)

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 ↗

(On Diff #270365)

if we could lookup arrays of values, then yes.

[pifon2a]

Address the comments.

[frgossen]

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

15

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

[herhut]

Neat!

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

15

Only for uses, not definitions.

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> {
   let summary = "Convert operations from the shape dialect into the standard "
                 "dialect";
   let constructor = "mlir::createConvertShapeToStandardPass()";
 }
 
 //===----------------------------------------------------------------------===//
+// ShapeToSCF
+//===----------------------------------------------------------------------===//
+
+def ConvertShapeToSCF : FunctionPass<"convert-shape-to-scf"> {
+  let summary = "Convert operations from the shape dialect to the SCF dialect";
+  let constructor = "mlir::createConvertShapeToSCFPass()";
+}
+
+//===----------------------------------------------------------------------===//
 // SPIRVToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertSPIRVToLLVM : Pass<"convert-spirv-to-llvm", "ModuleOp"> {
   let summary = "Convert SPIR-V dialect to LLVM dialect";
   let constructor = "mlir::createConvertSPIRVToLLVMPass()";
 }
 

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

+//===- ShapeToSCF.h - Conversion utils 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_SHAPETOSCF_SHAPETOSCF_H_
+#define MLIR_CONVERSION_SHAPETOSCF_SHAPETOSCF_H_
+
+#include <memory>
+
+namespace mlir {
+
+class MLIRContext;
+class FunctionPass;
+class OwningRewritePatternList;
+
+void populateShapeToSCFConversionPatterns(OwningRewritePatternList &patterns,
+                                          MLIRContext *ctx);
+
+std::unique_ptr<FunctionPass> createConvertShapeToSCFPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_SHAPETOSCF_SHAPETOSCF_H_

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.h"
+#include "mlir/Conversion/ShapeToSCF/ShapeToSCF.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AVX512ToLLVM)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToStandard)
+add_subdirectory(ShapeToSCF)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)

mlir/lib/Conversion/ShapeToSCF/CMakeLists.txt

+add_mlir_conversion_library(MLIRShapeToSCF
+  ShapeToSCF.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ShapeToSCF
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRShape
+  MLIRPass
+  MLIRSCF
+  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/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;
+
+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 size =
+      rewriter.create<DimOp>(loc, rewriter.getIndexType(), extentTensor, zero);
+
+  auto loop = rewriter.create<scf::ForOp>(
+      loc, zero, size, 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 *reduceBody = reduceOp.getBody();
+        mapping.map(reduceBody->getArguments(), mapped_values);
+        for (auto &nested : reduceBody->without_terminator())
+          b.clone(nested, mapping);
+
+        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();
+}
+
+namespace {
+struct ConvertShapeToSCFPass
+    : public ConvertShapeToSCFBase<ConvertShapeToSCFPass> {
+  void runOnFunction() override;
+};
+} // namespace
+
+void ConvertShapeToSCFPass::runOnFunction() {
+  MLIRContext &ctx = getContext();
+
+  OwningRewritePatternList patterns;
+  populateShapeToSCFConversionPatterns(patterns, &ctx);
+
+  ConversionTarget target(getContext());
+  target.addLegalDialect<ShapeDialect, scf::SCFDialect, StandardOpsDialect>();
+  target.addIllegalOp<ReduceOp>();
+  if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+    signalPassFailure();
+}
+
+void mlir::populateShapeToSCFConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  patterns.insert<ReduceOpConverter>(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 {
+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: }
+// CHECK-NEXT: return [[RESULT]] : !shape.size