[mlir] Add folding for shape.any

[tpopp]

Summary

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

[silvas]

mlir/test/Dialect/Shape/canonicalize.mlir

154

Can you expand a bit on this "yet"? Is that a direction we want to go? It seems like it would be a very small change to this patch to support that, suggesting there is a deeper reason why we want to avoid that, which would be good to document here or in AnyOp::fold

[jpienaar]

mlir/lib/Dialect/Shape/IR/Shape.cpp

106

So the folder only focuses on case where the output can be fully statically computed? [this overlaps with Sean's question below]

108

Nit: Perhaps expand the comment slightly: it is commutative and so constants have been moved to the trailing operands.

This would be known to someone who read the canonicalization doc or looked at FoldUtils.cpp , but I think it is useful to avoid needing that context especially as short.

[herhut]

mlir/test/Dialect/Shape/canonicalize.mlir

154

We would need to inspect operands and their types for this to work, which is more complex than just looking at constant operands. So it makes sense to land this first, if only to support static shape workloads, and expand later. A comment like Folding of any with partially constant operands is not yet implemented should convey this.

[tpopp]

Explain commutativity of any.

[tpopp]

mlir/lib/Dialect/Shape/IR/Shape.cpp

106

Yes, I started to add that as a comment, but then realized that's essentially the definition of folding, so I think it would be redundant to state.

mlir/test/Dialect/Shape/canonicalize.mlir

154

There's a not a specific reason to not support it. I would rather wait until there's a clear use case though to implement more of this as it's more risky given the inherent lack of safety in the operation.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

 }
 
 //===----------------------------------------------------------------------===//
 // Shape constraint related ops.
 //===----------------------------------------------------------------------===//
 
 //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
     their dimensions. This can be best seen with examples below.
 
     The result is undefined, but still side-effect free, in cases where the
     inputs have differing ranks or differ in extents of shared dimensions.
 
     Example:
     ```mlir
       %s0 = shape.any [2,?], [?,3] // [2,3]
       %s1 = shape.any [?,?], [1,2] // [1,2]
     ```
   }];
 
   let arguments = (ins Variadic<Shape_ShapeType>:$inputs);
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "$inputs attr-dict";
+
+  let hasFolder = 1;
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {
   let summary = "Return a logical AND of all witnesses.";
   let description = [{
     Used to simplify constraints as any single failing precondition is enough
     to prevent execution.
 

mlir/lib/Dialect/Shape/IR/Shape.cpp

default:
     llvm_unreachable("unexpected 'shape' type kind");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // 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,
                         RegionRange regions,
                         SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }

mlir/test/Dialect/Shape/canonicalize.mlir

func @f() {
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.true_witness
   %1 = "test.source"() : () -> !shape.witness
   %2 = shape.assuming_all %0, %1
   "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
+}