[mlir][affine] Add ValueBoundsOpInterface to [de]linearize_index #121833
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
Since a need for it came up dowstream (in proving that loops run at least once), this commit implements the ValueBoundsOpInterface for affine.delinearize_index and affine.linearize_index, using affine map representations of the operations they perform. For reasons that are unclear to me, attempting to provide the addition constraints that can be inferred from setting the outer bounds on a affine.delinearize_index doesn't work (that is, in the test, I know %0#0 is both %arg0 / 15 and < 2 (that is, that %arg0 < 30)) but I can't record this extra constraint. I've left this issue as-is for now.
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-affine Author: Krzysztof Drewniak (krzysz00) ChangesSince a need for it came up dowstream (in proving that loops run at least once), this commit implements the ValueBoundsOpInterface for affine.delinearize_index and affine.linearize_index, using affine map representations of the operations they perform. For reasons that are unclear to me, attempting to provide the addition constraints that can be inferred from setting the outer bounds on a affine.delinearize_index doesn't work (that is, in the test, I know %0#0 is both %arg0 / 15 and < 2 (that is, that %arg0 < 30)) but I can't record this extra constraint. I've left this issue as-is for now. Full diff: https://github.com/llvm/llvm-project/pull/121833.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index 82a9fb0d490882..d65840493897e3 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -91,6 +91,66 @@ struct AffineMaxOpInterface
};
};
+struct AffineDelinearizeIndexOpInterface
+ : public ValueBoundsOpInterface::ExternalModel<
+ AffineDelinearizeIndexOpInterface, AffineDelinearizeIndexOp> {
+ void populateBoundsForIndexValue(Operation *rawOp, Value value,
+ ValueBoundsConstraintSet &cstr) const {
+ auto op = cast<AffineDelinearizeIndexOp>(rawOp);
+ auto result = cast<OpResult>(value);
+ assert(result.getOwner() == rawOp &&
+ "bounded value isn't a result of this delinearize_index");
+ unsigned resIdx = result.getResultNumber();
+
+ AffineExpr linearIdx = cstr.getExpr(op.getLinearIndex());
+
+ SmallVector<OpFoldResult> basis = op.getPaddedBasis();
+ AffineExpr divisor = cstr.getExpr(1);
+ for (OpFoldResult basisElem :
+ ArrayRef<OpFoldResult>(basis).drop_front(resIdx + 1))
+ divisor = divisor * cstr.getExpr(basisElem);
+
+ auto resBound = cstr.bound(result);
+ if (resIdx == 0) {
+ resBound == linearIdx.floorDiv(divisor);
+ if (!basis.front().isNull())
+ resBound < cstr.getExpr(basis.front());
+ return;
+ }
+ AffineExpr thisBasis = cstr.getExpr(basis[resIdx]);
+ resBound == (linearIdx % (thisBasis * divisor)).floorDiv(divisor);
+ }
+};
+
+struct AffineLinearizeIndexOpInterface
+ : public ValueBoundsOpInterface::ExternalModel<
+ AffineLinearizeIndexOpInterface, AffineLinearizeIndexOp> {
+ void populateBoundsForIndexValue(Operation *rawOp, Value value,
+ ValueBoundsConstraintSet &cstr) const {
+ auto op = cast<AffineLinearizeIndexOp>(rawOp);
+ assert(value == op.getResult() &&
+ "value isn't the result of this linearize");
+
+ AffineExpr bound = cstr.getExpr(0);
+ AffineExpr stride = cstr.getExpr(1);
+ SmallVector<OpFoldResult> basis = op.getPaddedBasis();
+ OperandRange multiIndex = op.getMultiIndex();
+ for (auto [revArgNum, length] : llvm::enumerate(llvm::reverse(basis))) {
+ unsigned argNum = multiIndex.size() - (revArgNum + 1);
+ if (argNum == 0)
+ break;
+ OpFoldResult indexAsFoldRes = getAsOpFoldResult(multiIndex[argNum]);
+ bound = bound + cstr.getExpr(indexAsFoldRes) * stride;
+ stride = stride * cstr.getExpr(length);
+ }
+ bound = bound + cstr.getExpr(op.getMultiIndex().front()) * stride;
+ auto resBound = cstr.bound(value);
+ resBound == bound;
+ if (op.getDisjoint() && !basis.front().isNull()) {
+ resBound <= stride *cstr.getExpr(basis.front());
+ }
+ }
+};
} // namespace
} // namespace mlir
@@ -100,6 +160,10 @@ void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
+ AffineDelinearizeIndexOp::attachInterface<
+ AffineDelinearizeIndexOpInterface>(*ctx);
+ AffineLinearizeIndexOp::attachInterface<AffineLinearizeIndexOpInterface>(
+ *ctx);
});
}
diff --git a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
index 935c08aceff548..a6f5f8a5e45f5b 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -155,3 +155,41 @@ func.func @compare_maps(%a: index, %b: index) {
: (index, index, index, index) -> ()
return
}
+
+// -----
+
+// CHECK-DAG: #[[$map1:.+]] = affine_map<()[s0] -> (s0 floordiv 15)>
+// CHECK-DAG: #[[$map2:.+]] = affine_map<()[s0] -> ((s0 mod 15) floordiv 5)>
+// CHECK-DAG: #[[$map3:.+]] = affine_map<()[s0] -> (s0 mod 5)>
+// CHECK-LABEL: func.func @delinearize_static
+// CHECK-SAME: (%[[arg0:.+]]: index)
+// CHECK-DAG: %[[v1:.+]] = affine.apply #[[$map1]](}[%[[arg0]]]
+// CHECK-DAG: %[[v2:.+]] = affine.apply #[[$map2]](}[%[[arg0]]]
+// CHECK-DAG: %[[v3:.+]] = affine.apply #[[$map3]]()[%[[arg0]]]
+// CHECK: return %[[v1]], %[[v2]], %[[v3]]
+func.func @delinearize_static(%arg0: index) -> (index, index, index) {
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %0:3 = affine.delinearize_index %arg0 into (2, 3, 5) : index, index, index
+ %1 = "test.reify_bound"(%0#0) {type = "EQ"} : (index) -> (index)
+ %2 = "test.reify_bound"(%0#1) {type = "EQ"} : (index) -> (index)
+ %3 = "test.reify_bound"(%0#2) {type = "EQ"} : (index) -> (index)
+ // TODO: why doesn't this return true? I'm setting the bound.
+ "test.compaare"(%0#0, %c2) {cmp = "LT"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%0#1, %c3) {cmp = "LT"} : (index, index) -> ()
+ return %1, %2, %3 : index, index, index
+}
+
+// -----
+
+// CHECK: #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>
+// CHECK-LABEL: func.func @linearize_static
+// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index)
+// CHECK: %[[v1:.+]] = affine.apply #[[$map]]()[%[[arg0]], %[[arg1]]]
+// CHECK: return %[[v1]]
+func.func @linearize_static(%arg0: index, %arg1: index) -> index {
+ %0 = affine.linearize_index disjoint [%arg0, %arg1] by (2, 3) : index
+ %1 = "test.reify_bound"(%0) {type = "EQ"} : (index) -> (index)
+ return %1 : index
+}
|
kuhar
reviewed
Jan 6, 2025
qedawkins
approved these changes
Jan 7, 2025
| ArrayRef<OpFoldResult>(basis).drop_front(resIdx + 1)) | ||
| divisor = divisor * cstr.getExpr(basisElem); | ||
|
|
||
| auto resBound = cstr.bound(result); |
There was a problem hiding this comment.
nit: Style of other value bounds impls just calls cstr.bound multiple times rather than creating a local instance of the builder.
(I wonder if that could be the cause for why the extra < bound in delinearize_index is failing below?)
| %0 = affine.linearize_index disjoint [%arg0, %arg1] by (2, 3) : index | ||
| %1 = "test.reify_bound"(%0) {type = "EQ"} : (index) -> (index) | ||
| return %1 : index | ||
| } |
There was a problem hiding this comment.
Can you add tests for both without the leading basis elements? (getPaddedBasis = [null, ...])
| %1 = "test.reify_bound"(%0#0) {type = "EQ"} : (index) -> (index) | ||
| %2 = "test.reify_bound"(%0#1) {type = "EQ"} : (index) -> (index) | ||
| %3 = "test.reify_bound"(%0#2) {type = "EQ"} : (index) -> (index) | ||
| // TODO: why doesn't this return true? I'm setting the bound. |
There was a problem hiding this comment.
Does it work if you comment out the other bounds of resBound, e.g., resBound == linearIdx.floorDiv(divisor);?
There was a problem hiding this comment.
Does it return "false" or "unknown"?
There was a problem hiding this comment.
I just ran your code and this returns true.
There was a problem hiding this comment.
... turns out there was a typo in the tests
krzysz00
added a commit
to iree-org/llvm-project
that referenced
this pull request
Jan 10, 2025
…m#121833) Since a need for it came up dowstream (in proving that loops run at least once), this commit implements the ValueBoundsOpInterface for affine.delinearize_index and affine.linearize_index, using affine map representations of the operations they perform. These implementations also use information from outer bounds to impose additional constraints when those are available.
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.
Since a need for it came up dowstream (in proving that loops run at least once), this commit implements the ValueBoundsOpInterface for affine.delinearize_index and affine.linearize_index, using affine map representations of the operations they perform.
These implementations also use information from outer bounds to impose additional constraints when those are available.