[Mlir-commits] [mlir] 80be54c - [mlir] Lower Shape binary ops (AddOp, MulOp) to Standard.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2020-06-08T17:48:01+02:00
New Revision: 80be54c08f6c2185fde854c0aabb48971489fdb6

URL: https://github.com/llvm/llvm-project/commit/80be54c08f6c2185fde854c0aabb48971489fdb6
DIFF: https://github.com/llvm/llvm-project/commit/80be54c08f6c2185fde854c0aabb48971489fdb6.diff

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

Differential Revision: https://reviews.llvm.org/D81344

Added: 
    

Modified: 
    mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
    mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
index 5c74cdcaa241..81293a028489 100644
--- a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
+++ b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
@@ -15,10 +15,26 @@
 #include "mlir/Transforms/DialectConversion.h"
 
 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 FromExtentTensorOpConversion
     : public OpConversionPattern<shape::FromExtentTensorOp> {
 public:
@@ -128,6 +144,8 @@ void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
+      BinaryOpConversion<AddOp, AddIOp>,
+      BinaryOpConversion<MulOp, MulIOp>,
       FromExtentTensorOpConversion,
       IndexToSizeOpConversion,
       SizeToIndexOpConversion,

diff  --git a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
index de420a96a70f..fd1f980c5369 100644
--- a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
+++ b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
@@ -62,3 +62,16 @@ func @from_extent_tensor(%tensor : tensor<?xindex>) -> !shape.shape {
       : (tensor<?xindex>) -> !shape.shape
   return %shape : !shape.shape
 }
+
+// -----
+
+// 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
+}