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

[Aart Bik]

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

>From 4b9ad2c723e51e9814500eacb0ae44895f32413b Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 5 Feb 2024 15:53:13 -0800
Subject: [PATCH 1/2] [mlir][sparse] made sparse vectorizer more robust on
 position of invariants

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
---
 .../Dialect/SparseTensor/Transforms/SparseVectorization.cpp | 6 ++++++
 mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir        | 3 +--
 2 files changed, 7 insertions(+), 2 deletions(-)

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> {

>From d7f208ca8475eed262fe6bcc0763050fc0dc15f6 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 5 Feb 2024 15:59:07 -0800
Subject: [PATCH 2/2] typo

---
 .../lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index 011e66073ff3f..2b81d6cdc1eab 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -316,7 +316,7 @@ 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.
+      // Swap non-invariant.
       if (!isInvariantValue(inv, block)) {
         inv = idx;
         idx = load.getOperand(0);