[Mlir-commits] [mlir] f8d7bd9 - [MLIR][Shape] Remove empty extent tensor operands

[Frederik Gossen]

Author: Frederik Gossen
Date: 2021-04-27T14:51:43+02:00
New Revision: f8d7bd996f1e9a01b9d5cee39c579dd8ac71bae4

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

LOG: [MLIR][Shape] Remove empty extent tensor operands

Empty extent tensor operands were only removed when they were defined as a
constant. Additionally, we can remove them if they are known to be empty by
their type `tensor<0xindex>`.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index ac98bfef1b58..2e6677c12c79 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/mlir/include/mlir/IR/BuiltinAttributes.h
@@ -70,6 +70,9 @@ class ElementsAttr : public Attribute {
   /// Returns the number of elements held by this attribute.
   int64_t size() const { return getNumElements(); }
 
+  /// Returns if the number of elements held by this attribute is 0.
+  bool empty() const { return size() == 0; }
+
   /// Generates a new ElementsAttr by mapping each int value to a new
   /// underlying APInt. The new values can represent either an integer or float.
   /// This ElementsAttr should contain integers.

diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index ea2b10bdcf0e..96618d24747f 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -534,8 +534,14 @@ struct RemoveEmptyShapeOperandsPattern : public OpRewritePattern<OpTy> {
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
     auto isPotentiallyNonEmptyShape = [](Value shape) {
-      if (auto constShape = shape.getDefiningOp<ConstShapeOp>())
-        return constShape.shape().size() != 0;
+      if (auto extentTensorTy = shape.getType().dyn_cast<RankedTensorType>()) {
+        if (extentTensorTy.getDimSize(0) == 0)
+          return false;
+      }
+      if (auto constShape = shape.getDefiningOp<ConstShapeOp>()) {
+        if (constShape.shape().empty())
+          return false;
+      }
       return true;
     };
     auto newOperands = llvm::to_vector<8>(

diff  --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index a47db3044e0e..6dc48b1e732a 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -641,13 +641,13 @@ func @f() {
 // -----
 // 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-SAME:  (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>, %{{.*}}: tensor<0xindex>)
+func @f(%arg0 : tensor<?xindex>, %arg1 : tensor<?xindex>, %arg2 : tensor<0xindex>) {
   // 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
+  %1 = shape.cstr_broadcastable %arg0, %arg1, %0, %arg2
+      : tensor<?xindex>, tensor<?xindex>, !shape.shape, tensor<0xindex>
   "consume.witness"(%1) : (!shape.witness) -> ()
   return
 }