[Mlir-commits] [mlir] [MLIR][Affine] Fix/complete access index invariance, add isInvariantAccess (PR #84602)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Mar 8 20:39:02 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
Author: Uday Bondhugula (bondhugula)
<details>
<summary>Changes</summary>
isAccessIndexInvariant had outdated code and didn't handle IR with multiple
affine.apply ops, which is inconvenient when used as a utility. This is
addressed by switching to use the proper API on AffineValueMap. Add
mlir::affine::isInvariantAccess exposed for outside use and tested via
the test pass. Add a method on AffineValueMap. Add test cases to
exercise simplification and composition for invariant access analysis.
A TODO/FIXME has been added but this issue existed before.
---
Full diff: https://github.com/llvm/llvm-project/pull/84602.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h (+6)
- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h (+5)
- (modified) mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp (+22-20)
- (modified) mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp (+9)
- (modified) mlir/test/Dialect/Affine/access-analysis.mlir (+31-2)
- (modified) mlir/test/lib/Dialect/Affine/TestAccessAnalysis.cpp (+13-6)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
index 1f64b57cac5782..114723128d1754 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
@@ -48,6 +48,12 @@ std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);
/// this method is thus able to determine non-trivial divisors.
uint64_t getLargestDivisorOfTripCount(AffineForOp forOp);
+/// Checks if an affine load or store access depends on `forOp'. This also
+/// transitively checks if the load/store is dependent on another loop IV which
+/// in turn uses `forOp` is its loop bounds.
+template <typename LoadOrStoreOp>
+bool isInvariantAccess(LoadOrStoreOp memOp, AffineForOp forOp);
+
/// Given an induction variable `iv` of type AffineForOp and `indices` of type
/// IndexType, returns the set of `indices` that are independent of `iv`.
///
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h b/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h
index 8439930a87467c..7ad0e4a1e5ea04 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h
@@ -44,6 +44,11 @@ class AffineValueMap {
// Resets this AffineValueMap with 'map', 'operands', and 'results'.
void reset(AffineMap map, ValueRange operands, ValueRange results = {});
+ /// Composes all incoming affine.apply ops and then simplifies and
+ /// canonicalizes the map and operands. This can change the number of
+ /// operands, but the result count remains the same.
+ void composeSimplifyAndCanonicalize();
+
/// Return the value map that is the difference of value maps 'a' and 'b',
/// represented as an affine map and its operands. The output map + operands
/// are canonicalized and simplified.
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index fc0515ba95f4fe..9c6cf5a2f943c5 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -161,29 +161,31 @@ uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {
/// Returns false in cases with more than one AffineApplyOp, this is
/// conservative.
static bool isAccessIndexInvariant(Value iv, Value index) {
- assert(isAffineForInductionVar(iv) && "iv must be a AffineForOp");
- assert(isa<IndexType>(index.getType()) && "index must be of IndexType");
- SmallVector<Operation *, 4> affineApplyOps;
- getReachableAffineApplyOps({index}, affineApplyOps);
-
- if (affineApplyOps.empty()) {
- // Pointer equality test because of Value pointer semantics.
- return index != iv;
- }
-
- if (affineApplyOps.size() > 1) {
- affineApplyOps[0]->emitRemark(
- "CompositionAffineMapsPass must have been run: there should be at most "
- "one AffineApplyOp, returning false conservatively.");
- return false;
- }
+ assert(isAffineForInductionVar(iv) && "iv must be an affine.for iv");
+ assert(isa<IndexType>(index.getType()) && "index must be of 'index' type");
+ auto map = AffineMap::getMultiDimIdentityMap(/*numDims=*/1, iv.getContext());
+ SmallVector<Value> operands = {index};
+ AffineValueMap avm(map, operands);
+ avm.composeSimplifyAndCanonicalize();
+ return !avm.isFunctionOf(0, iv);
+}
- auto composeOp = cast<AffineApplyOp>(affineApplyOps[0]);
- // We need yet another level of indirection because the `dim` index of the
- // access may not correspond to the `dim` index of composeOp.
- return !composeOp.getAffineValueMap().isFunctionOf(0, iv);
+// Pre-requisite: Loop bounds should be in canonical form.
+template <typename LoadOrStoreOp>
+bool mlir::affine::isInvariantAccess(LoadOrStoreOp memOp, AffineForOp forOp) {
+ AffineValueMap avm(memOp.getAffineMap(), memOp.getMapOperands());
+ avm.composeSimplifyAndCanonicalize();
+ return !llvm::is_contained(avm.getOperands(), forOp.getInductionVar());
}
+// Explicitly instantiate the template so that the compiler knows we need them.
+template bool mlir::affine::isInvariantAccess(AffineReadOpInterface,
+ AffineForOp);
+template bool mlir::affine::isInvariantAccess(AffineWriteOpInterface,
+ AffineForOp);
+template bool mlir::affine::isInvariantAccess(AffineLoadOp, AffineForOp);
+template bool mlir::affine::isInvariantAccess(AffineStoreOp, AffineForOp);
+
DenseSet<Value> mlir::affine::getInvariantAccesses(Value iv,
ArrayRef<Value> indices) {
DenseSet<Value> res;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp b/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
index 2800237fd05ac6..6a52849186872e 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
@@ -24,6 +24,15 @@ void AffineValueMap::reset(AffineMap map, ValueRange operands,
this->results.assign(results.begin(), results.end());
}
+void AffineValueMap::composeSimplifyAndCanonicalize() {
+ AffineMap sMap = getAffineMap();
+ fullyComposeAffineMapAndOperands(&sMap, &operands);
+ // Full composition also canonicalizes and simplifies before returning. We
+ // need to canonicalize once more to drop unused operands.
+ canonicalizeMapAndOperands(&sMap, &operands);
+ this->map.reset(sMap);
+}
+
void AffineValueMap::difference(const AffineValueMap &a,
const AffineValueMap &b, AffineValueMap *res) {
assert(a.getNumResults() == b.getNumResults() && "invalid inputs");
diff --git a/mlir/test/Dialect/Affine/access-analysis.mlir b/mlir/test/Dialect/Affine/access-analysis.mlir
index 68310b9323535a..789de646a8f9e2 100644
--- a/mlir/test/Dialect/Affine/access-analysis.mlir
+++ b/mlir/test/Dialect/Affine/access-analysis.mlir
@@ -1,13 +1,14 @@
// RUN: mlir-opt %s -split-input-file -test-affine-access-analysis -verify-diagnostics | FileCheck %s
-// CHECK-LABEL: func @loop_1d
-func.func @loop_1d(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
+// CHECK-LABEL: func @loop_simple
+func.func @loop_simple(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
%c0 = arith.constant 0 : index
%M = memref.dim %A, %c0 : memref<?x?xf32>
affine.for %i = 0 to %M {
affine.for %j = 0 to %M {
affine.load %A[%c0, %i] : memref<?x?xf32>
// expected-remark at above {{contiguous along loop 0}}
+ // expected-remark at above {{invariant along loop 1}}
affine.load %A[%c0, 8 * %i + %j] : memref<?x?xf32>
// expected-remark at above {{contiguous along loop 1}}
// Note/FIXME: access stride isn't being checked.
@@ -15,6 +16,7 @@ func.func @loop_1d(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
// These are all non-contiguous along both loops. Nothing is emitted.
affine.load %A[%i, %c0] : memref<?x?xf32>
+ // expected-remark at above {{invariant along loop 1}}
// Note/FIXME: access stride isn't being checked.
affine.load %A[%i, 8 * %j] : memref<?x?xf32>
// expected-remark at above {{contiguous along loop 1}}
@@ -27,6 +29,22 @@ func.func @loop_1d(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
// -----
+// CHECK-LABEL: func @loop_unsimplified
+func.func @loop_unsimplified(%A : memref<100xf32>) {
+ affine.for %i = 0 to 100 {
+ affine.load %A[2 * %i - %i - %i] : memref<100xf32>
+ // expected-remark at above {{invariant along loop 0}}
+
+ %m = affine.apply affine_map<(d0) -> (-2 * d0)>(%i)
+ %n = affine.apply affine_map<(d0) -> (2 * d0)>(%i)
+ affine.load %A[(%m + %n) floordiv 2] : memref<100xf32>
+ // expected-remark at above {{invariant along loop 0}}
+ }
+ return
+}
+
+// -----
+
#map = affine_map<(d0) -> (d0 * 16)>
#map1 = affine_map<(d0) -> (d0 * 16 + 16)>
#map2 = affine_map<(d0) -> (d0)>
@@ -41,11 +59,19 @@ func.func @tiled(%arg0: memref<*xf32>) {
%alloc_0 = memref.alloc() : memref<1x16x1x16xf32>
affine.for %arg4 = #map(%arg1) to #map1(%arg1) {
affine.for %arg5 = #map(%arg3) to #map1(%arg3) {
+ // TODO: here and below, the access isn't really invariant
+ // along tile-space IVs where the intra-tile IVs' bounds
+ // depend on them.
%0 = affine.load %cast[%arg4] : memref<64xf32>
// expected-remark at above {{contiguous along loop 3}}
+ // expected-remark at above {{invariant along loop 0}}
+ // expected-remark at above {{invariant along loop 1}}
+ // expected-remark at above {{invariant along loop 2}}
+ // expected-remark at above {{invariant along loop 4}}
affine.store %0, %alloc_0[0, %arg1 * -16 + %arg4, 0, %arg3 * -16 + %arg5] : memref<1x16x1x16xf32>
// expected-remark at above {{contiguous along loop 4}}
// expected-remark at above {{contiguous along loop 2}}
+ // expected-remark at above {{invariant along loop 1}}
}
}
affine.for %arg4 = #map(%arg1) to #map1(%arg1) {
@@ -56,6 +82,9 @@ func.func @tiled(%arg0: memref<*xf32>) {
// expected-remark at above {{contiguous along loop 2}}
affine.store %0, %alloc[0, %arg5, %arg6, %arg4] : memref<1x224x224x64xf32>
// expected-remark at above {{contiguous along loop 3}}
+ // expected-remark at above {{invariant along loop 0}}
+ // expected-remark at above {{invariant along loop 1}}
+ // expected-remark at above {{invariant along loop 2}}
}
}
}
diff --git a/mlir/test/lib/Dialect/Affine/TestAccessAnalysis.cpp b/mlir/test/lib/Dialect/Affine/TestAccessAnalysis.cpp
index b38046299d504a..751302550092d7 100644
--- a/mlir/test/lib/Dialect/Affine/TestAccessAnalysis.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestAccessAnalysis.cpp
@@ -59,18 +59,25 @@ void TestAccessAnalysis::runOnOperation() {
enclosingOps.clear();
getAffineForIVs(*memOp, &enclosingOps);
for (unsigned d = 0, e = enclosingOps.size(); d < e; d++) {
+ AffineForOp loop = enclosingOps[d];
int memRefDim;
- bool isContiguous;
+ bool isContiguous, isInvariant;
if (auto read = dyn_cast<AffineReadOpInterface>(memOp)) {
- isContiguous = isContiguousAccess(enclosingOps[d].getInductionVar(),
- read, &memRefDim);
+ isContiguous =
+ isContiguousAccess(loop.getInductionVar(), read, &memRefDim);
+ isInvariant = isInvariantAccess(read, loop);
} else {
- isContiguous = isContiguousAccess(enclosingOps[d].getInductionVar(),
- cast<AffineWriteOpInterface>(memOp),
- &memRefDim);
+ auto write = cast<AffineWriteOpInterface>(memOp);
+ isContiguous =
+ isContiguousAccess(loop.getInductionVar(), write, &memRefDim);
+ isInvariant = isInvariantAccess(write, loop);
}
+ // Check for contiguity for the innermost memref dimension to avoid
+ // emitting too many diagnostics.
if (isContiguous && memRefDim == 0)
memOp->emitRemark("contiguous along loop ") << d << '\n';
+ if (isInvariant)
+ memOp->emitRemark("invariant along loop ") << d << '\n';
}
}
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/84602
More information about the Mlir-commits
mailing list