[mlir][tosa] Add folder for multiply like reduce_prod operation #128067
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Member
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-tosa Author: Tai Ly (Tai78641) ChangesThis commit uses mulBinaryFolder for reduce_prod operations that have a constant 1D input of two values. Full diff: https://github.com/llvm/llvm-project/pull/128067.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 9bfc2aae1d6a5..45bcacff19caa 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -963,10 +963,34 @@ REDUCE_FOLDER(ReduceAllOp)
REDUCE_FOLDER(ReduceAnyOp)
REDUCE_FOLDER(ReduceMaxOp)
REDUCE_FOLDER(ReduceMinOp)
-REDUCE_FOLDER(ReduceProdOp)
REDUCE_FOLDER(ReduceSumOp)
#undef REDUCE_FOLDER
+OpFoldResult ReduceProdOp::fold(FoldAdaptor adaptor) {
+ ShapedType inputTy = llvm::cast<ShapedType>(getInput().getType());
+ if (!inputTy.hasRank())
+ return {};
+ if (inputTy == getType() && (inputTy.getRank() == 0 || inputTy.getDimSize(getAxis()) == 1))
+ return getInput();
+
+ // Fold multiply like reduce_prod operators using mulBinaryFolder
+ if (inputTy.getRank() == 1 && inputTy.getDimSize(0) == 2) {
+ const auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
+ if (!resultTy)
+ return {};
+
+ const auto elements = llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput());
+ if (!elements)
+ return {};
+
+ const auto lhsAttr = DenseElementsAttr::get(resultTy, {elements.getValues<Attribute>()[0]});
+ const auto rhsAttr = DenseElementsAttr::get(resultTy, {elements.getValues<Attribute>()[1]});
+ return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, 0);
+ }
+
+ return {};
+}
+
OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
auto inputTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 0e177a076ee7a..316f22f88fc69 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1012,3 +1012,14 @@ func.func nested @do_not_fold_reciprocal_int() -> tensor<3x600x1200xi32> {
%2 = "tosa.reciprocal"(%1): (tensor<3x600x1200xi32>) -> tensor<3x600x1200xi32>
return %2 : tensor<3x600x1200xi32>
}
+
+// -----
+
+// CHECK-LABEL: @fold_reduce_prod_is_mul
+func.func @fold_reduce_prod_is_mul() -> tensor<1xi32> {
+ // CHECK-DAG: %[[VAL_0:.*]] = "tosa.const"() <{value = dense<77> : tensor<1xi32>}> : () -> tensor<1xi32>
+ // CHECK: return %[[VAL_0]] : tensor<1xi32>
+ %0 = "tosa.const"() <{value = dense<[1, 77]> : tensor<2xi32>}> : () -> tensor<2xi32>
+ %1 = "tosa.reduce_prod"(%0) <{axis = 0 : i32}> : (tensor<2xi32>) -> tensor<1xi32>
+ return %1 : tensor<1xi32>
+}
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
Tai78641
force-pushed
the
pr_add_mul_folder
branch
2 times, most recently
from
a4b507e to
62b2409
Compare
Jerry-Ge
changed the title
GeorgeARM
reviewed
Feb 24, 2025
| (inputTy.getRank() == 0 || inputTy.getDimSize(getAxis()) == 1)) | ||
| return getInput(); | ||
|
|
||
| // Fold multiply like reduce_prod operators using mulBinaryFolder |
Contributor
There was a problem hiding this comment.
This is quite limited special-cased. What is the motivation behind this?
Contributor
There was a problem hiding this comment.
Agreed it's quite special cased, it was observed while legalizing a model with dynamic shapes
FranklandJack
requested changes
Feb 27, 2025
Contributor
FranklandJack
left a comment
FranklandJack
left a comment
There was a problem hiding this comment.
Should we remove the change ID from the commit message?
| OpFoldResult ReduceProdOp::fold(FoldAdaptor adaptor) { | ||
| ShapedType inputTy = llvm::cast<ShapedType>(getInput().getType()); | ||
| if (!inputTy.hasRank()) | ||
| return {}; |
Contributor
There was a problem hiding this comment.
Could we add a newline after the return here to aid readability?
| #undef REDUCE_FOLDER | ||
|
|
||
| OpFoldResult ReduceProdOp::fold(FoldAdaptor adaptor) { | ||
| ShapedType inputTy = llvm::cast<ShapedType>(getInput().getType()); |
Contributor
There was a problem hiding this comment.
I think we can use auto here since the type is already explicit on the RHS.
| return getInput(); | ||
|
|
||
| // Fold multiply like reduce_prod operators using mulBinaryFolder | ||
| if (inputTy.getRank() == 1 && inputTy.getDimSize(0) == 2) { |
Contributor
There was a problem hiding this comment.
Could we invert the condition here and do an early exit to avoid the indention and big conditional block?
| return %2 : tensor<3x600x1200xi32> | ||
| } | ||
|
|
||
| // ----- |
Contributor
There was a problem hiding this comment.
Could we add some negative tests here? From the above logic it seems like the following cases should fail:
- input tensor of rank > 1
- input tensor of rank 1 with dim[0] > 1
Tai78641
force-pushed
the
pr_add_mul_folder
branch
from
62b2409 to
835c8db
Compare
Tai78641
force-pushed
the
pr_add_mul_folder
branch
3 times, most recently
from
98ff8a5 to
b2026c2
Compare
This commit uses mulBinaryFolder for reduce_prod operations that have a constant 1D input of two values. Change-Id: Icb234282c70898189083231506ed38a3ab40efb2 Signed-off-by: Luke Hutton <[email protected]>
Tai78641
force-pushed
the
pr_add_mul_folder
branch
from
b2026c2 to
a8fd688
Compare
Contributor
|
No longer needed |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit uses mulBinaryFolder for reduce_prod operations that have a constant 1D input of two values.