[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 72b7eaf7ea5b0b43984951226855d83a3e991937 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
---
 .../Dialect/Tosa/Utils/ConversionUtils.cpp    | 21 ++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
index d1a8732dac212..3a51939e07b5b 100644
--- a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
+++ b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
@@ -77,24 +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 = lowerRank - 1; i >= 0; i--) {
+    higherRankDim = higherRankShape[i + rankDiff];
+    lowerRankDim = lowerRankShape[i];
 
-  for (int64_t i = higherRank - 1, j = lowerRank - 1; i >= 0 && j >= 0;
-       i--, j--) {
-    higherRankDim = higherRankShape[i];
-    lowerRankDim = lowerRankShape[j];
-
-    if (lowerRankDim == 1 && higherRankDim > 1)
-      reshapeOutputShape[i] = 1;
-    else if ((lowerRankDim > 1 && higherRankDim == 1) ||
-             (lowerRankDim == higherRankDim))
-      reshapeOutputShape[i] = lowerRankDim;
-    else if (higherRankDim != lowerRankDim)
+    if (lowerRankDim != 1 && higherRankDim != 1 &&
+        lowerRankDim != higherRankDim)
       return failure();
+
+    reshapeOutputShape[i + rankDiff] = lowerRankDim == 1 ? 1 : lowerRankDim;
   }
   return success();
 }