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

[Hamza Qureshi]

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

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

>From dac597449ce0900048c69bd38728b815320153eb Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Sat, 15 Aug 2026 11:49:07 +0500
Subject: [PATCH] [mlir][sparse] Fix crash on linalg ops with buffer semantics

LinalgOpRewriter matched any linalg.generic whose structure looked like a
sparse GPU kernel, but the rewrites it dispatches to query the sparse tensor
encoding of the operands via getSparseTensorType, which casts the operand type
to RankedTensorType. None of the structural checks reject buffer semantics, so
a matching linalg.generic on memrefs made that cast assert.

Reject ops without pure tensor semantics up front, matching the guard already
used by the linalg patterns in SparseReinterpretMap and SparseTensorRewriting.
This covers all of the rewrites reached through this pattern, since they share
the same entry point.

Fixes #216215
---
 .../Transforms/SparseGPUCodegen.cpp           |  5 ++
 .../GPU/gpu_buffer_semantics.mlir             | 57 +++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 mlir/test/Dialect/SparseTensor/GPU/gpu_buffer_semantics.mlir

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
+}