Assembling a CSC matrix with Python bindings

Beginners
,
1

Hi!

I’m writing with a question regarding Python Bindings and tensor-sparse dialect.

I’m trying to write code that takes three arrays and assembles a CSR/CSC matrix. For CSR format the implementation below works well (with permutation [0,1]), but for CSC (permutation [1,0]) it fails with:

error: unknown: the sparse-tensor must have the identity mapping
 note: unknown: see current operation: %0 = "sparse_tensor.assemble"(%arg2, %arg0, %arg1) : (tensor<?xf64>, tensor<?xi64>, tensor<?xi64>) -> tensor<100x100xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d1 : dense, d0 : compressed), posWidth = 64, crdWidth = 64 }>>

Here’s full python code to reproduce:

import ctypes
import ctypes.util

import mlir.passmanager
from mlir import ir
from mlir.dialects import arith, func, linalg, sparse_tensor, tensor

with ir.Context(), ir.Location.unknown():
    module = ir.Module.create()
    f64 = ir.F64Type.get()
    i64 = ir.IntegerType.get_signless(64)

    levels = (sparse_tensor.LevelType.dense, sparse_tensor.LevelType.compressed)
    ordering = ir.AffineMap.get_permutation([1, 0])
    encoding = sparse_tensor.EncodingAttr.get(levels, ordering, ordering, 64, 64)
    csc_shaped = ir.RankedTensorType.get([100, 100], f64, encoding)

    tensor_1d_i64 = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], i64)
    tensor_1d_f64 = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], f64)

    with ir.InsertionPoint(module.body):
        @func.FuncOp.from_py_func(tensor_1d_i64, tensor_1d_i64, tensor_1d_f64)
        def assemble(pos, crd, data):
         return sparse_tensor.assemble(csc_shaped, data, (pos, crd))

    assemble.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
    pm = mlir.passmanager.PassManager.parse("builtin.module(sparsifier{})")
    pm.run(module.operation)

How can I assemble a CSC matrix with sparse_tensor.assemble here?

2

For me the solution was to move to sparse-assembler{direct-out=true} which takes care of assembling and disassembling tensors. Now I can pass constituent arrays directly.