[Mlir-commits] [mlir] [MLIR][VectorToGPU] Cache slice analysis in getOpToConvert (PR #221795)
Dmitriy Smirnov
llvmlistbot at llvm.org
Thu Sep 10 01:39:20 PDT 2026
https://github.com/d-smirnov updated https://github.com/llvm/llvm-project/pull/221795
>From c470b55cf536eb377ea3b23ab93778a3259c7815 Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Mon, 7 Sep 2026 12:38:58 +0100
Subject: [PATCH 1/3] [MLIR][VectorToGPU] Cache slice analysis in
getOpToConvert
ConvertVectorToGPU spent a lot of time rebuilding the same backward and
forward slices for overlapping vector/MMA use-def components, which was
especially visible on int8 cases.
This change caches backward slices, forward slices, and
supportsMMaMatrixType() results per operation, and marks fully expanded
components as already analyzed. That avoids repeatedly rediscovering and
revalidating the same component from multiple seeds without changing the
lowering itself.
The resulting conversion is unchanged, but compile time for the affected
int8 case drops significantly.
Change-Id: I742c2c85442ce6d2f8ff82322b85d0dc7693c681
---
.../Conversion/VectorToGPU/VectorToGPU.cpp | 105 +++++++++++-------
1 file changed, 63 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 975fe28399609..ffbb113758c48 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -325,30 +325,51 @@ static bool supportsMMaMatrixType(Operation *op, bool useNvGpu) {
return elementwiseSupportsMMAMatrixType(op);
}
-/// Return an unsorted slice handling scf.for region differently than
-/// `getSlice`. In scf.for we only want to include as part of the slice elements
-/// that are part of the use/def chain.
-static SetVector<Operation *>
-getSliceContract(Operation *op,
- const BackwardSliceOptions &backwardSliceOptions,
- const ForwardSliceOptions &forwardSliceOptions) {
- SetVector<Operation *> slice;
- slice.insert(op);
- unsigned currentIndex = 0;
- SetVector<Operation *> backwardSlice;
- SetVector<Operation *> forwardSlice;
- while (currentIndex != slice.size()) {
- auto *currentOp = (slice)[currentIndex];
- // Compute and insert the backwardSlice starting from currentOp.
- backwardSlice.clear();
+// Analyze slice of operations based on convert op to figure out if the whole
+// slice can be converted to MMA operations.
+static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
+ bool useNvGpu) {
+ auto hasVectorDest = [](Operation *op) {
+ return llvm::any_of(op->getResultTypes(), llvm::IsaPred<VectorType>);
+ };
+ BackwardSliceOptions backwardSliceOptions;
+ backwardSliceOptions.filter = hasVectorDest;
+
+ auto hasVectorSrc = [](Operation *op) {
+ return llvm::any_of(op->getOperandTypes(), llvm::IsaPred<VectorType>);
+ };
+ ForwardSliceOptions forwardSliceOptions;
+ forwardSliceOptions.filter = hasVectorSrc;
+
+ DenseMap<Operation *, SmallVector<Operation *>> backwardSliceCache;
+ DenseMap<Operation *, SmallVector<Operation *>> forwardSliceCache;
+ DenseMap<Operation *, bool> supportsMMAMatrixTypeCache;
+ // Once we have expanded a seed to its full use/def component, every op in
+ // that component would rediscover the same set again.
+ DenseSet<Operation *> analyzedOps;
+
+ auto getCachedBackwardSlice =
+ [&](Operation *currentOp) -> ArrayRef<Operation *> {
+ auto [it, inserted] = backwardSliceCache.try_emplace(currentOp);
+ if (!inserted)
+ return it->second;
+
+ SetVector<Operation *> backwardSlice;
LogicalResult result =
getBackwardSlice(currentOp, &backwardSlice, backwardSliceOptions);
assert(result.succeeded() && "expected a backward slice");
(void)result;
- slice.insert_range(backwardSlice);
+ it->second.append(backwardSlice.begin(), backwardSlice.end());
+ return it->second;
+ };
+
+ auto getCachedForwardSlice =
+ [&](Operation *currentOp) -> ArrayRef<Operation *> {
+ auto [it, inserted] = forwardSliceCache.try_emplace(currentOp);
+ if (!inserted)
+ return it->second;
- // Compute and insert the forwardSlice starting from currentOp.
- forwardSlice.clear();
+ SetVector<Operation *> forwardSlice;
// Special case for ForOp, we don't want to include the whole region but
// only the value using the region arguments.
// TODO: We should refine this to only care about the region arguments being
@@ -361,42 +382,42 @@ getSliceContract(Operation *op,
} else {
getForwardSlice(currentOp, &forwardSlice, forwardSliceOptions);
}
- slice.insert_range(forwardSlice);
- ++currentIndex;
- }
- return slice;
-}
-
-// Analyze slice of operations based on convert op to figure out if the whole
-// slice can be converted to MMA operations.
-static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
- bool useNvGpu) {
- auto hasVectorDest = [](Operation *op) {
- return llvm::any_of(op->getResultTypes(), llvm::IsaPred<VectorType>);
+ it->second.append(forwardSlice.begin(), forwardSlice.end());
+ return it->second;
};
- BackwardSliceOptions backwardSliceOptions;
- backwardSliceOptions.filter = hasVectorDest;
- auto hasVectorSrc = [](Operation *op) {
- return llvm::any_of(op->getOperandTypes(), llvm::IsaPred<VectorType>);
+ auto cachedSupportsMMAMatrixType = [&](Operation *currentOp) {
+ auto [it, inserted] =
+ supportsMMAMatrixTypeCache.try_emplace(currentOp, false);
+ if (inserted)
+ it->second = supportsMMaMatrixType(currentOp, useNvGpu);
+ return it->second;
};
- ForwardSliceOptions forwardSliceOptions;
- forwardSliceOptions.filter = hasVectorSrc;
SetVector<Operation *> opToConvert;
op->walk([&](Operation *nestedOp) {
if (!isa<vector::ContractionOp>(nestedOp) &&
!elementwiseSupportsMMAMatrixType(nestedOp))
return;
- if (opToConvert.contains(nestedOp))
+ if (analyzedOps.contains(nestedOp))
return;
- SetVector<Operation *> dependentOps =
- getSliceContract(nestedOp, backwardSliceOptions, forwardSliceOptions);
+
+ SetVector<Operation *> dependentOps;
+ dependentOps.insert(nestedOp);
+ unsigned currentIndex = 0;
+ while (currentIndex != dependentOps.size()) {
+ Operation *currentOp = dependentOps[currentIndex++];
+ dependentOps.insert_range(getCachedBackwardSlice(currentOp));
+ dependentOps.insert_range(getCachedForwardSlice(currentOp));
+ }
+
+ analyzedOps.insert_range(dependentOps);
+
// If any instruction cannot use MMA matrix type drop the whole
// chain. MMA matrix are stored in an opaque type so they cannot be used
// by all operations.
- if (llvm::any_of(dependentOps, [useNvGpu](Operation *op) {
- if (!supportsMMaMatrixType(op, useNvGpu)) {
+ if (llvm::any_of(dependentOps, [&](Operation *op) {
+ if (!cachedSupportsMMAMatrixType(op)) {
LDBG() << "cannot convert op: " << *op;
return true;
}
>From a87e82887c997a7c8ed2b5318c9cc3794c0c00a8 Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Wed, 9 Sep 2026 16:00:55 +0100
Subject: [PATCH 2/3] Addressed comments
Change-Id: I2a469269b3a9addf7f5a9db929bfbe0f9f64272c
---
mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index ffbb113758c48..f79402a25dc1d 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -359,7 +359,7 @@ static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
getBackwardSlice(currentOp, &backwardSlice, backwardSliceOptions);
assert(result.succeeded() && "expected a backward slice");
(void)result;
- it->second.append(backwardSlice.begin(), backwardSlice.end());
+ it->second = backwardSlice.takeVector();
return it->second;
};
@@ -382,7 +382,7 @@ static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
} else {
getForwardSlice(currentOp, &forwardSlice, forwardSliceOptions);
}
- it->second.append(forwardSlice.begin(), forwardSlice.end());
+ it->second = forwardSlice.takeVector();
return it->second;
};
>From 0dd5d0f9aebca318ea38995087e99f15fc914448 Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Thu, 10 Sep 2026 09:29:58 +0100
Subject: [PATCH 3/3] Addressed comments
Change-Id: I2d4a4d78006ff5a6b9e495cf3f38dc45e52a4681
---
mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp | 8 +-------
1 file changed, 1 insertion(+), 7 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index f79402a25dc1d..f19e8290987ed 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -344,9 +344,6 @@ static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
DenseMap<Operation *, SmallVector<Operation *>> backwardSliceCache;
DenseMap<Operation *, SmallVector<Operation *>> forwardSliceCache;
DenseMap<Operation *, bool> supportsMMAMatrixTypeCache;
- // Once we have expanded a seed to its full use/def component, every op in
- // that component would rediscover the same set again.
- DenseSet<Operation *> analyzedOps;
auto getCachedBackwardSlice =
[&](Operation *currentOp) -> ArrayRef<Operation *> {
@@ -399,7 +396,7 @@ static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
if (!isa<vector::ContractionOp>(nestedOp) &&
!elementwiseSupportsMMAMatrixType(nestedOp))
return;
- if (analyzedOps.contains(nestedOp))
+ if (backwardSliceCache.contains(nestedOp))
return;
SetVector<Operation *> dependentOps;
@@ -410,9 +407,6 @@ static SetVector<Operation *> getOpToConvert(mlir::Operation *op,
dependentOps.insert_range(getCachedBackwardSlice(currentOp));
dependentOps.insert_range(getCachedForwardSlice(currentOp));
}
-
- analyzedOps.insert_range(dependentOps);
-
// If any instruction cannot use MMA matrix type drop the whole
// chain. MMA matrix are stored in an opaque type so they cannot be used
// by all operations.
More information about the Mlir-commits
mailing list