[Mlir-commits] [mlir] [mlir][affine] Add ValueBoundsOpInterface to [de]linearize_index (PR #121833)
Krzysztof Drewniak
llvmlistbot at llvm.org
Tue Jan 7 11:06:04 PST 2025
https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/121833
>From b985a173b1cdaf589ef6afa8c0f0177dbf9dd3c3 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 6 Jan 2025 20:46:44 +0000
Subject: [PATCH 1/2] [mlir][affine] Add ValueBoundsOpInterface to
[de]linearize_index
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.
---
.../Affine/IR/ValueBoundsOpInterfaceImpl.cpp | 64 +++++++++++++++++++
.../value-bounds-op-interface-impl.mlir | 38 +++++++++++
2 files changed, 102 insertions(+)
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
+}
>From 426e3355ecf306bd17ae011156003132c8452b8a Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Tue, 7 Jan 2025 19:05:53 +0000
Subject: [PATCH 2/2] Review comments, fix typos in tests
---
.../Affine/IR/ValueBoundsOpInterfaceImpl.cpp | 18 +++----
.../value-bounds-op-interface-impl.mlir | 51 +++++++++++++++++--
2 files changed, 55 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index d65840493897e3..e93b99b4f49866 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -106,19 +106,17 @@ struct AffineDelinearizeIndexOpInterface
SmallVector<OpFoldResult> basis = op.getPaddedBasis();
AffineExpr divisor = cstr.getExpr(1);
- for (OpFoldResult basisElem :
- ArrayRef<OpFoldResult>(basis).drop_front(resIdx + 1))
+ for (OpFoldResult basisElem : llvm::drop_begin(basis, resIdx + 1))
divisor = divisor * cstr.getExpr(basisElem);
- auto resBound = cstr.bound(result);
if (resIdx == 0) {
- resBound == linearIdx.floorDiv(divisor);
+ cstr.bound(value) == linearIdx.floorDiv(divisor);
if (!basis.front().isNull())
- resBound < cstr.getExpr(basis.front());
+ cstr.bound(value) < cstr.getExpr(basis.front());
return;
}
AffineExpr thisBasis = cstr.getExpr(basis[resIdx]);
- resBound == (linearIdx % (thisBasis * divisor)).floorDiv(divisor);
+ cstr.bound(value) == (linearIdx % (thisBasis * divisor)).floorDiv(divisor);
}
};
@@ -135,8 +133,9 @@ struct AffineLinearizeIndexOpInterface
AffineExpr stride = cstr.getExpr(1);
SmallVector<OpFoldResult> basis = op.getPaddedBasis();
OperandRange multiIndex = op.getMultiIndex();
+ unsigned numArgs = multiIndex.size();
for (auto [revArgNum, length] : llvm::enumerate(llvm::reverse(basis))) {
- unsigned argNum = multiIndex.size() - (revArgNum + 1);
+ unsigned argNum = numArgs - (revArgNum + 1);
if (argNum == 0)
break;
OpFoldResult indexAsFoldRes = getAsOpFoldResult(multiIndex[argNum]);
@@ -144,10 +143,9 @@ struct AffineLinearizeIndexOpInterface
stride = stride * cstr.getExpr(length);
}
bound = bound + cstr.getExpr(op.getMultiIndex().front()) * stride;
- auto resBound = cstr.bound(value);
- resBound == bound;
+ cstr.bound(value) == bound;
if (op.getDisjoint() && !basis.front().isNull()) {
- resBound <= stride *cstr.getExpr(basis.front());
+ cstr.bound(value) < stride *cstr.getExpr(basis.front());
}
}
};
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 a6f5f8a5e45f5b..5354eb38d7b039 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -163,8 +163,8 @@ func.func @compare_maps(%a: index, %b: index) {
// 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: %[[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) {
@@ -174,7 +174,31 @@ func.func @delinearize_static(%arg0: index) -> (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.
+ // expected-remark @below{{true}}
+ "test.compare"(%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-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_no_outer_bound
+// 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_no_outer_bound(%arg0: index) -> (index, index, index) {
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %0:3 = affine.delinearize_index %arg0 into (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)
"test.compaare"(%0#0, %c2) {cmp = "LT"} : (index, index) -> ()
// expected-remark @below{{true}}
"test.compare"(%0#1, %c3) {cmp = "LT"} : (index, index) -> ()
@@ -186,10 +210,29 @@ func.func @delinearize_static(%arg0: index) -> (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: %[[v1:.+]] = affine.apply #[[$map]]()[%[[arg1]], %[[arg0]]]
// CHECK: return %[[v1]]
func.func @linearize_static(%arg0: index, %arg1: index) -> index {
+ %c6 = arith.constant 6 : index
%0 = affine.linearize_index disjoint [%arg0, %arg1] by (2, 3) : index
%1 = "test.reify_bound"(%0) {type = "EQ"} : (index) -> (index)
+ // expected-remark @below{{true}}
+ "test.compare"(%0, %c6) {cmp = "LT"} : (index, index) -> ()
+ return %1 : index
+}
+
+// -----
+
+// CHECK: #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>
+// CHECK-LABEL: func.func @linearize_static_no_outer_bound
+// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index)
+// CHECK: %[[v1:.+]] = affine.apply #[[$map]]()[%[[arg1]], %[[arg0]]]
+// CHECK: return %[[v1]]
+func.func @linearize_static_no_outer_bound(%arg0: index, %arg1: index) -> index {
+ %c6 = arith.constant 6 : index
+ %0 = affine.linearize_index disjoint [%arg0, %arg1] by (3) : index
+ %1 = "test.reify_bound"(%0) {type = "EQ"} : (index) -> (index)
+ // expected-error @below{{unknown}}
+ "test.compare"(%0, %c6) {cmp = "LT"} : (index, index) -> ()
return %1 : index
}
More information about the Mlir-commits
mailing list