[Mlir-commits] [mlir] f368b3a - [MLIR][Shape] Canonicalize duplicate operands in `shape.cstr_broadcastable`

[Frederik Gossen]

Author: Frederik Gossen
Date: 2021-03-23T12:23:22+01:00
New Revision: f368b3a029887eb09fe9c2b537b76bf6f4274ae4

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

LOG: [MLIR][Shape] Canonicalize duplicate operands in `shape.cstr_broadcastable`

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Shape/IR/Shape.cpp
    mlir/test/Dialect/Shape/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index f68478042e65..d2a10a9f5dcc 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -553,7 +553,8 @@ void CstrBroadcastableOp::getCanonicalizationPatterns(
   // Canonicalization patterns have overlap with the considerations during
   // folding in case additional shape information is inferred at some point that
   // does not result in folding.
-  patterns.add<CstrBroadcastableEqOps>(context);
+  patterns.add<CstrBroadcastableEqOps,
+               RemoveDuplicateOperandsPattern<CstrBroadcastableOp>>(context);
 }
 
 // Return true if there is exactly one attribute not representing a scalar

diff  --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 3399fe0f4e23..39f17e9d253f 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -648,7 +648,7 @@ func @f() {
   // CHECK: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [8, 1] : !shape.shape
+  %cs0 = shape.const_shape [8, 1] : !shape.shape  
   %cs1 = shape.const_shape [1, 8] : !shape.shape
   %cs2 = shape.const_shape [1, -1] : !shape.shape
   %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
@@ -665,7 +665,7 @@ func @f() {
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [8, 1] : !shape.shape
   %cs1 = shape.const_shape [1, -1] : !shape.shape
-  %cs2 = shape.const_shape [1, -1] : !shape.shape
+  %cs2 = shape.const_shape [8, -1] : !shape.shape
   %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
@@ -1097,6 +1097,19 @@ func @is_broadcastable_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)
 
 // -----
 
+// CHECK-LABEL: @cstr_broadcastable_on_duplicate_shapes
+// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)
+func @cstr_broadcastable_on_duplicate_shapes(%a : !shape.shape,
+    %b : !shape.shape) -> !shape.witness {
+  // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[A]], %[[B]] :
+  // CHECK: return %[[RES]]
+  %0 = shape.cstr_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape,
+      !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape
+  return %0 : !shape.witness
+}
+
+// -----
+
 // CHECK-LABEL: @broadcast_on_same_shape
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)
 func @broadcast_on_same_shape(%shape : !shape.shape) -> !shape.shape {