Released restriction that prevented implicit dynamic-to-static dimension type cast in TOSA ops.

[jpienaar]

Thanks! (Let me know if you need help landing)

[rafaelubalmw]

@jpienaar Thanks, Jacques. Yes, please, it'd be great if you could land this.

mlir/docs/Traits/Broadcastable.md

while len(shape) < rank:
     shape.prepend(1)
 ```
     
 Given the shapes of two ranked input operands, the result's shape is inferred by equalizing input ranks and inferring individual dimensions, as shown here:
 
 ```python
 InferShape(shape0, shape1):
 
-  # Equalize ranks
-  rank = max(GetRank(shape0), GetRank(shape1))
-  ExpandRank(shape0, rank)
-  ExpandRank(shape1, rank)
+  # Equalize ranks
+  rank = max(GetRank(shape0), GetRank(shape1))
+  ExpandRank(shape0, rank)
+  ExpandRank(shape1, rank)
   
-  # Infer shape
-  inferredShape = []
-  for (dim0, dim1) in zip(shape0, shape1):
-    inferredDim = InferDim(dim0, dim1)
-        inferredShape.append(inferredDim)
-  return inferredShape
+  # Infer shape
+  inferredShape = []
+  for (dim0, dim1) in zip(shape0, shape1):
+    inferredDim = InferDim(dim0, dim1)
+  return inferredShape
 ```
   
 The result shape for an operation with an arbitrary number of input operands is then inferred by discarding unranked operands, applying shape inference on the first ranked operand pair, and updating the inferred shape with each additional ranked operand. If the operation has no ranked operands, the result shape cannot be inferred. If the operation has exactly one ranked operand, its shape is directly provided as the inferred result shape. Formally:
 
 ```python
 InferResultShape(op):
 
   # Filter ranked operands

 If either the result is unranked or all input operands are unranked, no further verification steps are needed, and the process ends here successfully. If, on the contrary, both the result and at least one input operand are ranked, verification continues by checking for a matching rank between the previously inferred shape and the result.
 
 Once a rank match is guaranteed, each dimension of the inferred shape is compared with the corresponding dimension of the actual result shape according to the following table table:
 
 
 | `inferredDim` | `actualDim` | Verification outcome |
 | ------------- | ----------- | -------------------- |
 | ? | ? | **OK** |
-| ? | static | **Error** <br> An inferred dimension being dynamic indicates that its size cannot be inferred at compile time from its input operands. The presence of a static dimension in the actual result is counterintuitive and is therefore not allowed. |
+| ? | static | **OK** <br> A failure to guarantee that the runtime dimension size of the result is equal to `actualDim` causes undefined behavior. While unusual, this implicit dynamic-to-static cast is convenient in certain scenarios, such as an intermediate state of a shape inference pass. Ultimately, a static dimension in the result implies that all input dimension sizes are also known at compile time and may therefore become static as well, preferably. |
 | static | ? | **OK** <br> The actual result dimension may be dynamic even when a static size can be inferred at compile time. The programmer may choose to relax the specificity of the result dimension for forward compatibility of the result type. |
 | static | static | **OK if equal** <br> When both the inferred and actual dimensions are static, they must be set to the same size. |
 
 
 The full verification process can be formally specified as follows:
 
 ```python
 Verify(op):
 
   # Run shape inference
   inferredShape = InferResultShape(op.operands)
 
   # Done if result is unranked or all operands are unranked
   if not IsRanked(op.result) or inferredShape is None:
     return
   
   # Rank must match
   actualShape = GetShape(op.result):
   ERROR_IF(len(inferredShape) != len(actualShape))
   
   # Verify
   for (inferredDim, actualDim) in zip(inferredShape, actualShape):
-    ERROR_IF(IsDynamic(inferredDim) and IsStatic(actualDim))
     ERROR_IF(IsStatic(actualDim) and inferredDim != actualDim)
 ```
     
 ## Examples
 
 The following are correct uses of broadcastable ops:
 
 ```mlir

 // The inferred result shape is tensor<2xi32>, but the actual result shape is
 // tensor<4xi32>, which is not compatible.
 %result = "test.broadcastable"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32) -> tensor<4xi32>
 
 // The inferred result shape is tensor<1xi32>, but the actual result shape is
 // tensor<4xi32>. Broadcast semantics are not applicable for results.
 %result = "test.broadcastable"(%arg0, %arg1) : (tensor<1xi32>, tensor<1xi32) -> tensor<4xi32>
 ```
+
+

