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

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/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

>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] [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 3a487a3bd6a069..011e66073ff3ff 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 dfee2b1261b6cc..e25c3a02f91271 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> {