[mlir][Shape] Canonicalize assume_all with one input and tensor_cast of constant_shape

[bkramer]

Summary

This allows simplifying some more complicated shape expressions

[tpopp]

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

277–279

I don't think you really need these comments.

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

16

Do you actually need the AssumingAll predicate? I thought this would only match on AssumingAll with a single input. Do multiple inputs still get matched and stored as a list?

36

Should this be $ty?

mlir/test/Dialect/Shape/canonicalize.mlir

871

nit: new line after this line.

[bkramer]

• Address comments

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

16

Yes, $x is a list here, so we have to check the length.

36

No, this is working as intended (although a bit weird with the implicit type cast)

[tpopp]

mlir/test/Dialect/Shape/canonicalize.mlir

430–431

This needs to be updated to if more than one witness is not statically passing

[herhut]

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

36

This only works if the consumer can accept the new type of $c. This would break for example if the consumer is a select operation.

%s0 = shape.const_shape [2] : tensor<?xindex>
%s1 = shape.const_shape [1] : tensor<1xindex>
%c = tensor_cast %s1 : tensor<1xindex> to tensor<?xindex>
%r = select %p %s0 %c : tensor<?xindex>

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

def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
   }];
   let arguments = (ins IndexElementsAttr:$shape);
   let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
   // TODO: Move this to main so that all shape ops implement these.
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_ConstSizeOp : Shape_Op<"const_size", [
     ConstantLike,
     NoSideEffect,
     DeclareOpInterfaceMethods<OpAsmOpInterface>
   ]> {
   let summary = "Creates a constant of type `shape.size`";
def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {
   }];
 
   let arguments = (ins Variadic<Shape_WitnessType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$inputs attr-dict";
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 
   let verifier = [{ return ::verify(*this); }];
 }
 
 def Shape_AssumingOp : Shape_Op<"assuming",
                            [SingleBlockImplicitTerminator<"AssumingYieldOp">,
                             DeclareOpInterfaceMethods<RegionBranchOpInterface>,
                             RecursiveSideEffects]> {

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

void AssumingOp::inlineRegionIntoParent(AssumingOp &op,
   // AssumingOp.
   rewriter.mergeBlocks(assumingBlock, blockBeforeAssuming);
   rewriter.mergeBlocks(blockAfterAssuming, blockBeforeAssuming);
 }
 
 //===----------------------------------------------------------------------===//
 // AssumingAllOp
 //===----------------------------------------------------------------------===//
+
+void AssumingAllOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<AssumingAllOneOp>(context);
+}
+
 OpFoldResult AssumingAllOp::fold(ArrayRef<Attribute> operands) {
   // Iterate in reverse to first handle all constant operands. They are
   // guaranteed to be the tail of the inputs because this is commutative.
   for (int idx = operands.size() - 1; idx >= 0; idx--) {
     Attribute a = operands[idx];
     // Cannot fold if any inputs are not constant;
     if (!a)
       return nullptr;
static ParseResult parseConstShapeOp(OpAsmParser &parser,
   if (parser.parseColonType(resultTy))
     return failure();
   result.types.push_back(resultTy);
   return success();
 }
 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
 
+void ConstShapeOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<TensorCastConstShape>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // CstrBroadcastableOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 // Given an input shape Value, try to obtain the shape's values.
 LogicalResult getShapeVec(Value input, SmallVectorImpl<int64_t> &shapeValues) {
   if (auto inputOp = input.getDefiningOp<ShapeOfOp>()) {

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

 include "mlir/Dialect/Shape/IR/ShapeOps.td"
+include "mlir/Dialect/StandardOps/IR/Ops.td"
 
 def AllInputShapesEq : Constraint<CPred< [{
   llvm::all_of($0, [&](mlir::Value val) {
     return $0[0] == val;
   })
 }]>>;
 
+def HasSingleElement : Constraint<CPred< [{
+  $0.size() == 1
+}]>>;
+
 // Canonicalization patterns.
 
+def AssumingAllOneOp : Pat<(Shape_AssumingAllOp $args),
+                           (replaceWithValue $args),
+                           [(HasSingleElement $args)]>;
+
 def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $x, $x),
   (Shape_ConstWitnessOp ConstBoolAttrTrue)>;
 
 def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
 
 def IndexToSizeToIndexCanonicalization : Pat<
   (Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
   (replaceWithValue $arg)>;
 
 def SizeToIndexToSizeCanonicalization : Pat<
   (Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
   (replaceWithValue $arg)>;
 
+def TensorCastConstShape : Pat <
+  (TensorCastOp (Shape_ConstShapeOp:$c $ty)), (replaceWithValue $c)>;

mlir/test/Dialect/Shape/canonicalize.mlir

func @f() {
   %2 = shape.const_witness true
   %3 = shape.assuming_all %0, %1, %2
   "consume.witness"(%3) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
-// assuming_all should not be removed if not all witnesses are statically passing.
+// assuming_all should not be removed if more than one witness is not
+// statically passing
 //
 // Additionally check that the attribute is moved to the end as this op is
 // commutative.
 // CHECK-LABEL: func @f
 func @f() {
-  // CHECK-NEXT: %[[UNKNOWN:.*]] = "test.source"
-  // CHECK-NEXT: shape.assuming_all %[[UNKNOWN]]
+  // CHECK-NEXT: %[[UNKNOWN1:.*]] = "test.source"
+  // CHECK-NEXT: %[[UNKNOWN2:.*]] = "test.source"
+  // CHECK-NEXT: shape.assuming_all %[[UNKNOWN1]], %[[UNKNOWN2]]
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_witness true
   %1 = "test.source"() : () -> !shape.witness
-  %2 = shape.assuming_all %0, %1
-  "consume.witness"(%2) : (!shape.witness) -> ()
+  %2 = "test.source"() : () -> !shape.witness
+  %3 = shape.assuming_all %0, %1, %2
+  "consume.witness"(%3) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
 func @f(%arg : !shape.shape) -> !shape.shape {

 
 // Fold to_extent_tensor when already on tensor.
 // CHECK-LABEL: @fold_to_extent_tensor_on_tensor
 func @fold_to_extent_tensor_on_tensor(%arg: tensor<?xindex>) -> tensor<?xindex> {
   // CHECK-NOT: to_extent_tensor
   %casted = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<?xindex>
   return %casted : tensor<?xindex>
 }
+
+// -----
+
+// Fold assuming_all with a single input
+// CHECK-LABEL: @fold_assuming_all_single_element
+func @fold_assuming_all_single_element(%arg: tensor<?xindex>) {
+  // CHECK-NOT: assuming_all
+  %0 = "test.source"() : () -> (!shape.witness)
+  %1 = shape.assuming_all %0
+  "consume.witness"(%1) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+
+// Fold tensor_cast of a const_shape to const_shape
+// CHECK-LABEL: @fold_tensor_cast_of_const_shape
+func @fold_tensor_cast_of_const_shape(%arg: tensor<?xindex>) {
+  // CHECK-NOT: tensor_cast
+  %0 = shape.const_shape [2] : tensor<?xindex>
+  %1 = tensor_cast %0 : tensor<?xindex> to tensor<1xindex>
+  %2 = shape.cstr_broadcastable %1, %0 : tensor<1xindex>, tensor<?xindex>
+  "consume.witness"(%2) : (!shape.witness) -> ()
+  return
+}