[mlir] Lower shape.num_elements -> shape.reduce.

[pifon2a]

Update the comment.

[pifon2a]

Fix cmake.

[silvas]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

330

Why this rename? I think "body" sounds better.

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

15

Can we find a better name for this pass? The goal is to reduce shape ops to other more primitive shape ops. That this can simplify converting to standard is a byproduct but not essential.

Maybe "ShapeToShapeLowering?

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

3

nit: this looks like it got line wrapped.

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

17

Is AtomicRMWOp needed?

[jpienaar]

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

24

Could we expand here on what the pass does too?

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

2

Could we fit the next along with this line?

29

Prefer to keep anonymous namespaces small (https://llvm.org/docs/CodingStandards.html#anonymous-namespaces)

[pifon2a]

Address the comments.

[pifon2a]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

330

This is to avoid confusion, since SingleBlockImplicitTerminator generates a getBody() method that return Block*, i.e. body().front() if using the old naming.

scf.for and scf.parallel also call it just region.

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

15

Naming is hard. Thanks, I like ShapeToShapeLowering!

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

17

well, yes, but not here :)

[pifon2a]

Thank you, Jacques and Sean!

[mehdi_amini]

Reverted in a25f5cd70cef ; the build was broken with -DBUILD_SHARED_LIBS=ON

mlir/docs/Passes.md

 ## `loop` Dialect Passes
 
 [include "LoopPasses.md"]
 
 ## `quant` Dialect Passes
 
 [include "QuantPasses.md"]
 
+## `shape` Dialect Passes
+
+[include "Shape.md"]
+
 ## `spv` Dialect Passes
 
 [include "SPIRVPasses.md"]
 
 ## `standard` Dialect Passes
 
 [include "StandardPasses.md"]

mlir/include/mlir/Dialect/Shape/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let description = [{
     }
     ```
 
     If the shape is unranked, then the results of the op is also unranked.
   }];
 
   let arguments = (ins Shape_ShapeType:$shape, Variadic<AnyType>:$initVals);
   let results = (outs Variadic<AnyType>:$result);
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region SizedRegion<1>:$region);
 
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "Value shape, ValueRange initVals">,
   ];
 
   let verifier = [{ return ::verify(*this); }];
   let printer = [{ return ::print(p, *this); }];

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

+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls)
+add_public_tablegen_target(MLIRStandardTransformsIncGen)
+
+add_mlir_doc(Passes -gen-pass-doc StandardPasses ./)

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

+//===- Passes.h - Pass Entrypoints ------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes that expose pass constructors in the
+// shape transformation library.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
+#define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
+
+#include <memory>
+
+namespace mlir {
+
+class Pass;
+
+/// Creates an instance of the LegalizeShapeOpsForExport pass.
+std::unique_ptr<Pass> createLegalizeShapeOpsForExport();
+
+} // 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 LegalizeShapeOpsForExport : FunctionPass<"legalize-shape-ops-for-export"> {
+  let summary = "Legalize Shape dialect to be convertible to Standard";
+  let constructor = "mlir::createLegalizeShapeOpsForExport()";
+}
+
+#endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Quant/Passes.h"
 #include "mlir/Dialect/SCF/Passes.h"
 #include "mlir/Dialect/SPIRV/Passes.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Transforms/LocationSnapshot.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/ViewOpGraph.h"
 #include "mlir/Transforms/ViewRegionGraph.h"
 
 #include <cstdlib>
 
#include "mlir/Dialect/Quant/Passes.h.inc"
 
   // SPIR-V
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SPIRV/Passes.h.inc"
 
   // Standard
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
+
+  // Shape
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/Shape/Transforms/Passes.h.inc"
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLPASSES_H_

mlir/lib/Dialect/Shape/CMakeLists.txt

add_mlir_dialect_library(MLIRShape
 
   LINK_LIBS PUBLIC
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   )
+
+add_subdirectory(Transforms)

mlir/lib/Dialect/Shape/IR/Shape.cpp

void ReduceOp::build(OpBuilder &builder, OperationState &result, Value shape,
   for (Type initValType : initVals.getTypes()) {
     bodyBlock.addArgument(initValType);
     result.addTypes(initValType);
   }
 }
 
 static LogicalResult verify(ReduceOp op) {
   // Verify block arg types.
-  Block &block = op.body().front();
+  Block &block = op.region().front();
 
   auto blockArgsCount = op.initVals().size() + 2;
   if (block.getNumArguments() != blockArgsCount)
     return op.emitOpError() << "ReduceOp body is expected to have "
                             << blockArgsCount << " arguments";
 
   if (block.getArgument(0).getType() != IndexType::get(op.getContext()))
     return op.emitOpError(
static ParseResult parseReduceOp(OpAsmParser &parser, OperationState &result) {
 
   return success();
 }
 
 static void print(OpAsmPrinter &p, ReduceOp op) {
   p << op.getOperationName() << '(' << op.shape() << ", " << op.initVals()
     << ") ";
   p.printOptionalArrowTypeList(op.getResultTypes());
-  p.printRegion(op.body());
+  p.printRegion(op.region());
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 namespace mlir {
 namespace shape {
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
 
 } // namespace shape
 } // namespace mlir

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

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

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

+//===- LegalizeShapeOpsForExport.cpp - Prepare for lowering to Standard
+//----===//
+//
+// 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/Shape/IR/Shape.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.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.num_elements` to `shape.reduce`.
+struct NumElementsOpConverter : public OpRewritePattern<NumElementsOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(NumElementsOp op,
+                                PatternRewriter &rewriter) const final {
+    auto loc = op.getLoc();
+    Value init = rewriter.create<ConstSizeOp>(loc, rewriter.getIndexAttr(1));
+    ReduceOp reduce = rewriter.create<ReduceOp>(loc, op.shape(), init);
+
+    // Generate reduce operator.
+    Block *body = reduce.getBody();
+    OpBuilder b = OpBuilder::atBlockEnd(body);
+    Value product =
+        b.create<MulOp>(loc, body->getArgument(1), body->getArgument(2));
+    b.create<YieldOp>(loc, product);
+
+    rewriter.replaceOp(op, reduce.result());
+    return success();
+  }
+};
+
+struct LegalizeShapeOpsForExport
+    : public LegalizeShapeOpsForExportBase<LegalizeShapeOpsForExport> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    patterns.insert<NumElementsOpConverter>(&getContext());
+
+    ConversionTarget target(getContext());
+    target.addLegalDialect<ShapeDialect>();
+    target.addIllegalOp<NumElementsOp>();
+    if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::createLegalizeShapeOpsForExport() {
+  return std::make_unique<LegalizeShapeOpsForExport>();
+}

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

+//===- PassDetail.h - Shape 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_SHAPE_TRANSFORMS_PASSDETAIL_H_
+#define DIALECT_SHAPE_TRANSFORMS_PASSDETAIL_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+class AtomicRMWOp;
+
+#define GEN_PASS_CLASSES
+#include "mlir/Dialect/Shape/Transforms/Passes.h.inc"
+
+} // end namespace mlir
+
+#endif // DIALECT_SHAPE_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/Shape/legalize-for-export.mlir

+// RUN: mlir-opt -legalize-shape-ops-for-export -split-input-file %s | FileCheck %s --dump-input-on-failure
+
+// 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]]
+