[MLIR] TosaToLinalg: Prefer to emit identity maps #123295
Merged
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
mgehre-amd
changed the title
mgehre-amd
changed the title
|
@llvm/pr-subscribers-mlir-tosa Author: Matthias Gehre (mgehre-amd) ChangesWhen deciding whether to emit a map like This helps later transformations which often require the affine map to be a projected permuatation. Full diff: https://github.com/llvm/llvm-project/pull/123295.diff 2 Files Affected:
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 9295afd36e3ab1..a183c27abf62ae 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -882,8 +882,14 @@ emitElementwiseComputation(ConversionPatternRewriter &rewriter, Location loc,
auto shape = cast<ShapedType>(operand.getType()).getShape();
SmallVector<AffineExpr> affineExprs;
for (auto it : llvm::enumerate(shape)) {
- auto affineExpr = it.value() == 1 ? rewriter.getAffineConstantExpr(0)
- : rewriter.getAffineDimExpr(it.index());
+ // Prefer producting identity maps whenever possible (i.e. no broadcasting
+ // needed) because some transforms (like reshape folding)
+ // do not support affine constant exprs.
+ bool requiresBroadcast =
+ (it.value() == 1 && resultType.getDimSize(it.index()) != 1);
+ auto affineExpr = requiresBroadcast
+ ? rewriter.getAffineConstantExpr(0)
+ : rewriter.getAffineDimExpr(it.index());
affineExprs.push_back(affineExpr);
}
return AffineMap::get(rank, 0, affineExprs, rewriter.getContext());
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 1d235092b71d55..f36f449da8dbc4 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -253,6 +253,26 @@ func.func @test_add_1d_broadcast_static_to_static(%arg0: tensor<1xf32>, %arg1: t
// -----
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: @test_add_1d_matching_no_broadcast
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+func.func @test_add_1d_matching_no_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
+
+ // CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<1xf32>
+ // CHECK: %[[RESULT:.*]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[ARG1]] : tensor<1xf32>, tensor<1xf32>) outs(%[[VAL_0]] : tensor<1xf32>) {
+ // CHECK: ^bb0(%[[VAL_1:.*]]: f32, %[[VAL_2:.*]]: f32, %[[VAL_3:.*]]: f32):
+ // CHECK: %[[VAL_4:.*]] = arith.addf %[[VAL_1]], %[[VAL_2]] : f32
+ // CHECK: linalg.yield %[[VAL_4]] : f32
+ // CHECK: } -> tensor<1xf32>
+ %0 = tosa.add %arg0, %arg1 : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+
+ // CHECK: return %[[RESULT]] : tensor<1xf32>
+ return %0 : tensor<1xf32>
+}
+
+// -----
+
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_add_1d_matching_static
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
|
|
@llvm/pr-subscribers-mlir Author: Matthias Gehre (mgehre-amd) ChangesWhen deciding whether to emit a map like This helps later transformations which often require the affine map to be a projected permuatation. Full diff: https://github.com/llvm/llvm-project/pull/123295.diff 2 Files Affected:
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 9295afd36e3ab1..a183c27abf62ae 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -882,8 +882,14 @@ emitElementwiseComputation(ConversionPatternRewriter &rewriter, Location loc,
auto shape = cast<ShapedType>(operand.getType()).getShape();
SmallVector<AffineExpr> affineExprs;
for (auto it : llvm::enumerate(shape)) {
- auto affineExpr = it.value() == 1 ? rewriter.getAffineConstantExpr(0)
- : rewriter.getAffineDimExpr(it.index());
+ // Prefer producting identity maps whenever possible (i.e. no broadcasting
+ // needed) because some transforms (like reshape folding)
+ // do not support affine constant exprs.
+ bool requiresBroadcast =
+ (it.value() == 1 && resultType.getDimSize(it.index()) != 1);
+ auto affineExpr = requiresBroadcast
+ ? rewriter.getAffineConstantExpr(0)
+ : rewriter.getAffineDimExpr(it.index());
affineExprs.push_back(affineExpr);
}
return AffineMap::get(rank, 0, affineExprs, rewriter.getContext());
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 1d235092b71d55..f36f449da8dbc4 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -253,6 +253,26 @@ func.func @test_add_1d_broadcast_static_to_static(%arg0: tensor<1xf32>, %arg1: t
// -----
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: @test_add_1d_matching_no_broadcast
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+func.func @test_add_1d_matching_no_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
+
+ // CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<1xf32>
+ // CHECK: %[[RESULT:.*]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[ARG1]] : tensor<1xf32>, tensor<1xf32>) outs(%[[VAL_0]] : tensor<1xf32>) {
+ // CHECK: ^bb0(%[[VAL_1:.*]]: f32, %[[VAL_2:.*]]: f32, %[[VAL_3:.*]]: f32):
+ // CHECK: %[[VAL_4:.*]] = arith.addf %[[VAL_1]], %[[VAL_2]] : f32
+ // CHECK: linalg.yield %[[VAL_4]] : f32
+ // CHECK: } -> tensor<1xf32>
+ %0 = tosa.add %arg0, %arg1 : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+
+ // CHECK: return %[[RESULT]] : tensor<1xf32>
+ return %0 : tensor<1xf32>
+}
+
+// -----
+
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_add_1d_matching_static
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
|
|
@llvm/pr-subscribers-mlir-linalg Author: Matthias Gehre (mgehre-amd) ChangesWhen deciding whether to emit a map like This helps later transformations which often require the affine map to be a projected permuatation. Full diff: https://github.com/llvm/llvm-project/pull/123295.diff 2 Files Affected:
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 9295afd36e3ab1..a183c27abf62ae 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -882,8 +882,14 @@ emitElementwiseComputation(ConversionPatternRewriter &rewriter, Location loc,
auto shape = cast<ShapedType>(operand.getType()).getShape();
SmallVector<AffineExpr> affineExprs;
for (auto it : llvm::enumerate(shape)) {
- auto affineExpr = it.value() == 1 ? rewriter.getAffineConstantExpr(0)
- : rewriter.getAffineDimExpr(it.index());
+ // Prefer producting identity maps whenever possible (i.e. no broadcasting
+ // needed) because some transforms (like reshape folding)
+ // do not support affine constant exprs.
+ bool requiresBroadcast =
+ (it.value() == 1 && resultType.getDimSize(it.index()) != 1);
+ auto affineExpr = requiresBroadcast
+ ? rewriter.getAffineConstantExpr(0)
+ : rewriter.getAffineDimExpr(it.index());
affineExprs.push_back(affineExpr);
}
return AffineMap::get(rank, 0, affineExprs, rewriter.getContext());
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 1d235092b71d55..f36f449da8dbc4 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -253,6 +253,26 @@ func.func @test_add_1d_broadcast_static_to_static(%arg0: tensor<1xf32>, %arg1: t
// -----
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: @test_add_1d_matching_no_broadcast
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+func.func @test_add_1d_matching_no_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
+
+ // CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<1xf32>
+ // CHECK: %[[RESULT:.*]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[ARG1]] : tensor<1xf32>, tensor<1xf32>) outs(%[[VAL_0]] : tensor<1xf32>) {
+ // CHECK: ^bb0(%[[VAL_1:.*]]: f32, %[[VAL_2:.*]]: f32, %[[VAL_3:.*]]: f32):
+ // CHECK: %[[VAL_4:.*]] = arith.addf %[[VAL_1]], %[[VAL_2]] : f32
+ // CHECK: linalg.yield %[[VAL_4]] : f32
+ // CHECK: } -> tensor<1xf32>
+ %0 = tosa.add %arg0, %arg1 : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+
+ // CHECK: return %[[RESULT]] : tensor<1xf32>
+ return %0 : tensor<1xf32>
+}
+
+// -----
+
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_add_1d_matching_static
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
|
When deciding whether to emit a map like `#map = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>` or `#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>` for an operand of a `linalg.generic` when lowering element-wise TOSA ops, prefer the latter unless broadcasting of the operand is really needed. This helps later transformations which often require the affine map to be a projected permuatation.
mgehre-amd
force-pushed
the
matthias.tosa_to_linalg_prefer_identity
branch
from
2bc038e to
07bd346
Compare
GeorgeARM
approved these changes
Jan 17, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
When deciding whether to emit a map like
#map = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>or#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>for an operand of alinalg.genericwhen lowering element-wise TOSA ops, prefer the latter unless broadcasting of the operand is really needed.This helps later transformations which often require the affine map to be a projected permuatation.