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

[llvmlistbot at llvm.org]

Author: Aart Bik
Date: 2024-02-05T16:12:47-08:00
New Revision: 5a9af39aab40bba52d4e46cabf4b1ab47f614fa2

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

LOG: [mlir][sparse] made sparse vectorizer more robust on position of invariants (#80766)

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
    mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index 3a487a3bd6a06..2b81d6cdc1eab 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.
+      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> {