[MLIR][Shape] Add error message to `cstr_broadcastable` operation

[frgossen]

Summary

Add error message to concretize the context in which the constraint fails.

[jpienaar]

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

847

If we made the default value "required broadcastable shapes", would that reduce the changes needed below?

851

This would seem to need an update?

[frgossen]

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

847

That's what I wanted but it causes problems with the variadic shapes. The builder functions are then often ambiguous, even in the generated code.

[frgossen]

Address comments

[frgossen]

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

847

Correction: The ambiguity derives not from the variadic shapes but from the default builders. When neither the message nor an attribute dict are passed there are two candidate function.

[herhut]

+1 for the general direction. I agree a default would be useful, so that we do not need to pass a message when we do not care about the specific wording. If that requires fighting with the infra, lets land this for now.

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

 }
 
 def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable",
                                          [Commutative, InferTypeOpInterface]> {
   let summary = "Determines if 2+ shapes can be successfully broadcasted";
   let description = [{
     Given input shapes or extent tensors, return a witness specifying if they
     are broadcastable. This broadcastable follows the same logic as what
-    shape.broadcast documents.
+    shape.broadcast documents. The op carries an error message to concretize the
+    context in which the constraint fails.
 
     "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
+    %w0 = shape.cstr_broadcastable [2,2], [3,1,2] msg("require broadcastable shapes") // Passing
+    %w1 = shape.cstr_broadcastable [2,2], [3,2] msg("require broadcastable shapes") // Failure
     ```
   }];
 
-  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
+  let arguments = (ins
+    Variadic<Shape_ShapeOrExtentTensorType>:$shapes,
+    StrAttr:$msg
+  );
   let results = (outs Shape_WitnessType:$result);
 
-  let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
+  let assemblyFormat = "$shapes `msg` `(` $msg `)` attr-dict `:` type($shapes)";
 
   let builders = [
   OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
     [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
   ];
 
   let extraClassDeclaration = [{
     // TODO: This should really be automatic. Figure out how to not need this defined.

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_SHAPETOSTANDARD_TD
 #define MLIR_CONVERSION_SHAPETOSTANDARD_TD
 
 include "mlir/Dialect/Shape/IR/ShapeOps.td"
 
-def BroadcastableStringAttr : NativeCodeCall<[{
-  $_builder.getStringAttr("required broadcastable shapes")
-}]>;
-
-def CstrBroadcastableToRequire : Pat<(Shape_CstrBroadcastableOp $shapes),
-            (Shape_CstrRequireOp
-              (Shape_IsBroadcastableOp $shapes),
-              (BroadcastableStringAttr))>;
+def CstrBroadcastableToRequire : Pat<(Shape_CstrBroadcastableOp $shapes, $msg),
+    $msg
+  )>;
 
 def EqStringAttr : NativeCodeCall<[{
   $_builder.getStringAttr("required equal shapes")
 }]>;
 
 def CstrEqToRequire : Pat<(Shape_CstrEqOp $shapes),
   (Shape_CstrRequireOp (Shape_ShapeEqOp $shapes), (EqStringAttr))>;
 
 #endif // MLIR_CONVERSION_SHAPETOSTANDARD_TD

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

auto canonicalizeOperand = [&](Value operand) {
       return operand;
     };
     auto newOperands = llvm::to_vector<8>(
         llvm::map_range(op.getOperands(), canonicalizeOperand));
 
     // Rewrite op if any change required.
     if (!anyChange)
       return failure();
-    rewriter.replaceOpWithNewOp<OpTy>(op, op->getResultTypes(), newOperands);
+    rewriter.replaceOpWithNewOp<OpTy>(op, op->getResultTypes(), newOperands,
+                                      op->getAttrs());
     return success();
   }
 };
 
 struct BroadcastConcretizeResultTypePattern
     : public OpRewritePattern<BroadcastOp> {
   using OpRewritePattern<BroadcastOp>::OpRewritePattern;
 

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

 }]>>;
 
 // Canonicalization patterns.
 
 def AssumingAllOneOp : Pat<(Shape_AssumingAllOp $args),
                            (replaceWithValue $args),
                            [(HasSingleElement $args)]>;
 
-def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $shapes),
+def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $shapes, $msg),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
 
 def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
 
 def IndexToSizeToIndexCanonicalization : Pat<

mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir

 // RUN: mlir-opt -convert-shape-constraints <%s | FileCheck %s
 
 // There's not very much useful to check here other than pasting the output.
 // CHECK-LABEL:   func @cstr_broadcastable(
 // CHECK-SAME:                             %[[LHS:.*]]: tensor<?xindex>,
 // CHECK-SAME:                             %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
 // CHECK:           %[[RET:.*]] = shape.const_witness true
 // CHECK:           %[[BROADCAST_IS_VALID:.*]] = shape.is_broadcastable %[[LHS]], %[[RHS]]
-// CHECK:           assert %[[BROADCAST_IS_VALID]], "required broadcastable shapes"
+// CHECK:           assert %[[BROADCAST_IS_VALID]], "require broadcastable shapes"
 // CHECK:           return %[[RET]] : !shape.witness
 // CHECK:         }
 func @cstr_broadcastable(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> !shape.witness {
-  %witness = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  %witness = shape.cstr_broadcastable %arg0, %arg1 msg("require broadcastable shapes") : tensor<?xindex>, tensor<?xindex>
   return %witness : !shape.witness
 }
 
 // CHECK-LABEL:   func @cstr_eq(
 // CHECK-SAME:                             %[[LHS:.*]]: tensor<?xindex>,
 // CHECK-SAME:                             %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
 // CHECK:           %[[RET:.*]] = shape.const_witness true
 // CHECK:           %[[EQUAL_IS_VALID:.*]] = shape.shape_eq %[[LHS]], %[[RHS]]

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // CHECK:                %[[AND_REDUCTION:.*]] = and %[[SECOND_REDUCTION]], %[[GOOD]] : i1
 // CHECK:                scf.yield %[[AND_REDUCTION]] : i1
 // CHECK:             }
 // CHECK:             scf.yield %[[FINAL_RESULT]] : i1
 
 // -----
 
 func @broadcast(%a : tensor<2xindex>, %b : tensor<3xindex>, %c : tensor<2xindex>) -> !shape.witness {
-  %0 = shape.cstr_broadcastable %a, %b, %c : tensor<2xindex>, tensor<3xindex>, tensor<2xindex>
+  %0 = shape.cstr_broadcastable %a, %b, %c msg("require broadcastable shapes")
+      : tensor<2xindex>, tensor<3xindex>, tensor<2xindex>
   return %0 : !shape.witness
 }
 // CHECK-LABEL:   func @broadcast(
 // CHECK-SAME:          %[[ARG0:.*]]: tensor<2xindex>,
 // CHECK-SAME:          %[[ARG1:.*]]: tensor<3xindex>,
 // CHECK-SAME:          %[[ARG2:.*]]: tensor<2xindex>)
 // CHECK:           %[[C0:.*]] = constant 0 : index
 // CHECK:           %[[C1:.*]] = constant 1 : index

 // CHECK:                %[[EQUALS_1:.*]] = cmpi eq, %[[EXTRACTED:.*]], %c1 : index
 // CHECK:                %[[EQUALS_BROADCASTED:.*]] = cmpi eq, %[[EXTRACTED:.*]], %[[DIM2]] : index
 // CHECK:                %[[GOOD:.*]] = or %[[EQUALS_1:.*]], %[[EQUALS_BROADCASTED:.*]] : i1
 // CHECK:                %[[AND_REDUCTION:.*]] = and %[[SECOND_REDUCTION]], %[[GOOD]] : i1
 // CHECK:                scf.yield %[[AND_REDUCTION]] : i1
 // CHECK:             }
 // CHECK:             scf.yield %[[FINAL_RESULT]] : i1
 
-// CHECK:           %[[RESULT:.*]] = shape.cstr_require %[[ALL_RESULT]], "required broadcastable shapes"
+// CHECK:           %[[RESULT:.*]] = shape.cstr_require %[[ALL_RESULT]], "require broadcastable shapes"
 // CHECK:           return %[[RESULT]] : !shape.witness
 // CHECK:         }
 
 // -----
 
 func @broadcast_3_shapes_different_extents(%a : tensor<2xindex>,
                                            %b : tensor<3xindex>,
                                            %c : tensor<2xindex>) {

mlir/test/Dialect/Shape/canonicalize.mlir

 // -----
 // `assuming_all` with duplicate operands.
 // CHECK-LABEL: func @assuming_all_duplicate_operands
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>)
 func @assuming_all_duplicate_operands(%arg0 : tensor<?xindex>,
     %arg1 : tensor<?xindex>) -> !shape.witness {
   // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]
   // CHECK: return %[[RES]]
-  %0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  %0 = shape.cstr_broadcastable %arg0, %arg1 msg("require broadcastable shapes")
+      : tensor<?xindex>, tensor<?xindex>
   %1 = shape.assuming_all %0, %0, %0
   return %1 : !shape.witness
 }
 
 // -----
 // `assuming_all` with all `cstr_eq` but disjoint operands cannot be collapsed.
 // CHECK-LABEL: func @assuming_all_to_cstr_eq
 // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<3xindex>, %[[D:.*]]: tensor<3xindex>)

 // 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] : !shape.shape
   %cs1 = shape.const_shape [1, 5] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs1 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Empty shape arguments can be removed from broadcastable ops.
 // CHECK-LABEL: func @f
 // CHECK-SAME:  (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>, %{{.*}}: tensor<0xindex>)
 func @f(%arg0 : tensor<?xindex>, %arg1 : tensor<?xindex>, %arg2 : tensor<0xindex>) {
   // CHECK-NOT: const_shape
-  // CHECK: cstr_broadcastable %[[ARG0]], %[[ARG1]] : tensor<?xindex>, tensor<?xindex>
+  // CHECK: cstr_broadcastable %[[ARG0]], %[[ARG1]] msg("{{.*}}") : tensor<?xindex>, tensor<?xindex>
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.cstr_broadcastable %arg0, %arg1, %0, %arg2
+      msg("require broadcastable shapes")
       : tensor<?xindex>, tensor<?xindex>, !shape.shape, tensor<0xindex>
   "consume.witness"(%1) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with non-broadcastable constant shapes is always false
 // CHECK-LABEL: func @static_non_broadcastable
 func @static_non_broadcastable() {
   // 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] : !shape.shape
   %cs1 = shape.const_shape [1, 5] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs1 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   "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] : !shape.shape
-  %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %arg0, %cs0 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   "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 : !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %arg0, %arg0 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Broadcastable canonicalization also works on extent tensors.
 // CHECK-LABEL: func @broadcastable_on_extent_tensors
 func @broadcastable_on_extent_tensors(%arg : tensor<?xindex>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.cstr_broadcastable %arg, %arg : tensor<?xindex>, tensor<?xindex>
+  %0 = shape.cstr_broadcastable %arg, %arg msg("require broadcastable shapes")
+      : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Fold ternary broadcastable
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [8, 1] : !shape.shape
   %cs1 = shape.const_shape [1, 8] : !shape.shape
   %cs2 = shape.const_shape [1, 1] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2
+      msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Fold ternary broadcastable with dynamic ranks
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [8, 1] : !shape.shape
   %cs1 = shape.const_shape [1, -1] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs0, %cs1 : !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs0, %cs1
+      msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // One scalar and one non-scalar and one unknown cannot be broadcasted at compile time
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [8, 1] : !shape.shape
   %cs1 = shape.const_shape [1, 8] : !shape.shape
   %cs2 = shape.const_shape [1, -1] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2
+      msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // One scalar and two unknowns cannot be broadcasted at compile time
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [8, 1] : !shape.shape
   %cs1 = shape.const_shape [1, -1] : !shape.shape
   %cs2 = shape.const_shape [8, -1] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2
+      msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with scalars and a non-scalar can be constant folded
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs0, %arg0 : !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs0, %arg0
+      msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // One scalar and one non-scalar and one unknown cannot be folded.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {
   // CHECK: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [] : !shape.shape
   %cs1 = shape.const_shape [2] : !shape.shape
-  %0 = shape.cstr_broadcastable %cs0, %cs1, %arg0 : !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %cs0, %cs1, %arg0
+      msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Fold `rank` based on constant shape.
 // CHECK-LABEL: @fold_rank

 // Canonicalize scalar cstr_broadcastable checks
 // CHECK-LABEL: @cstr_broadcastable_scalar
 func @cstr_broadcastable_scalar(%arg0 : tensor<?xf32>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
-  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 msg("require broadcastable shapes")
+      : !shape.shape, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Do not canonicalize cstr_broadcastable checks with 2 unknowns
 // CHECK-LABEL: @cstr_broadcastable_unknown
 func @cstr_broadcastable_unknown(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>) {
   // CHECK-NEXT: shape.shape_of %arg0
   // CHECK-NEXT: shape.shape_of %arg1
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
   %1 = shape.shape_of %arg1 : tensor<?xf32> -> tensor<?xindex>
-  %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 msg("require broadcastable shapes")
+      : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Scalars are safe to broadcast to unranked sizes.
 // CHECK-LABEL: @cstr_broadcastable_scalar_unranked
 func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>, %arg1 : tensor<index>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.shape_of %arg1 : tensor<index> -> tensor<?xindex>
   %1 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
-  %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 msg("require broadcastable shapes")
+      : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Fold `shape_eq` for equal and constant shapes.
 // CHECK-LABEL: @shape_eq_fold_1

 }
 
 // -----
 
 // CHECK-LABEL: @cstr_broadcastable_on_duplicate_shapes
 // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)
 func @cstr_broadcastable_on_duplicate_shapes(%a : !shape.shape,
     %b : !shape.shape) -> !shape.witness {
-  // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[A]], %[[B]] :
+  // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[A]], %[[B]] msg("{{.*}}")
   // CHECK: return %[[RES]]
-  %0 = shape.cstr_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape,
-      !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %a, %b, %a, %a, %a, %b
+      msg("require broadcastable shapes") : !shape.shape, !shape.shape,
+      !shape.shape, !shape.shape, !shape.shape, !shape.shape
   return %0 : !shape.witness
 }
 
 // -----
 
 // CHECK-LABEL: @broadcast_on_same_shape
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)
 func @broadcast_on_same_shape(%shape : !shape.shape) -> !shape.shape {

 }
 // ----
 
 // CHECK-LABEL: @cstr_broadcastable_folding
 func @cstr_broadcastable_folding(%arg : tensor<?x4xf32>) {
   // CHECK: const_witness true
   %0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex>
   %1 = shape.const_shape [4] : tensor<1xindex>
-  %2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex>
+  %2 = shape.cstr_broadcastable %0, %1 msg("require broadcastable shapes")
+      : tensor<2xindex>, tensor<1xindex>
   "use"(%2) : (!shape.witness) -> ()
 }
 
 // -----
 
 // CHECK-LABEL: @cast_extent_tensor_operands
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<3xindex>)
 func @cast_extent_tensor_operands(%arg0 : tensor<?xindex>,
     %arg1 : tensor<3xindex>) -> (!shape.witness, tensor<?xindex>) {
   // CHECK: %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?xindex> to tensor<3xindex>
-  // CHECK: %[[WIT:.*]] = shape.cstr_broadcastable %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex>
+  // CHECK: %[[WIT:.*]] = shape.cstr_broadcastable %[[CAST_ARG0]], %[[ARG1]] msg("{{.*}}") : tensor<3xindex>, tensor<3xindex>
   // CHECK: %[[UNCAST_RES:.*]] = shape.broadcast %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex> -> tensor<3xindex>
   // CHECK: %[[RES:.*]] = tensor.cast %[[UNCAST_RES]] : tensor<3xindex> to tensor<?xindex>
   // CHECK: return %[[WIT]], %[[RES]]
   %0 = tensor.cast %arg0 : tensor<?xindex> to tensor<3xindex>
   %1 = tensor.cast %arg1 : tensor<3xindex> to tensor<?xindex>
-  %2 = shape.cstr_broadcastable %0, %1 : tensor<3xindex>, tensor<?xindex>
-  %3 = shape.broadcast %0, %1 :tensor<3xindex>, tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 msg("require broadcastable shapes")
+      : tensor<3xindex>, tensor<?xindex>
+  %3 = shape.broadcast %0, %1 : tensor<3xindex>, tensor<?xindex>
       -> tensor<?xindex>
   return %2, %3 : !shape.witness, tensor<?xindex>
 }
 
 // -----
 
 // CHECK-LABEL: @concretize_broadcast_result_type
 // CHECK-SAME:  (%[[ARG0:.*]]: tensor<2xindex>, %[[ARG1:.*]]: tensor<3xindex>)

mlir/test/Dialect/Shape/invalid.mlir

 
 // expected-error@+1 {{@fn not found}}
 module attributes {shape.lib = @fn} { }
 
 // -----
 
 func @fn(%arg: !shape.shape) -> !shape.witness {
   // expected-error@+1 {{required at least 2 input shapes}}
-  %0 = shape.cstr_broadcastable %arg : !shape.shape
+  %0 = shape.cstr_broadcastable %arg msg("require broadcastable shapes")
+      : !shape.shape
   return %0 : !shape.witness
 }
 
 // -----
 
 // Test that type inference flags the wrong return type.
 
 func @const_shape() {
   // expected-error@+1 {{'tensor<3xindex>' are incompatible with return type(s) of operation 'tensor<2xindex>'}}
   %0 = shape.const_shape [4, 5, 6] : tensor<2xindex>
   return
 }

mlir/test/Dialect/Shape/ops.mlir

func @test_shape_of(%arg0: tensor<?xf32>) -> tensor<?xindex> {
   %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
 
 func @test_constraints() {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %true = constant true
-  %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
+  %w0 = shape.cstr_broadcastable %0, %1 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   %w1 = shape.cstr_eq %0, %1 : !shape.shape, !shape.shape
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false
   %w4 = shape.cstr_require %true, "msg"
   %w_all = shape.assuming_all %w0, %w1, %w2, %w3, %w4
   shape.assuming %w_all -> !shape.shape {
     %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
 
 func @eq_on_extent_tensors(%lhs : tensor<?xindex>,
                            %rhs : tensor<?xindex>) {
   %w0 = shape.cstr_eq %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
   return
 }
 
 func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
                                       %rhs : tensor<?xindex>) {
-  %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
+  %w0 = shape.cstr_broadcastable %lhs, %rhs
+      msg("require broadcastable shapes")
+      : tensor<?xindex>, tensor<?xindex>
   return
 }
 
 func @mul(%size_arg : !shape.size, %index_arg : index) {
   %size_prod = shape.mul %size_arg, %size_arg
       : !shape.size, !shape.size -> !shape.size
   %index_prod = shape.mul %index_arg, %index_arg : index, index -> index
   %mixed_prod = shape.mul %size_arg, %index_arg

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints -canonicalize <%s | FileCheck %s --dump-input=fail --check-prefixes=CANON,CHECK-BOTH
 // RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints <%s | FileCheck %s --dump-input=fail --check-prefixes=REPLACE,CHECK-BOTH
 
 // -----
 // Check that cstr_broadcastable is removed.
 //
 // CHECK-BOTH: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
   // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
   // REPLACE-NOT: shape.cstr_eq
   // REPLACE: shape.assuming %[[WITNESS]]
   // CANON-NEXT: test.source
   // CANON-NEXT: return
-  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %arg0, %arg1 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   %1 = shape.assuming %0 -> index {
     %2 = "test.source"() : () -> (index)
     shape.assuming_yield %2 : index
   }
   return %1 : index
 }
 
 // -----

 // -----
 // With a non-const value, we cannot fold away the code, but all constraints
 // should be removed still.
 //
 // CHECK-BOTH: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
   // CANON-NEXT: test.source
   // CANON-NEXT: return
-  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
+  %0 = shape.cstr_broadcastable %arg0, %arg1 msg("require broadcastable shapes")
+      : !shape.shape, !shape.shape
   %1 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape
   %2 = shape.assuming_all %0, %1
   %3 = shape.assuming %0 -> index {
     %4 = "test.source"() : () -> (index)
     shape.assuming_yield %4 : index
   }
   return %3 : index
 }