[Mlir-commits] [mlir] [mlir][tosa] Fix ability to expand ranks with dynamic shape support (PR #128037)

[llvmlistbot at llvm.org]

https://github.com/Jerry-Ge updated https://github.com/llvm/llvm-project/pull/128037

>From 16462f491dfc368604a0675699d9781ea4f30eeb Mon Sep 17 00:00:00 2001
From: Suraj Sudhir <suraj.sudhir at arm.com>
Date: Thu, 21 Mar 2024 23:02:32 +0000
Subject: [PATCH] [mlir][tosa] Fix ability to expand ranks with dynamic shape
 support

- The use of != 1 accommodates the use of kDynamicDim
- Simplified the for loop to iterate only on lowerRank and access the
  higherRank dim by using the rankDiff

Signed-off-by: Suraj Sudhir <suraj.sudhir at arm.com>
Change-Id: I0f223f335667b2e32c43d4370f0a4b11b0617694
---
 mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
index d1a8732dac212..a6df4694bc851 100644
--- a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
+++ b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
@@ -77,22 +77,21 @@ computeReshapeOutput(ArrayRef<int64_t> higherRankShape,
   // Initialize new shapes with [1] * higherRank.
   int64_t higherRank = higherRankShape.size();
   int64_t lowerRank = lowerRankShape.size();
-
   reshapeOutputShape.assign(higherRank, 1);
 
   int64_t higherRankDim;
   int64_t lowerRankDim;
+  const int64_t rankDiff = higherRank - lowerRank;
 
-  for (int64_t i = higherRank - 1, j = lowerRank - 1; i >= 0 && j >= 0;
-       i--, j--) {
-    higherRankDim = higherRankShape[i];
-    lowerRankDim = lowerRankShape[j];
+  for (int64_t i = lowerRank - 1; i >= 0; i--) {
+    higherRankDim = higherRankShape[i + rankDiff];
+    lowerRankDim = lowerRankShape[i];
 
-    if (lowerRankDim == 1 && higherRankDim > 1)
-      reshapeOutputShape[i] = 1;
-    else if ((lowerRankDim > 1 && higherRankDim == 1) ||
+    if (lowerRankDim == 1 && higherRankDim != 1)
+      reshapeOutputShape[i + rankDiff] = 1;
+    else if ((lowerRankDim != 1 && higherRankDim == 1) ||
              (lowerRankDim == higherRankDim))
-      reshapeOutputShape[i] = lowerRankDim;
+      reshapeOutputShape[i + rankDiff] = lowerRankDim;
     else if (higherRankDim != lowerRankDim)
       return failure();
   }