[mlir] Lower Shape binary ops (AddOp, MulOp) to Standard.

[frgossen]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

23

nit: If these were called SrcOp/DstOp it could be clearer that these are operations, not types.

[pifon2a]

Addressed the comment.

[herhut]

Nice! Thanks.

[jpienaar]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

31

What does the adaptor give you which the generated op doesn't?

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 //===- LinalgToStandard.cpp - conversion from Linalg to Standard 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/ShapeToStandard/ShapeToStandard.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/Transforms/DialectConversion.h"
 
-namespace mlir {
+using namespace mlir;
+using namespace mlir::shape;
+
 namespace {
 
 /// Conversion patterns.
+template <typename SrcOpTy, typename DstOpTy>
+class BinaryOpConversion : public OpConversionPattern<SrcOpTy> {
+public:
+  using OpConversionPattern<SrcOpTy>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    typename SrcOpTy::OperandAdaptor adaptor(operands);
+    rewriter.replaceOpWithNewOp<DstOpTy>(op.getOperation(), adaptor.lhs(),
+                                         adaptor.rhs());
+    return success();
+  }
+};
+
 class SizeToIndexOpConversion
     : public OpConversionPattern<shape::SizeToIndexOp> {
 public:
   using OpConversionPattern<shape::SizeToIndexOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(shape::SizeToIndexOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
void runOnOperation() override {
     auto module = getOperation();
     if (failed(applyFullConversion(module, target, patterns, &typeConverter)))
       signalPassFailure();
   }
 };
 
 } // namespace
 
-void populateShapeToStandardConversionPatterns(
+void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
+      BinaryOpConversion<AddOp, AddIOp>,
+      BinaryOpConversion<MulOp, MulIOp>,
       IndexToSizeOpConversion,
       SizeToIndexOpConversion>(ctx);
   // clang-format on
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> createConvertShapeToStandardPass() {
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }
-
-} // namespace mlir

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // Lower `index_to_size` conversion to no-op.
 // CHECK-LABEL: @index_to_size
 // CHECK-SAME: (%[[INDEX:.*]]: index) -> index
 func @index_to_size(%index : index) -> !shape.size {
   // CHECK-NEXT: return %[[INDEX]] : index
   %size = shape.index_to_size %index
   return %size : !shape.size
 }
+
+// -----
+
+// Lower binary ops.
+// CHECK-LABEL: @binary_ops
+// CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
+func @binary_ops(%lhs : !shape.size, %rhs : !shape.size) {
+  %sum = "shape.add"(%lhs, %rhs) : (!shape.size, !shape.size) -> !shape.size
+  // CHECK-NEXT: addi %[[LHS]], %[[RHS]] : index
+  %product = shape.mul %lhs, %rhs
+  // CHECK-NEXT: muli %[[LHS]], %[[RHS]] : index
+  return
+}