[Mlir-commits] [mlir] [mlir][sparse] Fix loop-invariant detection for block arguments (PR #236231)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Jayanth Kandi (jayanth260)

<details>
<summary>Changes</summary>

## Summary

Fix sparse vectorization failing to recognize function arguments as loop-invariant values.

Previously, `isInvariantValue()` relied on `getDefiningOp()` to determine whether a value was defined outside the loop. Since block arguments have no defining operation, function arguments were incorrectly classified as non-invariant.

This prevented vectorization of supported patterns such as `a[base + i]`, where `base` is a function argument, as well as shift operations with invariant shift amounts supplied as function arguments.

## Changes

- Extend `isInvariantValue()` to handle `BlockArgument` values using their owning block.
- Preserve the existing handling of operation results.
- Add regression tests covering invariant base offsets, reversed addition operands, invariant shift amounts, and non-contiguous memory access.

## Testing

- Verified that loops using `a[base + i]` are vectorized with `vl=4`.
- Verified that invariant shift amounts are broadcast and vectorized.
- Verified that non-contiguous accesses remain scalar.
- Ran the SparseTensor regression test suite: **121/121 tests passed**.

Fixes #<!-- -->228126

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp (+9-5) 
- (added) mlir/test/Dialect/SparseTensor/sparse_vector_invariant_args.mlir (+89) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index b5ec17c88b2050..68cbc3c74fafbf 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -44,15 +44,19 @@ struct VL {
   bool enableSIMDIndex32;
 };
 
-/// Helper test for invariant value (defined outside given block).
-static bool isInvariantValue(Value val, Block *block) {
-  return val.getDefiningOp() && val.getDefiningOp()->getBlock() != block;
-}
-
 /// Helper test for invariant argument (defined outside given block).
 static bool isInvariantArg(BlockArgument arg, Block *block) {
   return arg.getOwner() != block;
 }
+/// Helper test for invariant value (defined outside given block).
+/// Handles both block arguments and operation results.
+static bool isInvariantValue(Value val, Block *block) {
+  if (auto arg = dyn_cast<BlockArgument>(val))
+    return isInvariantArg(arg, block);
+    
+  return val.getDefiningOp()->getBlock() != block;
+}
+
 
 /// Returns true when `mem`'s most minor dimension has a statically known
 /// non-unit stride.
diff --git a/mlir/test/Dialect/SparseTensor/sparse_vector_invariant_args.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_invariant_args.mlir
new file mode 100644
index 00000000000000..c01f2635758f19
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_invariant_args.mlir
@@ -0,0 +1,89 @@
+// RUN: mlir-opt %s --sparse-vectorization='vl=4' | FileCheck %s
+
+// CHECK-LABEL: func.func @invariant_base(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.load
+// CHECK: arith.muli {{.*}} : vector<4xi32>
+// CHECK: vector.store
+
+func.func @invariant_base(%a: memref<?xi32>, %base: index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+  %two = arith.constant 2 : i32
+
+  scf.for %i = %c0 to %c64 step %c1 {
+    %p = arith.addi %base, %i : index
+    %x = memref.load %a[%p] : memref<?xi32>
+    %y = arith.muli %x, %two : i32
+    memref.store %y, %a[%p] : memref<?xi32>
+  } {"Emitted from" = "linalg.generic"}
+
+  return
+}
+
+// CHECK-LABEL: func.func @invariant_shift(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.load
+// CHECK: vector.broadcast {{.*}} : i32 to vector<4xi32>
+// CHECK: arith.shli {{.*}} : vector<4xi32>
+// CHECK: vector.store
+
+func.func @invariant_shift(%a: memref<?xi32>, %shift: i32) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+
+  scf.for %i = %c0 to %c64 step %c1 {
+    %x = memref.load %a[%i] : memref<?xi32>
+    %y = arith.shli %x, %shift : i32
+    memref.store %y, %a[%i] : memref<?xi32>
+  } {"Emitted from" = "linalg.generic"}
+
+  return
+}
+
+// CHECK-LABEL: func.func @invariant_base_reversed(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: vector.load
+// CHECK: arith.muli {{.*}} : vector<4xi32>
+// CHECK: vector.store
+
+func.func @invariant_base_reversed(%a: memref<?xi32>, %base: index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+  %two = arith.constant 2 : i32
+
+  scf.for %i = %c0 to %c64 step %c1 {
+    %p = arith.addi %i, %base : index
+    %x = memref.load %a[%p] : memref<?xi32>
+    %y = arith.muli %x, %two : i32
+    memref.store %y, %a[%p] : memref<?xi32>
+  } {"Emitted from" = "linalg.generic"}
+
+  return
+}
+
+// CHECK-LABEL: func.func @non_contiguous(
+// CHECK: scf.for {{.*}} step %{{.*}} {
+// CHECK: arith.muli {{.*}} : index
+// CHECK: memref.load
+// CHECK: arith.muli {{.*}} : i32
+// CHECK: memref.store
+
+func.func @non_contiguous(%a: memref<?xi32>, %stride: index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+  %two = arith.constant 2 : i32
+
+  scf.for %i = %c0 to %c64 step %c1 {
+    %p = arith.muli %stride, %i : index
+    %x = memref.load %a[%p] : memref<?xi32>
+    %y = arith.muli %x, %two : i32
+    memref.store %y, %a[%p] : memref<?xi32>
+  } {"Emitted from" = "linalg.generic"}
+
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/236231