[Mlir-commits] [mlir] [Vector] Do not sink elementwise on dynamic position (PR #148245)

[Kunwar Grover]

https://github.com/Groverkss created https://github.com/llvm/llvm-project/pull/148245

Propagating vector.extract when a dynamic position is present can cause dominance issues and needs better handling. For now, disable propagation if there is a dynamic position present.

>From 6b95277e28da536df235551975c2074c045df2ae Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Fri, 11 Jul 2025 15:04:03 +0100
Subject: [PATCH] [Vector] Do not sink elementwise on dynamic position

---
 .../Dialect/Vector/Transforms/VectorTransforms.cpp   |  6 ++++++
 mlir/test/Dialect/Vector/vector-sink.mlir            | 12 ++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index bcaea1c79471f..fe2707629d82e 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1088,6 +1088,12 @@ class ExtractOpFromElementwise final
     if (!llvm::all_equal(eltwise->getOperandTypes()))
       return rewriter.notifyMatchFailure(op, "operand types are different");
 
+    // Dynamic position can cause dominance issues, so conservatively fail for
+    // now.
+    if (!op.getDynamicPosition().empty())
+      return rewriter.notifyMatchFailure(
+          op, "dynamic position not yet implemented");
+
     Type dstType = op.getType();
 
     OpBuilder::InsertionGuard g(rewriter);
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index 900ad99bb4a4c..b826cdca134e6 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -514,6 +514,18 @@ func.func @negative_extract_vec_fma(%arg0: vector<4xf32>, %arg1: vector<4xf32>,
   return %1 : f32
 }
 
+// CHECK-LABEL: @negative_extract_dynamic_pos
+func.func @negative_extract_dynamic_pos(%arg0: vector<4xf32>, %arg1 : vector<4xf32>, %idx : vector<4xindex>) -> f32 {
+  // CHECK-NOT: vector.extract
+  // CHECK: arith.addf %{{.*}}, %{{.*}} : vector<4xf32>
+  // CHECK: vector.extract
+  // CHECK: vector.extract
+  %0 = arith.addf %arg0, %arg1 : vector<4xf32>
+  %1 = vector.extract %idx[0] : index from vector<4xindex>
+  %2 = vector.extract %0[%1] : f32 from vector<4xf32>
+  return %2 : f32
+}
+
 //-----------------------------------------------------------------------------
 // [Pattern: ExtractOpFromLoad]
 //-----------------------------------------------------------------------------