PassManager fails on simple COO addition example

MLIR
1

Hi!
I’m trying to run a simple COO elemwise addition example. Inputs and outputs of the operation are (COO, COO) -> COO. But when running PassManager I get an error:

mlir._mlir_libs._site_initialize.<locals>.MLIRError: Failure while executing pass pipeline:
error: "-":6:5: input/output trailing COO level-ranks don't match
 note: "-":6:5: see current operation: %0 = "sparse_tensor.assemble"(%arg0, %arg1, %arg2) : (tensor<?xi64>, tensor<?xi64>, tensor<?xf64>) -> tensor<3x4xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton), posWidth = 64, crdWidth = 64 }>>
error: "-":6:5: input/output trailing COO level-ranks don't match
 note: "-":6:5: see current operation: %1 = "sparse_tensor.assemble"(%arg3, %arg4, %arg5) : (tensor<?xi64>, tensor<?xi64>, tensor<?xf64>) -> tensor<3x4xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton), posWidth = 64, crdWidth = 64 }>>
error: "-":6:5: type of return operand 1 ('memref<?xi64, strided<[?], offset: ?>>') doesn't match function result type ('memref<?xi64>') in function @add
 note: "-":6:5: see current operation: "func.return"(%3, %4, %5) : (memref<?xi64>, memref<?xi64, strided<[?], offset: ?>>, memref<?xf64>) -> ()

When replacing COO format with CSR the parsing completes successfully. I’m using LLVM 19.1.0-rc3.

Here’s a small GitHub repository that reproduces this issue: GitHub - mtsokol/mlir-coo-reproduce

Here’s a script: mlir-coo-reproduce/reproduce.py at main · mtsokol/mlir-coo-reproduce · GitHub

And a CI job that shows the issue: Initial commit · mtsokol/mlir-coo-reproduce@a57d5ef · GitHub

2

Could you try with singleton(soa) ?

3

@kromancer It looks like singleton(soa) fixes my issue - thank you! I can run passes and call my COO+COO=>COO operation from Python (after adjusting COO C structure to SoA convention).

But it looks like SoA attribute isn’t available in Python - here’s a PR that adds SoA to bindings: [MLIR][sparse] Add `soa` property to `sparse_tensor` Python bindings by mtsokol · Pull Request #109135 · llvm/llvm-project · GitHub