[Mlir-commits] [mlir] d287e55 - [MLIR][VectorToGPU] Cache slice analysis in getOpToConvert (#221795)

[llvmlistbot at llvm.org]

Author: Dmitriy Smirnov
Date: 2026-10-01T11:34:27+01:00
New Revision: d287e55d03029586399adab02977679e74c51f2d

URL: https://github.com/llvm/llvm-project/commit/d287e55d03029586399adab02977679e74c51f2d
DIFF: https://github.com/llvm/llvm-project/commit/d287e55d03029586399adab02977679e74c51f2d.diff

LOG: [MLIR][VectorToGPU] Cache slice analysis in getOpToConvert (#221795)

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.

---------

Co-authored-by: Thomas Preud'homme <thomas.preudhomme at arm.com>

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 4ed868eca25c7..194b47011b86c 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -324,30 +324,48 @@ static bool supportsMMaMatrixType(Operation *op, bool useNvGpu) {
   return elementwiseSupportsMMAMatrixType(op);
 }
 
-/// Return an unsorted slice handling scf.for region 
diff erently 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;
+
+  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 = backwardSlice.takeVector();
+    return it->second;
+  };
 
-    // Compute and insert the forwardSlice starting from currentOp.
-    forwardSlice.clear();
+  auto getCachedForwardSlice =
+      [&](Operation *currentOp) -> ArrayRef<Operation *> {
+    auto [it, inserted] = forwardSliceCache.try_emplace(currentOp);
+    if (!inserted)
+      return it->second;
+
+    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
@@ -360,42 +378,39 @@ 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 = forwardSlice.takeVector();
+    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 (backwardSliceCache.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));
+    }
     // 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;
           }