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

[llvmlistbot at llvm.org]

Author: Sayan Saha
Date: 2025-11-19T12:48:53-05:00
New Revision: 6ad162393cf8ab2989e158576877e4570e091bbf

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

LOG: [tosa] : Enhance EqualizeRanks to handle dynamic dimensions. (#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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
index 62c015a85ee36..36e89405210a6 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 isStaticDimAndNotEqualToOne = [](int64_t dim) {
+      return dim != 1 && dim != ShapedType::kDynamic;
+    };
+
+    if (isStaticDimAndNotEqualToOne(lowerRankDim) &&
+        isStaticDimAndNotEqualToOne(higherRankDim) &&
         lowerRankDim != higherRankDim)
       return failure();