Skip to content

Commit 5a9af39

Browse files
authored
[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: #79897
1 parent dbed898 commit 5a9af39

File tree

2 files changed

+7
-2
lines changed

2 files changed

+7
-2
lines changed

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,12 @@ static bool vectorizeSubscripts(PatternRewriter &rewriter, scf::ForOp forOp,
316316
if (auto load = cast.getDefiningOp<arith::AddIOp>()) {
317317
Value inv = load.getOperand(0);
318318
Value idx = load.getOperand(1);
319+
// Swap non-invariant.
320+
if (!isInvariantValue(inv, block)) {
321+
inv = idx;
322+
idx = load.getOperand(0);
323+
}
324+
// Inspect.
319325
if (isInvariantValue(inv, block)) {
320326
if (auto arg = llvm::dyn_cast<BlockArgument>(idx)) {
321327
if (isInvariantArg(arg, block) || !innermost)

mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
// FIXME: re-enable.
21
// RUN: mlir-opt %s -sparsifier="vl=8" | FileCheck %s
32

43
#Dense = #sparse_tensor.encoding<{
@@ -16,7 +15,7 @@
1615
}
1716

1817
// CHECK-LABEL: llvm.func @kernel_matvec
19-
// C_HECK: llvm.intr.vector.reduce.fadd
18+
// CHECK: llvm.intr.vector.reduce.fadd
2019
func.func @kernel_matvec(%arga: tensor<?x?xf32, #Dense>,
2120
%argb: tensor<?xf32>,
2221
%argx: tensor<?xf32>) -> tensor<?xf32> {

0 commit comments

Comments
 (0)