[Mlir-commits] [mlir] [tosa] : Enhance EqualizeRanks to handle dynamic dimensions. (PR #168564)

[Sayan Saha]

https://github.com/sahas3 created https://github.com/llvm/llvm-project/pull/168564

Legalizing following IR to `tosa` using `tf-tosa-opt` from `tensorflow` repo:
```
func.func @main(%arg0: tensor<?x?x?x?xf32>) -> tensor<?x?x?x5xf32> {
    %0 = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<5xf32>}> : () -> tensor<5xf32>
    %1 = tfl.add(%arg0, %0) <{fused_activation_function = "NONE"}> : (tensor<?x?x?x?xf32>, tensor<5xf32>) -> tensor<?x?x?x5xf32>
    return %1 : tensor<?x?x?x5xf32>
  }
```
fails with
```
error: 'tosa.add' op operands don't have matching ranks
    %1 = tfl.add(%arg0, %0) <{fused_activation_function = "NONE"}> : (tensor<?x?x?x?xf32>, tensor<5xf32>) -> tensor<?x?x?x5xf32>
         ^
tfl.mlir:3:10: note: see current operation: %1 = "tosa.add"(%arg0, %0) : (tensor<?x?x?x?xf32>, tensor<5xf32>) -> tensor<?x?x?x5xf32>
// -----// IR Dump After TosaLegalizeTFLPass Failed (tosa-legalize-tfl) //----- //
"func.func"() <{function_type = (tensor<?x?x?x?xf32>) -> tensor<?x?x?x5xf32>, sym_name = "main"}> ({
^bb0(%arg0: tensor<?x?x?x?xf32>):
  %0 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<5xf32>}> : () -> tensor<5xf32>
  %1 = "tosa.add"(%arg0, %0) : (tensor<?x?x?x?xf32>, tensor<5xf32>) -> tensor<?x?x?x5xf32>
  "func.return"(%1) : (tensor<?x?x?x5xf32>) -> ()
}) : () -> ()
```

This is because of the following check in `computeReshapeOutput` called from `EqualizeRanks` function:
```
if (lowerRankDim != 1 && higherRankDim != 1 &&
        lowerRankDim != higherRankDim)
      return failure();
```

Based on the broadcast semantics defined in https://mlir.llvm.org/docs/Traits/Broadcastable/#dimension-inference I think it's legal to allow `lowerRankDim != higherRankDim` if one of them is dynamic. At runtime verifier should enforce that 
1. if lowerRankDim is dynamic and higherRankDim is static then the dynamic dim matches the static dim and vice-versa
2. if both are dynamic, they should match
It's not necessary to error out during the op construction time.

>From 54c3b3b4375195914c369f8ed631e73cbd2d5177 Mon Sep 17 00:00:00 2001
From: Sayan Saha <sayans at mathworks.com>
Date: Tue, 18 Nov 2025 11:23:09 -0500
Subject: [PATCH] [tosa] : Enhance EqualizeRanks to handle dynamic dimensions.

---
 mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
index 62c015a85ee36..bb52d15026367 100644
--- a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
+++ b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
@@ -70,6 +70,8 @@ namespace {
 // If lower=[a], higher=[a, a], [a] reshaped into [1, a].
 // If lower=[a], target=[a, b, a], [a] reshaped into [1, 1, a].
 // If lower=[], target=[a, b, c], [] reshaped into [1, 1, 1].
+// If lower=[c], higher=[?, ?, c], [c] reshaped into [1, 1, c].
+// If lower=[?], higher=[?, ?, ?], [?] reshaped into [1, 1, ?].
 LogicalResult
 computeReshapeOutput(ArrayRef<int64_t> higherRankShape,
                      ArrayRef<int64_t> lowerRankShape,
@@ -87,7 +89,12 @@ computeReshapeOutput(ArrayRef<int64_t> higherRankShape,
     higherRankDim = higherRankShape[i + rankDiff];
     lowerRankDim = lowerRankShape[i];
 
-    if (lowerRankDim != 1 && higherRankDim != 1 &&
+    auto isKnownStaticShapeNotEqualToOne = [](int64_t dim) {
+      return dim != 1 && dim != ShapedType::kDynamic;
+    };
+
+    if (isKnownStaticShapeNotEqualToOne(lowerRankDim) &&
+        isKnownStaticShapeNotEqualToOne(higherRankDim) &&
         lowerRankDim != higherRankDim)
       return failure();