[MLIR][Shape] Support transforming shape.num_elements on tensors

[herhut]

Summary

The current transformation to shape.reduce does not support tensor values.
This adds the required changes to make that work, including fixing the builder
for shape.reduce.

[frgossen]

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

839

I think these braces can be removed.

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

68

In most other places, the legal dialects are listed in the template parameters.
addLegalDialect<ShapeDialect, StandardOpsDialect>()

[herhut]

Comments on formatting

[herhut]

Thanks!

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

void ReduceOp::build(OpBuilder &builder, OperationState &result, Value shape,
                      ValueRange initVals) {
   result.addOperands(shape);
   result.addOperands(initVals);
 
   Region *bodyRegion = result.addRegion();
   bodyRegion->push_back(new Block);
   Block &bodyBlock = bodyRegion->front();
   bodyBlock.addArgument(builder.getIndexType());
-  bodyBlock.addArgument(SizeType::get(builder.getContext()));
+
+  Type elementType;
+  if (auto tensorType = shape.getType().dyn_cast<TensorType>())
+    elementType = tensorType.getElementType();
+  else
+    elementType = SizeType::get(builder.getContext());
+  bodyBlock.addArgument(elementType);
 
   for (Type initValType : initVals.getTypes()) {
     bodyBlock.addArgument(initValType);
     result.addTypes(initValType);
   }
 }
 
 static LogicalResult verify(ReduceOp op) {

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

 //===- ShapeToShapeLowering.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/Dialect/StandardOps/IR/Ops.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;
 };
 } // namespace
 
 LogicalResult
 NumElementsOpConverter::matchAndRewrite(NumElementsOp op,
                                         PatternRewriter &rewriter) const {
   auto loc = op.getLoc();
-  Value init = rewriter.create<ConstSizeOp>(loc, rewriter.getIndexAttr(1));
+  Type valueType = op.getResult().getType();
+  Value init = op.getDialect()
+                   ->materializeConstant(rewriter, rewriter.getIndexAttr(1),
+                                         valueType, loc)
+                   ->getResult(0);
   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, b.getType<SizeType>(),
-                                  body->getArgument(1), body->getArgument(2));
+  Value product = b.create<MulOp>(loc, valueType, body->getArgument(1),
+                                  body->getArgument(2));
   b.create<YieldOp>(loc, product);
 
   rewriter.replaceOp(op, reduce.result());
   return success();
 }
 
 namespace {
 struct ShapeToShapeLowering
     : public ShapeToShapeLoweringBase<ShapeToShapeLowering> {
   void runOnFunction() override;
 };
 } // namespace
 
 void ShapeToShapeLowering::runOnFunction() {
   MLIRContext &ctx = getContext();
 
   OwningRewritePatternList patterns;
   populateShapeRewritePatterns(&ctx, patterns);
 
   ConversionTarget target(getContext());
-  target.addLegalDialect<ShapeDialect>();
+  target.addLegalDialect<ShapeDialect, StandardOpsDialect>();
   target.addIllegalOp<NumElementsOp>();
   if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
     signalPassFailure();
 }
 
 void mlir::populateShapeRewritePatterns(MLIRContext *context,
                                         OwningRewritePatternList &patterns) {
   patterns.insert<NumElementsOpConverter>(context);
 }
 
 std::unique_ptr<Pass> mlir::createShapeToShapeLowering() {
   return std::make_unique<ShapeToShapeLowering>();
 }

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

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