[mlir] Canonicalization of shape.cstr_broadcastable

[tpopp]

Summary

This allows replacing of this op with a true_witness in the case of both
inputs being const_shapes and being found to be broadcastable.

[silvas]

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

274

As far as the constant folding logic, you can just use an i1 for a witness. You just need to add a case in ShapeDialect::materializeConstant know how to materialize it into a true_witness op or false_witness op (or just a single const_witness?).

Maybe you can rename this patch "introduce basic constant folding for witnesses" and use this fold as the first example?

[herhut]

I am fine with this to land and make the transition to folding later.

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

272

Would a simple op.lhs() == op.rhs() work?

281

getOperand(0)->lhs()?
getOperand(1)->rhs()?

[tpopp]

Switch to folding and also handle repeated inputs.

[tpopp]

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

274

It took me a long time to figure out that I could use a BoolAttr for this. In the spirit of small commits, I'm leaving the ConstWitnessOp creation separate, but I'm not folding this for constant false witnesses because that needs to be carefully thought through as that is removing an assert that can be triggered.

[silvas]

Looks great! Thanks!

[isuruf]

This commit breaks building MLIR for me. You need to create the directory IR in the build directory in mlir/lib/Dialect/Shape/CMakeLists.txt. Otherwise mlir-tblgen fails with a file not found error.

[isuruf]

https://reviews.llvm.org/D81283

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

def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result",
     [{ /* nothing to do */ }]>
   ];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
-def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
+def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", [Commutative]> {
   let summary = "Determines if 2 shapes can be successfully broadcasted";
   let description = [{
     Given 2 input shapes, return a witness specifying if they are broadcastable.
     This broadcastable follows the same logic as what shape.broadcast documents.
 
     "cstr" operations represent runtime assertions.
 
     Example:
     ```mlir
       %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
       %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
     ```
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
+
+  let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 def Shape_CstrEqOp : Shape_Op<"cstr_eq", []> {
   let summary = "Determines if all input shapes are equal";
   let description = [{
     Given 1 or more input shapes, determine if all shapes are the exact same.
 
     "cstr" operations represent runtime assertions.

mlir/lib/Dialect/Shape/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS IR/ShapeCanonicalization.td)
+mlir_tablegen(IR/ShapeCanonicalization.inc -gen-rewriters)
+add_public_tablegen_target(MLIRShapeCanonicalizationIncGen)
+
 add_mlir_dialect_library(MLIRShape
   IR/Shape.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
 
   DEPENDS
   MLIRShapeOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   )

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

 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
+namespace {
+#include "IR/ShapeCanonicalization.inc"
+}
+
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
            WitnessType>();
static ParseResult parseConstShapeOp(OpAsmParser &parser,
 
   result.types.push_back(ShapeType::get(builder.getContext()));
   return success();
 }
 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
 
 //===----------------------------------------------------------------------===//
+// CstrBroadcastableOp
+//===----------------------------------------------------------------------===//
+
+void CstrBroadcastableOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  // If inputs are equal, return passing witness
+  patterns.insert<CstrBroadcastableEqOps>(context);
+}
+
+OpFoldResult CstrBroadcastableOp::fold(ArrayRef<Attribute> operands) {
+  if (!operands[0] || !operands[1])
+    return nullptr;
+  auto lhsShape = llvm::to_vector<6>(
+      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  auto rhsShape = llvm::to_vector<6>(
+      operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  SmallVector<int64_t, 6> resultShape;
+  if (OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
+    return BoolAttr::get(true, getContext());
+
+  // Because a failing witness result here represents an eventual assertion
+  // failure, we do not replace it with a constant witness.
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
 
 //===----------------------------------------------------------------------===//
 // ConstWitnessOp
 //===----------------------------------------------------------------------===//

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

+include "mlir/Dialect/Shape/IR/ShapeOps.td"
+
+def EqualBinaryOperands : Constraint<CPred<"$0 == $1">>;
+
+// Canonicalization patterns.
+def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $lhs, $rhs),
+  (Shape_ConstWitnessOp ConstBoolAttrTrue),
+  [(EqualBinaryOperands $lhs, $rhs)]>;

mlir/test/Dialect/Shape/canonicalize.mlir

 // 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
 }
+
+// -----
+// Broadcastable with broadcastable constant shapes can be removed.
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: shape.const_witness true
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [3, 1]
+  %cs1 = shape.const_shape [1, 5]
+  %0 = shape.cstr_broadcastable %cs0, %cs1
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// Broadcastable with non-broadcastable constant shapes is always false
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: shape.cstr_broadcastable
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [1, 3]
+  %cs1 = shape.const_shape [1, 5]
+  %0 = shape.cstr_broadcastable %cs0, %cs1
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// Broadcastable without guaranteed broadcastable shapes cannot be removed.
+// CHECK-LABEL: func @f
+func @f(%arg0 : !shape.shape) {
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: shape.cstr_broadcastable
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [1,3]
+  %0 = shape.cstr_broadcastable %arg0, %cs0
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// Broadcastable with non-constant but known equal shapes can be removed.
+// CHECK-LABEL: func @f
+func @f(%arg0 : !shape.shape) {
+  // CHECK-NEXT: shape.const_witness true
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.cstr_broadcastable %arg0, %arg0
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}