[mlir][sparse] Use vector.step for index vector generation

[c-rhodes]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Cullen Rhodes (c-rhodes)

Changes

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp (+1-12)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index 2b81d6cdc1eab..7a2400b41b045 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -381,18 +381,7 @@ static bool vectorizeExpr(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
       if (codegen) {
         VectorType vtp = vectorType(vl, arg.getType());
         Value veci = rewriter.create<vector::BroadcastOp>(loc, vtp, arg);
-        Value incr;
-        if (vl.enableVLAVectorization) {
-          Type stepvty = vectorType(vl, rewriter.getI64Type());
-          Value stepv = rewriter.create<LLVM::StepVectorOp>(loc, stepvty);
-          incr = rewriter.create<arith::IndexCastOp>(loc, vtp, stepv);
-        } else {
-          SmallVector<APInt> integers;
-          for (unsigned i = 0, l = vl.vectorLength; i < l; i++)
-            integers.push_back(APInt(/*width=*/64, i));
-          auto values = DenseElementsAttr::get(vtp, integers);
-          incr = rewriter.create<arith::ConstantOp>(loc, vtp, values);
-        }
+        Value incr = rewriter.create<vector::StepOp>(loc, vtp);
         vexp = rewriter.create<arith::AddIOp>(loc, veci, incr);
       }
       return true;

[MacDue]

LG, I guess there's no unit tests for this?

[c-rhodes]

LG, I guess there's no unit tests for this?

I did check there's coverage by adding an unreachable, there's mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir for fixed and also an integration test (can't remember which). The fold must kick in so the unit test remains unchanged.

[banach-space]

This is a very nice improvement, thanks!

There are no tests, but this PR is not adding any new functionality. Also, SparseTensor is very thoroughly tested through integration tests. I think that it's safe to assume that this codepath is indeed already covered by tests.

LGTM!