[Mlir-commits] [mlir] [mlir][sparse] made sparse vectorizer more robust on position of invariants (PR #80766)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Aart Bik (aartbik)

<details>
<summary>Changes</summary>

Because the sparse vectorizer relies on the code coming out of the sparsifier, the "patterns" are not always made very general. However, a recent change in the generated code revealed an obvious situation where the subscript analysis could be made a bit more robust.

Fixes:
https://github.com/llvm/llvm-project/issues/79897

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp (+6) 
- (modified) mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir (+1-2) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index 3a487a3bd6a06..011e66073ff3f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -316,6 +316,12 @@ static bool vectorizeSubscripts(PatternRewriter &rewriter, scf::ForOp forOp,
     if (auto load = cast.getDefiningOp<arith::AddIOp>()) {
       Value inv = load.getOperand(0);
       Value idx = load.getOperand(1);
+      // Swap non-invariant left.
+      if (!isInvariantValue(inv, block)) {
+        inv = idx;
+        idx = load.getOperand(0);
+      }
+      // Inspect.
       if (isInvariantValue(inv, block)) {
         if (auto arg = llvm::dyn_cast<BlockArgument>(idx)) {
           if (isInvariantArg(arg, block) || !innermost)
diff --git a/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir
index dfee2b1261b6c..e25c3a02f9127 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir
@@ -1,4 +1,3 @@
-// FIXME: re-enable.
 // RUN: mlir-opt %s -sparsifier="vl=8" |  FileCheck %s
 
 #Dense = #sparse_tensor.encoding<{
@@ -16,7 +15,7 @@
 }
 
 // CHECK-LABEL: llvm.func @kernel_matvec
-// C_HECK:       llvm.intr.vector.reduce.fadd
+// CHECK:       llvm.intr.vector.reduce.fadd
 func.func @kernel_matvec(%arga: tensor<?x?xf32, #Dense>,
                          %argb: tensor<?xf32>,
 			 %argx: tensor<?xf32>) -> tensor<?xf32> {

``````````

</details>


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