Illegal operation when slicing CSR/CSC/COO tensor

MLIR
1

Hi! I’m trying to access a sub-tensor out of a sparse tensor with a extract_slice, while keeping the same sparse format for the result.

I followed one example in the test suite:

Here’s my Python implementation, which extracts a sub-tensor out of CSR. Where I want to keep it as CSR:

import ctypes, ctypes.util, pathlib
from mlir.ir import Context, Module
from mlir import execution_engine, passmanager

MLIR_C_RUNNER_UTILS = ctypes.util.find_library("mlir_c_runner_utils")

with Context():
    module = Module.parse(
    """
    #sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64 }>
    module {
    func.func @getitem(%arg0: tensor<80x100xf64, #sparse>) -> tensor<2x2xf64, #sparse> attributes {llvm.emit_c_interface} {
        %extracted_slice = tensor.extract_slice %arg0[0, 0] [20, 20] [1, 1] : tensor<80x100xf64, #sparse> to tensor<20x20xf64, #sparse>
        return %extracted_slice : tensor<20x20xf64, #sparse>
    }
    }
    """
    )

    CWD = pathlib.Path(".")
    (CWD / "module.mlir").write_text(str(module))

    pm = passmanager.PassManager.parse("builtin.module(sparse-assembler{direct-out=true}, sparsifier{create-sparse-deallocs=1 enable-runtime-library=false})")
    pm.run(module.operation)

    execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])

Unfortunately it fails with:

Traceback (most recent call last):
  File "/home/mtsokol/sparse/slice_reproduce.py", line 24, in <module>
    pm.run(module.operation)
mlir._mlir_libs._site_initialize.<locals>.MLIRError: Failure while executing pass pipeline:
error: "-":6:9: failed to legalize operation 'func.return' that was explicitly marked illegal
 note: "-":6:9: see current operation: "func.return"(%1) : (tensor<20x20xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64 }>>) -> ()

saying that the operation is illegal. Is it possible to use extract_slice for extracting CSR out of CSR tensor? (without #sparse_tensor<slice(...)> attribute)

[EDIT] It also happens for CSC, COO and Dense formats.