[Mlir-commits] [mlir] 4a255bb - [mlir] Add folding for shape.any

[Tres Popp]

Author: Tres Popp
Date: 2020-06-05T11:00:19+02:00
New Revision: 4a255bbd29698d9d5fbf05fd21ab78ca540365f2

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

LOG: [mlir] Add folding for shape.any

If any input to shape.any is a const_shape, shape.any can be replaced
with that input.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index b66ea38d03da..075050a8c2b1 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -438,7 +438,7 @@ def Shape_ConcatOp : Shape_Op<"concat", []> {
 //===----------------------------------------------------------------------===//
 
 //TODO(tpopp): Move the code below and witnesses to a 
diff erent file.
-def Shape_AnyOp : Shape_Op<"any", [NoSideEffect]> {
+def Shape_AnyOp : Shape_Op<"any", [Commutative, NoSideEffect]> {
   let summary = "Return any combination of the input shapes";
   let description = [{
     This operation takes multiple input shapes and returns some combination of
@@ -458,6 +458,8 @@ def Shape_AnyOp : Shape_Op<"any", [NoSideEffect]> {
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "$inputs attr-dict";
+
+  let hasFolder = 1;
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {

diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index a4a8b2de59fd..2b05c4c65bab 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -101,6 +101,16 @@ void ShapeDialect::printType(Type type, DialectAsmPrinter &os) const {
 // AnyOp
 //===----------------------------------------------------------------------===//
 
+// TODO: Canonicalization should be implemented for shapes that can be
+// determined through mixtures of the known dimensions of the inputs.
+OpFoldResult AnyOp::fold(ArrayRef<Attribute> operands) {
+  // Only the last operand is checked because AnyOp is commutative.
+  if (operands.back())
+    return operands.back();
+
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // AssumingOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 646700f8d6bf..0f92c18c25a8 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -245,3 +245,25 @@ func @f() {
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
+
+// -----
+// any can be replaced with a constant input if it has one.
+// CHECK-LABEL: func @f
+func @f(%arg0 : !shape.shape) -> !shape.shape {
+  // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
+  // CHECK-NEXT: return %[[CS]]
+  %0 = shape.const_shape [2, 3, 4]
+  %1 = shape.any %0, %arg0
+  return %1 : !shape.shape
+}
+
+
+// -----
+// Folding of any with partially constant operands is not yet implemented.
+// CHECK-LABEL: func @f
+func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
+  // CHECK-NEXT: shape.any
+  // CHECK-NEXT: return %[[CS]]
+  %1 = shape.any %arg0, %arg1
+  return %1 : !shape.shape
+}