[PATCH] D80305: [mlir] Add folding for shape.any

[Theodore Popp via Phabricator via llvm-commits]

tpopp created this revision.
Herald added subscribers: llvm-commits, jurahul, Kayjukh, frgossen, grosul1, Joonsoo, stephenneuendorffer, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: jpienaar.
Herald added a reviewer: silvas.
Herald added a project: LLVM.
tpopp added a reviewer: herhut.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D80305

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


Index: mlir/test/Dialect/Shape/canonicalize.mlir
===================================================================
--- mlir/test/Dialect/Shape/canonicalize.mlir
+++ mlir/test/Dialect/Shape/canonicalize.mlir
@@ -137,3 +137,25 @@
   "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
+}
+
+
+// -----
+// any is not yet replaced without a constant input.
+// 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
+}
Index: mlir/lib/Dialect/Shape/IR/Shape.cpp
===================================================================
--- mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -101,6 +101,16 @@
 // 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;
+}
+
 LogicalResult
 AnyOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                         ValueRange operands, DictionaryAttr attributes,
Index: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
===================================================================
--- mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -370,7 +370,7 @@
 
 //TODO(tpopp): Move the code below and witnesses to a different file.
 def Shape_AnyOp : Shape_Op<"any",
-    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+    [Commutative, NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Return any combination of the input shapes.";
   let description = [{
     This operation takes multiple input shapes and returns some combination of
@@ -390,6 +390,8 @@
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "$inputs attr-dict";
+
+  let hasFolder = 1;
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D80305.265269.patch
Type: text/x-patch
Size: 2628 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200520/45e52624/attachment.bin>