[mlir][affine] Make [de]linearize_index a valid source of dims #138929
+53
−1
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
There's a sense in which affine.linearize_index and affine.delinearize_index are special-cases of affine.apply (which get their own ops to enable better code generation and more accurate canonicalization). Therefore, allow these operations to be dimension operands for operations like affine.load just like affine.apply can be.
|
@llvm/pr-subscribers-mlir-affine Author: Krzysztof Drewniak (krzysz00) ChangesThere's a sense in which affine.linearize_index and affine.delinearize_index are special-cases of affine.apply (which get their own ops to enable better code generation and more accurate canonicalization). Therefore, allow these operations to be dimension operands for operations like affine.load just like affine.apply can be. Full diff: https://github.com/llvm/llvm-project/pull/138929.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 1dfe2a57df587..65711c88dd582 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -1153,6 +1153,10 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> {
/// there is no outer bound specified, the leading entry of this result will be
/// nullptr.
SmallVector<OpFoldResult> getPaddedBasis();
+
+ /// Returns true if the result of this operation can be used as dimension id
+ /// within 'region', i.e., for all its uses with `region`.
+ bool isValidDim(Region *region);
}];
let hasVerifier = 1;
@@ -1253,6 +1257,10 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
/// If there is no outer bound specified, the leading entry of this basis will be
/// nullptr.
SmallVector<OpFoldResult> getPaddedBasis();
+
+ /// Returns true if the result of this operation can be used as dimension id
+ /// within 'region', i.e., for all its uses with `region`.
+ bool isValidDim(Region *region);
}];
let hasVerifier = 1;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 65f85444e70db..a01cbb1a85eb3 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -307,6 +307,8 @@ bool mlir::affine::isValidDim(Value value) {
// *) It is valid as a symbol.
// *) It is an induction variable.
// *) It is the result of an affine apply operation with dimension id operands.
+// *) It is the result of a more specialized index transformation (ex.
+// delinearize_index or linearize_index) with dimension id operands.
bool mlir::affine::isValidDim(Value value, Region *region) {
// The value must be an index type.
if (!value.getType().isIndex())
@@ -326,6 +328,10 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
// Affine apply operation is ok if all of its operands are ok.
if (auto applyOp = dyn_cast<AffineApplyOp>(op))
return applyOp.isValidDim(region);
+ if (auto delinearizeOp = dyn_cast<AffineDelinearizeIndexOp>(op))
+ return delinearizeOp.isValidDim(region);
+ if (auto linearizeOp = dyn_cast<AffineLinearizeIndexOp>(op))
+ return linearizeOp.isValidDim(region);
// The dim op is okay if its operand memref/tensor is defined at the top
// level.
if (auto dimOp = dyn_cast<ShapedDimOpInterface>(op))
@@ -4636,6 +4642,14 @@ void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
hasOuterBound);
}
+// The result of the affine apply operation can be used as a dimension id if all
+// its operands are valid dimension ids with the parent operation of `region`
+// defining the polyhedral scope for symbols.
+bool AffineDelinearizeIndexOp::isValidDim(Region *region) {
+ return llvm::all_of(getOperands(),
+ [&](Value op) { return ::isValidDim(op, region); });
+}
+
void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
OperationState &odsState,
Value linearIndex, ArrayRef<int64_t> basis,
@@ -5023,6 +5037,14 @@ LogicalResult AffineLinearizeIndexOp::verify() {
return success();
}
+// The result of the affine apply operation can be used as a dimension id if all
+// its operands are valid dimension ids with the parent operation of `region`
+// defining the polyhedral scope for symbols.
+bool AffineLinearizeIndexOp::isValidDim(Region *region) {
+ return llvm::all_of(getOperands(),
+ [&](Value op) { return ::isValidDim(op, region); });
+}
+
OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
std::optional<SmallVector<int64_t>> maybeStaticBasis =
foldCstValueToCstAttrBasis(getMixedBasis(), getDynamicBasisMutable(),
diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index 9703c05fff8f6..c4bb22f9a8dae 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -544,6 +544,34 @@ func.func @dynamic_dimension_index() {
// -----
+func.func @dynamic_linearized_index() {
+ "unknown.region"() ({
+ %idx = "unknown.test"() : () -> (index)
+ %memref = "unknown.test"() : () -> memref<?xf32>
+ %pos = affine.linearize_index [%idx, %idx] by (8) : index
+ // expected-error@below {{op operand cannot be used as a dimension id}}
+ affine.load %memref[%pos] : memref<?xf32>
+ "unknown.terminator"() : () -> ()
+ }) : () -> ()
+ return
+}
+
+// -----
+
+func.func @dynamic_delinearized_index() {
+ "unknown.region"() ({
+ %idx = "unknown.test"() : () -> (index)
+ %memref = "unknown.test"() : () -> memref<?x?xf32>
+ %pos0, %pos1 = affine.delinearize_index %idx into (8) : index, index
+ // expected-error@below {{op operand cannot be used as a dimension id}}
+ affine.load %memref[%pos0, %pos1] : memref<?x?xf32>
+ "unknown.terminator"() : () -> ()
+ }) : () -> ()
+ return
+}
+
+// -----
+
#map = affine_map<() -> ()>
#map1 = affine_map<() -> (1)>
func.func @no_lower_bound() {
diff --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir
index 233c18c82831c..8a3f41d1d9b05 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -148,6 +148,23 @@ func.func @valid_symbol_affine_scope(%n : index, %A : memref<?xf32>) {
// -----
+// Test dimension constraints for linearize_index and delinearize_index
+
+// CHECK-LABEL: func @valid_dim_linearize_delinearize
+func.func @valid_dim_linearize_delinearize(%m : index, %n : index, %A : memref<?xf32>, %B: memref<?x32x?xf32>) {
+ affine.for %0 = 0 to %m {
+ affine.for %1 = 0 to %n {
+ %load_idx = affine.linearize_index disjoint [%0, %1] by (%m, %n) : index
+ %store_idx0, %store_idx1 = affine.delinearize_index %n into (32) : index, index
+ %v = affine.load %A[%load_idx] : memref<?xf32>
+ affine.store %v, %B[%0, %store_idx1, %store_idx0] : memref<?x32x?xf32>
+ }
+ }
+ return
+}
+
+// -----
+
// Test the fact that module op always provides an affine scope.
%idx = "test.foo"() : () -> (index)
@@ -309,7 +326,7 @@ func.func @linearize_mixed(%index0: index, %index1: index, %index2: index, %basi
module {
func.func @gpu_launch_affine() {
%c1 = arith.constant 1 : index
- gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
+ gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
%thread_id_x = gpu.thread_id x
%c128 = arith.constant 128 : index
|
|
@llvm/pr-subscribers-mlir Author: Krzysztof Drewniak (krzysz00) ChangesThere's a sense in which affine.linearize_index and affine.delinearize_index are special-cases of affine.apply (which get their own ops to enable better code generation and more accurate canonicalization). Therefore, allow these operations to be dimension operands for operations like affine.load just like affine.apply can be. Full diff: https://github.com/llvm/llvm-project/pull/138929.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 1dfe2a57df587..65711c88dd582 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -1153,6 +1153,10 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> {
/// there is no outer bound specified, the leading entry of this result will be
/// nullptr.
SmallVector<OpFoldResult> getPaddedBasis();
+
+ /// Returns true if the result of this operation can be used as dimension id
+ /// within 'region', i.e., for all its uses with `region`.
+ bool isValidDim(Region *region);
}];
let hasVerifier = 1;
@@ -1253,6 +1257,10 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
/// If there is no outer bound specified, the leading entry of this basis will be
/// nullptr.
SmallVector<OpFoldResult> getPaddedBasis();
+
+ /// Returns true if the result of this operation can be used as dimension id
+ /// within 'region', i.e., for all its uses with `region`.
+ bool isValidDim(Region *region);
}];
let hasVerifier = 1;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 65f85444e70db..a01cbb1a85eb3 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -307,6 +307,8 @@ bool mlir::affine::isValidDim(Value value) {
// *) It is valid as a symbol.
// *) It is an induction variable.
// *) It is the result of an affine apply operation with dimension id operands.
+// *) It is the result of a more specialized index transformation (ex.
+// delinearize_index or linearize_index) with dimension id operands.
bool mlir::affine::isValidDim(Value value, Region *region) {
// The value must be an index type.
if (!value.getType().isIndex())
@@ -326,6 +328,10 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
// Affine apply operation is ok if all of its operands are ok.
if (auto applyOp = dyn_cast<AffineApplyOp>(op))
return applyOp.isValidDim(region);
+ if (auto delinearizeOp = dyn_cast<AffineDelinearizeIndexOp>(op))
+ return delinearizeOp.isValidDim(region);
+ if (auto linearizeOp = dyn_cast<AffineLinearizeIndexOp>(op))
+ return linearizeOp.isValidDim(region);
// The dim op is okay if its operand memref/tensor is defined at the top
// level.
if (auto dimOp = dyn_cast<ShapedDimOpInterface>(op))
@@ -4636,6 +4642,14 @@ void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
hasOuterBound);
}
+// The result of the affine apply operation can be used as a dimension id if all
+// its operands are valid dimension ids with the parent operation of `region`
+// defining the polyhedral scope for symbols.
+bool AffineDelinearizeIndexOp::isValidDim(Region *region) {
+ return llvm::all_of(getOperands(),
+ [&](Value op) { return ::isValidDim(op, region); });
+}
+
void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
OperationState &odsState,
Value linearIndex, ArrayRef<int64_t> basis,
@@ -5023,6 +5037,14 @@ LogicalResult AffineLinearizeIndexOp::verify() {
return success();
}
+// The result of the affine apply operation can be used as a dimension id if all
+// its operands are valid dimension ids with the parent operation of `region`
+// defining the polyhedral scope for symbols.
+bool AffineLinearizeIndexOp::isValidDim(Region *region) {
+ return llvm::all_of(getOperands(),
+ [&](Value op) { return ::isValidDim(op, region); });
+}
+
OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
std::optional<SmallVector<int64_t>> maybeStaticBasis =
foldCstValueToCstAttrBasis(getMixedBasis(), getDynamicBasisMutable(),
diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index 9703c05fff8f6..c4bb22f9a8dae 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -544,6 +544,34 @@ func.func @dynamic_dimension_index() {
// -----
+func.func @dynamic_linearized_index() {
+ "unknown.region"() ({
+ %idx = "unknown.test"() : () -> (index)
+ %memref = "unknown.test"() : () -> memref<?xf32>
+ %pos = affine.linearize_index [%idx, %idx] by (8) : index
+ // expected-error@below {{op operand cannot be used as a dimension id}}
+ affine.load %memref[%pos] : memref<?xf32>
+ "unknown.terminator"() : () -> ()
+ }) : () -> ()
+ return
+}
+
+// -----
+
+func.func @dynamic_delinearized_index() {
+ "unknown.region"() ({
+ %idx = "unknown.test"() : () -> (index)
+ %memref = "unknown.test"() : () -> memref<?x?xf32>
+ %pos0, %pos1 = affine.delinearize_index %idx into (8) : index, index
+ // expected-error@below {{op operand cannot be used as a dimension id}}
+ affine.load %memref[%pos0, %pos1] : memref<?x?xf32>
+ "unknown.terminator"() : () -> ()
+ }) : () -> ()
+ return
+}
+
+// -----
+
#map = affine_map<() -> ()>
#map1 = affine_map<() -> (1)>
func.func @no_lower_bound() {
diff --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir
index 233c18c82831c..8a3f41d1d9b05 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -148,6 +148,23 @@ func.func @valid_symbol_affine_scope(%n : index, %A : memref<?xf32>) {
// -----
+// Test dimension constraints for linearize_index and delinearize_index
+
+// CHECK-LABEL: func @valid_dim_linearize_delinearize
+func.func @valid_dim_linearize_delinearize(%m : index, %n : index, %A : memref<?xf32>, %B: memref<?x32x?xf32>) {
+ affine.for %0 = 0 to %m {
+ affine.for %1 = 0 to %n {
+ %load_idx = affine.linearize_index disjoint [%0, %1] by (%m, %n) : index
+ %store_idx0, %store_idx1 = affine.delinearize_index %n into (32) : index, index
+ %v = affine.load %A[%load_idx] : memref<?xf32>
+ affine.store %v, %B[%0, %store_idx1, %store_idx0] : memref<?x32x?xf32>
+ }
+ }
+ return
+}
+
+// -----
+
// Test the fact that module op always provides an affine scope.
%idx = "test.foo"() : () -> (index)
@@ -309,7 +326,7 @@ func.func @linearize_mixed(%index0: index, %index1: index, %index2: index, %basi
module {
func.func @gpu_launch_affine() {
%c1 = arith.constant 1 : index
- gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
+ gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
%thread_id_x = gpu.thread_id x
%c128 = arith.constant 128 : index
|
|
PRs that depend on this: |
|
Ping |
MaheshRavishankar
requested changes
May 10, 2025
MaheshRavishankar
approved these changes
May 13, 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.
There's a sense in which affine.linearize_index and affine.delinearize_index are special-cases of affine.apply (which get their own ops to enable better code generation and more accurate canonicalization). Therefore, allow these operations to be dimension operands for operations like affine.load just like affine.apply can be.