[mlir] Broadcastable relax verification.

[jpienaar]

Summary

Relax verification to enable specialization (e.g., shape inference and
the like) to be run on a verified IR. This moves to verify failing only
for what is known to be incorrect, but the expectation of the trait is
unchanged.

[rafaelubalmw]

Hi Jacques. I was addressing these issues this morning myself. I submitted a separate patch with these and some additional changes here:

https://reviews.llvm.org/D156714

I will leave a comment in the RFC as well, as I no longer seem to be able to edit it:

https://discourse.llvm.org/t/tosa-to-linalg-lowering-of-element-wise-ops/71559

mlir/lib/Dialect/Traits.cpp

static bool isCompatibleInferredReturnShape(ArrayRef<int64_t> inferred,
                                             ArrayRef<int64_t> existing) {
   auto isCompatible = [](int64_t inferredDim, int64_t existingDim) {
     // The following criterion is used to determine the validity of an existing
     // dimension:
     //
     // inferredDim  existingDim  Behavior
     // -----------  -----------  --------
     // dynamic      dynamic      OK
-    // dynamic      static       Error
+    // dynamic      static       OK^1
     // static       dynamic      OK
     // static       static       OK if equal
-    return ShapedType::isDynamic(existingDim) || inferredDim == existingDim;
+    //
+    // [1] This allows for implicit dynamic-to-static cast with undefined
+    // behavior if the resulting runtime dimension size does not match the given
+    // static size. This is meant to only flag known invalid cases while the
+    // expectation is that these should be equal dynamically.
+    return ShapedType::isDynamic(existingDim) ||
+           ShapedType::isDynamic(inferredDim) || inferredDim == existingDim;
   };
   if (inferred.size() != existing.size())
     return false;
   for (auto [inferredDim, existingDim] : llvm::zip(inferred, existing))
     if (!isCompatible(inferredDim, existingDim))
       return false;
   return true;
 }

mlir/test/Dialect/traits.mlir

 func.func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<?xi32>) -> tensor<4x3x2xi32> {
 ^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<?xi32>):
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<?xi32>) -> tensor<4x3x2xi32>
   return %0 : tensor<4x3x2xi32>
 }
 
 // -----
 
-// Error for inferred dynamic dimension but existing static dimensions
+// Error for inferred dynamic dimension but existing static dimensions, this
+// has moved to runtime UB if not matching.
 func.func @broadcast_tensor_tensor_tensor(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<2xi32> {
-  // expected-error @+1 {{op result type '2' not broadcast compatible with broadcasted operands's shapes '?'}}
+  // not-error @+1 {{op result type '2' not broadcast compatible with broadcasted operands's shapes '?'}}
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<?xi32>, tensor<?xi32>) -> tensor<2xi32>
   return %0 : tensor<2xi32>
 }
 
 // -----
 
 func.func @broadcast_tensor_tensor_tensor(%arg0: tensor<?x6x1xi32>, %arg1: tensor<*xi32>) -> tensor<?x6x?xi32> {
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<?x6x1xi32>, tensor<*xi32>) -> tensor<?x6x?xi32>