Skip to content

[MLIR][Affine] Reject fully consumed bounded delinearize split #217472

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joker-eph wants to merge 1 commit into
base: main
Choose a base branch
from
+28 −2
Open

[MLIR][Affine] Reject fully consumed bounded delinearize split #217472

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions mlir/lib/Dialect/Affine/IR/AffineOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5233,6 +5233,9 @@ struct CancelDelinearizeOfLinearizeDisjointExactTail
/// last k > 1 components of the delinearization basis multiply to the
/// last component of the linearization basis, break the linearization and
/// delinearization into two parts, peeling off the last input to linearization.
/// The split does not apply when it would consume an entire outer-bounded
/// delinearization basis because earlier linearization inputs still contribute
/// to the first delinearized result.
///
/// For example:
/// %0 = affine.linearize_index [%z, %y, %x] by (3, 2, 32) : index
Expand Down Expand Up @@ -5297,6 +5300,10 @@ struct SplitDelinearizeSpanningLastLinearizeArg final
delinearizeOp,
"need at least two elements to form the basis product");

if (elemsToSplit == basis.size() && delinearizeOp.hasOuterBound())
return rewriter.notifyMatchFailure(
delinearizeOp, "split would consume entire bounded basis");

Value linearizeWithoutBack = affine::AffineLinearizeIndexOp::create(
rewriter, linearizeOp.getLoc(), linearizeOp.getLinearIndex().getType(),
linearizeOp.getMultiIndex().drop_back(), linearizeOp.getDynamicBasis(),
Expand Down
23 changes: 21 additions & 2 deletions mlir/test/Dialect/Affine/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence top-down=0" | FileCheck %s --check-prefix=CHECK-BOTTOM-UP

// XFAIL: mlir-expensive-checks

// -----

// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 - 1)>
Expand Down Expand Up @@ -1881,6 +1879,27 @@ func.func @split_delinearize_empty_linearize_basis(%arg0: index) -> (index, inde

// -----

// A split that consumes an entire bounded delinearization basis would lose the
// contribution of earlier linearization inputs to the first result.
// CHECK-LABEL: func @dont_split_fully_consumed_bounded_basis
// CHECK-SAME: (%[[A:.+]]: index, %[[B:.+]]: index)
// CHECK: %[[LIN:.+]] = affine.linearize_index disjoint [%[[A]], %[[B]]] by (2, 4) : index
// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[LIN]] into (2, 2) : index, index
// CHECK: return %[[DELIN]]#0, %[[DELIN]]#1
// CHECK-BOTTOM-UP-LABEL: func @dont_split_fully_consumed_bounded_basis
// CHECK-BOTTOM-UP-SAME: (%[[A:.+]]: index, %[[B:.+]]: index)
// CHECK-BOTTOM-UP: %[[LIN:.+]] = affine.linearize_index disjoint [%[[A]], %[[B]]] by (2, 4) : index
// CHECK-BOTTOM-UP: %[[DELIN:.+]]:2 = affine.delinearize_index %[[LIN]] into (2, 2) : index, index
// CHECK-BOTTOM-UP: return %[[DELIN]]#0, %[[DELIN]]#1
func.func @dont_split_fully_consumed_bounded_basis(%a: index, %b: index)
-> (index, index) {
%0 = affine.linearize_index disjoint [%a, %b] by (2, 4) : index
%1:2 = affine.delinearize_index %0 into (2, 2) : index, index
return %1#0, %1#1 : index, index
}

// -----

// CHECK-LABEL: @linearize_unit_basis_disjoint
// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index)
// CHECK: %[[ret:.+]] = affine.linearize_index disjoint [%[[arg0]], %[[arg2]]] by (3, %[[arg3]]) : index
Expand Down
Loading