Commit 6ad1623
authored
[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.1 parent 87a1fd1 commit 6ad1623
1 file changed
+8
-1
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
70 | 70 | | |
71 | 71 | | |
72 | 72 | | |
| 73 | + | |
| 74 | + | |
73 | 75 | | |
74 | 76 | | |
75 | 77 | | |
| |||
87 | 89 | | |
88 | 90 | | |
89 | 91 | | |
90 | | - | |
| 92 | + | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
91 | 98 | | |
92 | 99 | | |
93 | 100 | | |
| |||
0 commit comments