From 72b7eaf7ea5b0b43984951226855d83a3e991937 Mon Sep 17 00:00:00 2001 From: Suraj Sudhir 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 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 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(); }