[Mlir-commits] [mlir] [mlir][affine] ValueBoundsConstraintSet: Fully compose affine.apply (PR #68899)
Matthias Springer
llvmlistbot at llvm.org
Thu Oct 12 18:00:01 PDT 2023
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/68899
>From cce0c62dfc6c3556b059062091260b1ffc45f2c4 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Fri, 13 Oct 2023 02:57:24 +0200
Subject: [PATCH] [mlir][affine] ValueBoundsConstraintSet: Fully compose
affine.apply
Fully compose `affine.apply` ops before adding them to the underlying `FlatLinearConstraints`. This works around a limitation of `FlatLinearConstraints`, which cannot deduce a constant bound if it involves two identical local variables.
Details for future improvements of `FlatLinearConstraints`: The constraint set infrastructure fails to compute a constant bound of -8 for the first variable.
```
Domain: 0, Range: 1, Symbols: 4, Locals: 2
8 constraints
(None None None None None Local Local const)
1 -1 0 0 0 0 0 0 = 0
0 1 -1 1 0 0 0 0 = 0
0 0 1 0 0 0 -16 0 = 0
0 0 0 1 0 -16 0 -8 = 0
0 0 0 0 -1 0 32 31 >= 0
0 0 0 0 1 0 -32 0 >= 0
0 0 0 0 -1 32 0 31 >= 0
0 0 0 0 1 -32 0 0 >= 0
```
---
.../Affine/IR/ValueBoundsOpInterfaceImpl.h | 14 ++++++
.../Affine/IR/ValueBoundsOpInterfaceImpl.cpp | 47 +++++++++++++++++--
.../value-bounds-op-interface-impl.mlir | 32 +++++++++++++
.../Dialect/Affine/TestReifyValueBounds.cpp | 11 ++++-
4 files changed, 97 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h
index 2abbabc5bb2868c..5d4774861bdfd37 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h
@@ -9,11 +9,25 @@
#ifndef MLIR_DIALECT_AFFINE_IR_VALUEBOUNDSOPINTERFACEIMPL_H
#define MLIR_DIALECT_AFFINE_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#include "mlir/Support/LogicalResult.h"
+
namespace mlir {
class DialectRegistry;
+class Value;
namespace affine {
void registerValueBoundsOpInterfaceExternalModels(DialectRegistry ®istry);
+
+/// Compute whether the given values are equal. Return "failure" if equality
+/// could not be determined. `value1`/`value2` must be index-typed.
+///
+/// This function is similar to `ValueBoundsConstraintSet::areEqual`. To work
+/// around limitations in `FlatLinearConstraints`, this function fully composes
+/// `value1` and `value2` (if they are the result of affine.apply ops) before
+/// populating the constraint set. The folding/composing logic can see
+/// opportunities for simplifications that the constraint set implementation
+/// cannot see.
+FailureOr<bool> fullyComposeAndCheckIfEqual(Value value1, Value value2);
} // namespace affine
} // namespace mlir
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index 97dd70e4f1d2b7e..d47c8eb8ccb4272 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -27,12 +27,22 @@ struct AffineApplyOpInterface
assert(applyOp.getAffineMap().getNumResults() == 1 &&
"expected single result");
+ // Fully compose this affine.apply with other ops because the folding logic
+ // can see opportunities for simplifying the affine map that
+ // `FlatLinearConstraints` can currently not see.
+ AffineMap map = applyOp.getAffineMap();
+ SmallVector<Value> operands = llvm::to_vector(applyOp.getOperands());
+ fullyComposeAffineMapAndOperands(&map, &operands);
+
// Align affine map result with dims/symbols in the constraint set.
- AffineExpr expr = applyOp.getAffineMap().getResult(0);
- SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
- applyOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
- SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
- applyOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
+ AffineExpr expr = map.getResult(0);
+ SmallVector<AffineExpr> dimReplacements, symReplacements;
+ for (int64_t i = 0, e = map.getNumDims(); i < e; ++i)
+ dimReplacements.push_back(cstr.getExpr(operands[i]));
+ for (int64_t i = map.getNumDims(),
+ e = map.getNumDims() + map.getNumSymbols();
+ i < e; ++i)
+ symReplacements.push_back(cstr.getExpr(operands[i]));
AffineExpr bound =
expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
cstr.bound(value) == bound;
@@ -92,3 +102,30 @@ void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
});
}
+
+FailureOr<bool> mlir::affine::fullyComposeAndCheckIfEqual(Value value1,
+ Value value2) {
+ assert(value1.getType().isIndex() && "expected index type");
+ assert(value2.getType().isIndex() && "expected index type");
+
+ // Subtract the two values/dimensions from each other. If the result is 0,
+ // both are equal.
+ Builder b(value1.getContext());
+ AffineMap map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0,
+ b.getAffineDimExpr(0) - b.getAffineDimExpr(1));
+ // Fully compose the affine map with other ops because the folding logic
+ // can see opportunities for simplifying the affine map that
+ // `FlatLinearConstraints` can currently not see.
+ SmallVector<Value> mapOperands;
+ mapOperands.push_back(value1);
+ mapOperands.push_back(value2);
+ affine::fullyComposeAffineMapAndOperands(&map, &mapOperands);
+ ValueDimList valueDims;
+ for (Value v : mapOperands)
+ valueDims.push_back({v, std::nullopt});
+ FailureOr<int64_t> bound = ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::EQ, map, valueDims);
+ if (failed(bound))
+ return failure();
+ return *bound == 0;
+}
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 338c48c5b210bc1..8acf358c887a987 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -58,3 +58,35 @@ func.func @affine_min_lb(%a: index) -> (index) {
%2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
return %2 : index
}
+
+// -----
+
+// CHECK-LABEL: func @composed_affine_apply(
+// CHECK: %[[cst:.*]] = arith.constant -8 : index
+// CHECK: return %[[cst]]
+func.func @composed_affine_apply(%i1 : index) -> (index) {
+ // The ValueBoundsOpInterface implementation of affine.apply fully composes
+ // the affine map (and its operands) with other affine.apply ops drawn from
+ // its operands before adding it to the constraint set. This is to work
+ // around a limitation in `FlatLinearConstraints`, which can currently not
+ // compute a constant bound for %s. (The affine map simplification logic can
+ // simplify %s to -8.)
+ %i2 = affine.apply affine_map<(d0) -> ((d0 floordiv 32) * 16)>(%i1)
+ %i3 = affine.apply affine_map<(d0) -> ((d0 floordiv 32) * 16 + 8)>(%i1)
+ %s = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%i2, %i3]
+ %reified = "test.reify_constant_bound"(%s) {type = "EQ"} : (index) -> (index)
+ return %reified : index
+}
+
+
+// -----
+
+// Test for affine::fullyComposeAndCheckIfEqual
+func.func @composed_are_equal(%i1 : index) {
+ %i2 = affine.apply affine_map<(d0) -> ((d0 floordiv 32) * 16)>(%i1)
+ %i3 = affine.apply affine_map<(d0) -> ((d0 floordiv 32) * 16 + 8)>(%i1)
+ %s = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%i2, %i3]
+ // expected-remark @below{{different}}
+ "test.are_equal"(%i2, %i3) {compose} : (index, index) -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index ad017cef1b9bace..6e3c3dff759a2ed 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/Arith/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -186,8 +187,14 @@ static LogicalResult testEquality(func::FuncOp funcOp) {
op->emitOpError("invalid op");
return WalkResult::skip();
}
- FailureOr<bool> equal = ValueBoundsConstraintSet::areEqual(
- op->getOperand(0), op->getOperand(1));
+ FailureOr<bool> equal = failure();
+ if (op->hasAttr("compose")) {
+ equal = affine::fullyComposeAndCheckIfEqual(op->getOperand(0),
+ op->getOperand(1));
+ } else {
+ equal = ValueBoundsConstraintSet::areEqual(op->getOperand(0),
+ op->getOperand(1));
+ }
if (failed(equal)) {
op->emitError("could not determine equality");
} else if (*equal) {
More information about the Mlir-commits
mailing list