[Mlir-commits] [mlir] d20b90a - [mlir][tosa] Enhance `resolveBroadcastShape` to handle dynamic dims (#175932)

[llvmlistbot at llvm.org]

Author: Luke Hutton
Date: 2026-02-04T14:15:08Z
New Revision: d20b90a2da8ab69ccf846a450e0beadea70cf6f8

URL: https://github.com/llvm/llvm-project/commit/d20b90a2da8ab69ccf846a450e0beadea70cf6f8
DIFF: https://github.com/llvm/llvm-project/commit/d20b90a2da8ab69ccf846a450e0beadea70cf6f8.diff

LOG: [mlir][tosa] Enhance `resolveBroadcastShape` to handle dynamic dims (#175932)

This commit improves `resolveBroadcastShape` when some input dims are
dynamic. Previously, this would result in shape inference failing,
meaning the output type of the operation would not be updated. Now the
output shape is correctly inferred, selecting the static dim if > 1,
otherwise selecting the dynamic dim.

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
    mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 5438cd73aef54..fe107699cc3d3 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -1481,6 +1481,19 @@ static void buildVariableOp(OpBuilder &builder, OperationState &result,
 //===----------------------------------------------------------------------===//
 // TOSA Operator Return Type Inference.
 //===----------------------------------------------------------------------===//
+static FailureOr<int64_t> resolveBroadcastDim(const int64_t dim1,
+                                              const int64_t dim2) {
+  if (dim1 == 1)
+    return dim2;
+  if (dim2 == 1)
+    return dim1;
+
+  if (ShapedType::isStatic(dim1) && ShapedType::isStatic(dim2) && dim1 != dim2)
+    return failure();
+
+  // Prefer static dimension over dynamic
+  return ShapedType::isDynamic(dim1) ? dim2 : dim1;
+}
 
 static LogicalResult resolveBroadcastShape(const ValueShapeRange &operands,
                                            SmallVector<int64_t> &outShape) {
@@ -1504,15 +1517,12 @@ static LogicalResult resolveBroadcastShape(const ValueShapeRange &operands,
     for (size_t i = 0, e = shape.getRank(); i < e; ++i) {
       auto dim1 = outShape[i + rankDiff];
       auto dim2 = shape.getDimSize(i);
-      auto resolvedDim = dim1;
 
-      if (dim1 == 1) {
-        resolvedDim = dim2;
-      } else if (dim2 == 1) {
-        resolvedDim = dim1;
-      } else if (dim1 != dim2) {
+      const FailureOr<int64_t> maybeResolvedDim =
+          resolveBroadcastDim(dim1, dim2);
+      if (failed(maybeResolvedDim))
         return failure();
-      }
+      const int64_t resolvedDim = *maybeResolvedDim;
       outShape[i + rankDiff] = resolvedDim;
     }
   }

diff  --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index bc5f41b1af304..0deb31dab1c6f 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -251,6 +251,24 @@ func.func @test_binary_i1(%arg0 : tensor<4xi1>, %arg1 : tensor<1xi1>) -> () {
 
 // -----
 
+// CHECK-LABEL: @test_dynamic_binary_broadcast
+func.func @test_dynamic_binary_broadcast(%arg0: tensor<1x?xf32>, %arg1: tensor<2x1xf32>) -> tensor<*xi1> {
+  // CHECK tosa.equal %arg0, %arg1 : (tensor<1x?xf32>, tensor<2x1xf32>) -> tensor<2x?xi1>
+  %0 = tosa.equal %arg0, %arg1 : (tensor<1x?xf32>, tensor<2x1xf32>) -> tensor<*xi1>
+  return %0 : tensor<*xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_resolvable_dynamic_binary_broadcast
+func.func @test_resolvable_dynamic_binary_broadcast(%arg0: tensor<1x?xf32>, %arg1: tensor<2x4xf32>) -> tensor<*xi1> {
+  // CHECK tosa.equal %arg0, %arg1 : (tensor<1x?xf32>, tensor<2x4xf32>) -> tensor<2x4xi1>
+  %0 = tosa.equal %arg0, %arg1 : (tensor<1x?xf32>, tensor<2x4xf32>) -> tensor<*xi1>
+  return %0 : tensor<*xi1>
+}
+
+// -----
+
 // CHECK-LABEL: @test_select_i32
 func.func @test_select_i32(%arg0 : tensor<4xi1>, %arg1 : tensor<1xi32>, %arg2 : tensor<4xi32>) -> () {
   // CHECK: tosa.select %arg0, %arg1, %arg2 : (tensor<4xi1>, tensor<1xi32>, tensor<4xi32>) -> tensor<4xi32>