[Mlir-commits] [mlir] [mlir][sparse] Fix crash on linalg ops with buffer semantics (PR #216458)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir-gpu

Author: Hamza Qureshi (hamzaqureshi5)

<details>
<summary>Changes</summary>

Fixes #<!-- -->216215

## The problem

`mlir-opt --sparse-gpu-codegen` crashes on valid input:

```
Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.
```

The pass looks for `linalg.generic` ops shaped like a matmul so it can turn them into GPU library calls. Its checks only look at the *shape* of the operation (loop count, iterator types, indexing maps, the multiply-add body) — none of them check whether the operands are tensors or memrefs.

Sparsity only exists on tensors, so when the pass then asks "is this operand a dense tensor?", it casts the operand type to `RankedTensorType`. If the op has buffer semantics, that operand is a `memref`, and the cast asserts.

## The fix

Reject ops with buffer semantics before any of that happens:

```cpp
if (!op.hasPureTensorSemantics())
  return failure();
```

`return failure()` just means "this pattern doesn't apply here", so the op is left alone — the correct outcome for memrefs.

This is the same guard the linalg patterns in `SparseReinterpretMap.cpp` and `SparseTensorRewriting.cpp` already use. Putting it at the entry point covers every rewrite reached through this pattern (SpMV, SpMM, SpGEMM, SDDMM, 2:4-SpMM), since they all share it.

## Testing

Added `mlir/test/Dialect/SparseTensor/GPU/gpu_buffer_semantics.mlir` with the matmul case from the issue plus a matvec case, checking the ops are left untouched.

Verified with an assertions-enabled build:
- the reproducer aborts before the fix and exits cleanly after
- the new test fails without the fix and passes with it
- `mlir/test/Dialect/SparseTensor` is 114/114, no regressions

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp (+5) 
- (added) mlir/test/Dialect/SparseTensor/GPU/gpu_buffer_semantics.mlir (+57) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index d7f46e1676243..14418b09f8354 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
@@ -1298,6 +1298,11 @@ struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
                                 PatternRewriter &rewriter) const override {
     if (op.getNumDpsInits() != 1)
       return failure(); // reject multi-output
+    // The rewrites below query the sparse tensor encoding of the operands,
+    // which is only defined for ranked tensor types. Reject buffer semantics
+    // rather than casting a memref operand to a tensor type.
+    if (!op.hasPureTensorSemantics())
+      return failure();
 
     const unsigned numLoops = op.getNumLoops();
     const unsigned numTensors = op->getNumOperands();
diff --git a/mlir/test/Dialect/SparseTensor/GPU/gpu_buffer_semantics.mlir b/mlir/test/Dialect/SparseTensor/GPU/gpu_buffer_semantics.mlir
new file mode 100644
index 0000000000000..b16cb9526083d
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/GPU/gpu_buffer_semantics.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-opt %s --split-input-file --sparse-gpu-codegen="num-threads=0" | FileCheck %s
+
+// Verifies that the sparse GPU rewrites leave linalg operations with buffer
+// semantics alone. Their operands have memref types, which do not carry a
+// sparse tensor encoding, so the ops must not be matched by the rewrites.
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func.func @matmul_buffer_semantics
+// CHECK-NOT:     gpu.
+// CHECK:         linalg.generic
+// CHECK:         return
+func.func @matmul_buffer_semantics(%arga: memref<4x8xf32>,
+                                   %argb: memref<8x6xf32>,
+                                   %argc: memref<4x6xf32>) {
+  linalg.generic {
+    indexing_maps = [#map, #map1, #map2],
+    iterator_types = ["parallel", "parallel", "reduction"]
+  }
+  ins(%arga, %argb : memref<4x8xf32>, memref<8x6xf32>)
+  outs(%argc : memref<4x6xf32>) {
+  ^bb0(%a: f32, %b: f32, %c: f32):
+    %0 = arith.mulf %a, %b : f32
+    %1 = arith.addf %c, %0 : f32
+    linalg.yield %1 : f32
+  }
+  return
+}
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+#map2 = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL: func.func @matvec_buffer_semantics
+// CHECK-NOT:     gpu.
+// CHECK:         linalg.generic
+// CHECK:         return
+func.func @matvec_buffer_semantics(%arga: memref<4x8xf32>,
+                                   %argb: memref<8xf32>,
+                                   %argc: memref<4xf32>) {
+  linalg.generic {
+    indexing_maps = [#map, #map1, #map2],
+    iterator_types = ["parallel", "reduction"]
+  }
+  ins(%arga, %argb : memref<4x8xf32>, memref<8xf32>)
+  outs(%argc : memref<4xf32>) {
+  ^bb0(%a: f32, %b: f32, %c: f32):
+    %0 = arith.mulf %a, %b : f32
+    %1 = arith.addf %c, %0 : f32
+    linalg.yield %1 : f32
+  }
+  return
+}

``````````

</details>


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