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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Kunwar Grover (Groverkss)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/148245.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+6) 
- (modified) mlir/test/Dialect/Vector/vector-sink.mlir (+12) 


``````````diff
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]
 //-----------------------------------------------------------------------------

``````````

</details>


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