[Mlir-commits] [mlir] 88b8b88 - [MLIR] Remove empty shape operands from `cstr_broadcastable` ops

[Frederik Gossen]

Author: Frederik Gossen
Date: 2021-04-26T18:34:18+02:00
New Revision: 88b8b88035faf2752b6d9f75e89333526a19b102

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

LOG: [MLIR] Remove empty shape operands from `cstr_broadcastable` ops

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

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 392bcddb920a..c19acf33f535 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -695,7 +695,8 @@ void CstrBroadcastableOp::getCanonicalizationPatterns(
   // folding in case additional shape information is inferred at some point that
   // does not result in folding.
   patterns.add<CstrBroadcastableEqOps,
-               RemoveDuplicateOperandsPattern<CstrBroadcastableOp>>(context);
+               RemoveDuplicateOperandsPattern<CstrBroadcastableOp>,
+               RemoveEmptyShapeOperandsPattern<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 95d0de628b75..68545d161374 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -638,6 +638,20 @@ func @f() {
   return
 }
 
+// -----
+// Empty shape arguments can be removed from broadcastable ops.
+// CHECK-LABEL: func @f
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>)
+func @f(%arg0 : tensor<?xindex>, %arg1 : tensor<?xindex>) {
+  // CHECK-NOT: const_shape
+  // CHECK: cstr_broadcastable %[[ARG0]], %[[ARG1]] : tensor<?xindex>, tensor<?xindex>
+  %0 = shape.const_shape [] : !shape.shape
+  %1 = shape.cstr_broadcastable %arg0, %arg1, %0
+      : tensor<?xindex>, tensor<?xindex>, !shape.shape
+  "consume.witness"(%1) : (!shape.witness) -> ()
+  return
+}
+
 // -----
 // Broadcastable with non-broadcastable constant shapes is always false
 // CHECK-LABEL: func @static_non_broadcastable