mlir/lib/Dialect/Traits.cpp

 static std::tuple<bool, bool> hasTensorOrVectorType(iterator_range types) {
   return std::make_tuple(
       llvm::any_of(types, [](Type t) { return isa<TensorType>(t); }),
       llvm::any_of(types, [](Type t) { return isa<VectorType>(t); }));
 }
 
 static bool isCompatibleInferredReturnShape(ArrayRef<int64_t> inferred,
                                             ArrayRef<int64_t> existing) {
+  // If both interred and existing dimensions are static, they must be equal.
   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
-    // static       dynamic      OK
-    // static       static       OK if equal
-    return ShapedType::isDynamic(existingDim) || inferredDim == existingDim;
+    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/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func.func @test_abs_scalar(%arg0: tensor<f32>) -> tensor<f32> {
 
   // CHECK: return [[GENERIC]] : tensor<f32>
   return %0 : tensor<f32>
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
-// CHECK-LABEL: @test_abs_1d_cast_result
+// CHECK-LABEL: @test_abs_1d_cast_static_to_dynamic
 // CHECK-SAME: ([[ARG0:%[0-9a-zA-Z_]*]]
-func.func @test_abs_1d_cast_result(%arg0: tensor<5xf32>) -> tensor<?xf32> {
+func.func @test_abs_1d_cast_static_to_dynamic(%arg0: tensor<5xf32>) -> tensor<?xf32> {
   // CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<5xf32>
   // CHECK: [[RESULT:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins([[ARG0]] : tensor<5xf32>) outs([[EMPTY]] : tensor<5xf32>) {
   // CHECK: ^bb0([[IN0:%.+]]: f32, [[OUT0:%.+]]: f32):
   // CHECK:   [[ABS:%.+]] = math.absf [[IN0]] : f32
   // CHECK:   linalg.yield [[ABS]] : f32
   // CHECK: } -> tensor<5xf32>
+  // CHECK: [[CAST_RESULT:%.+]] = tensor.cast [[RESULT]] : tensor<5xf32> to tensor<?xf32>
   %0 = "tosa.abs"(%arg0) : (tensor<5xf32>) -> tensor<?xf32>
 
-  // CHECK: [[CAST_RESULT:%.+]] = tensor.cast [[RESULT]] : tensor<5xf32> to tensor<?xf32>
   // CHECK: return [[CAST_RESULT]] : tensor<?xf32>
   return %0 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: @test_abs_1d_cast_dynamic_to_static
+// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
+func.func @test_abs_1d_cast_dynamic_to_static(%arg0: tensor<?xf32>) -> tensor<5xf32> {
+  // CHECK: %[[ZERO:.*]] = arith.constant 0 : index
+  // CHECK: %[[DIM_SIZE:.*]] = tensor.dim %[[ARG0]], %[[ZERO]] : tensor<?xf32>
+  // CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM_SIZE]]) : tensor<?xf32>
+  // CHECK: %[[RESULT:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%[[ARG0]] : tensor<?xf32>) outs(%[[EMPTY]] : tensor<?xf32>) {
+  // CHECK: ^bb0(%[[VAL_0:.*]]: f32, %[[VAL_1:.*]]: f32):
+  // CHECK:   %[[VAL_2:.*]] = math.absf %[[VAL_0]] : f32
+  // CHECK:   linalg.yield %[[VAL_2]] : f32
+  // CHECK: } -> tensor<?xf32>
+  // CHECK: %[[CAST_RESULT:.*]] = tensor.cast %[[RESULT]] : tensor<?xf32> to tensor<5xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<?xf32>) -> tensor<5xf32>
+
+  // CHECK: return %[[CAST_RESULT]] : tensor<5xf32>
+  return %0 : tensor<5xf32>
+}
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 // CHECK-LABEL: @test_abs_1d_dynamic
 // CHECK-SAME: ([[ARG0:%[0-9a-zA-Z_]*]]
 func.func @test_abs_1d_dynamic(%arg0: tensor<?xf32>) -> tensor<?xf32> {
 
   // CHECK: [[ZERO:%.+]] = arith.constant 0 : index
   // CHECK: [[DIM:%.+]] = tensor.dim [[ARG0]], [[ZERO]] : tensor<?xf32>
   // CHECK: [[EMPTY:%.+]] = tensor.empty([[DIM]]) : tensor<?xf32>
   // CHECK: [[RESULT:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<?xf32>) outs([[EMPTY]] : tensor<?xf32>) {

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
+// It is alright to have an implicit dynamic-to-static cast in a dimension size
+// as long as the runtime result size is consistent with the result tensor's
+// static dimension.
 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 '?'}}
   %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>