This page was generated from dialect/graphblas_dialect_tutorials/graphblas_lower/sparse_layouts.ipynb.

Working with Sparse Layouts in the GraphBLAS Dialect

This example will go over how to use the --graphblas-lower pass from graphblas-opt to lower the GraphBLAS dialect ops that directly manipulate the layouts of sparse tensors. In particular, we’ll focus on the graphblas.convert_layout and graphblas.transpose ops.

Since the ops reference already documents these ops with examples, we’ll only briefly describe them here.

Let’s first import some necessary modules and generate an instance of our JIT engine.

import mlir_graphblas
from mlir_graphblas.tools.utils import sparsify_array
import numpy as np

engine = mlir_graphblas.MlirJitEngine()

Using development graphblas-opt: /Users/pnguyen/code/mlir-graphblas/mlir_graphblas/src/build/bin/graphblas-opt

Overview of graphblas.convert_layout

Here, we’ll show how to use the graphblas.convert_layout op.

This op takes 1 sparse matrix in CSR or CSC format and creates a new sparse matrix of the desired format.

We’ll give several examples below of how this will work.

First, we’ll define an example input CSR matrix.

dense_matrix = np.array(
    [
        [1.1, 0. , 0. , 0. ],
        [0. , 0. , 2.2, 0. ],
        [0. , 0. , 0. , 0. ],
        [0. , 0. , 0. , 0. ]
    ],
    dtype=np.float64,
)

csr_matrix = sparsify_array(dense_matrix, [False, True])

graphblas.convert_layout (CSR->CSC)

Let’s convert this matrix to CSC format.

mlir_text = """
#CSR64 = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (i,j)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

#CSC64 = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (j,i)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

func @csr_to_csc(%sparse_tensor: tensor<?x?xf64, #CSR64>) -> tensor<?x?xf64, #CSC64> {
    %answer = graphblas.convert_layout %sparse_tensor : tensor<?x?xf64, #CSR64> to tensor<?x?xf64, #CSC64>
    return %answer : tensor<?x?xf64, #CSC64>
}
"""

Here are the passes we’ll use.

passes = [
    "--graphblas-structuralize",
    "--graphblas-dwim",
    "--graphblas-optimize",
    "--graphblas-lower",
    "--sparsification",
    "--sparse-tensor-conversion",
    "--linalg-bufferize",
    "--func-bufferize",
    "--tensor-bufferize",
    "--finalizing-bufferize",
    "--convert-linalg-to-loops",
    "--convert-vector-to-llvm",
    "--convert-math-to-llvm",
    "--convert-math-to-libm",
    "--convert-scf-to-cf",
    "--convert-memref-to-llvm",
    "--convert-openmp-to-llvm",
    "--convert-arith-to-llvm",
    "--convert-std-to-llvm",
    "--reconcile-unrealized-casts",
]

engine.add(mlir_text, passes)

['csr_to_csc']

csc_matrix = engine.csr_to_csc(csr_matrix)

Let’s look at the native arrays for csc_matrix, as well as the dense version.

print(csc_matrix.pointers)
print(csc_matrix.indices)
print(csc_matrix.values)

(array([], dtype=uint64), array([0, 1, 1, 2, 2], dtype=uint64))
(array([], dtype=uint64), array([0, 1], dtype=uint64))
[1.1 2.2]

csc_matrix.toarray()

array([[1.1, 0. , 0. , 0. ],
       [0. , 0. , 2.2, 0. ],
       [0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. ]])

np.all(dense_matrix == csc_matrix.toarray())

True

Overview of graphblas.transpose

Here, we’ll show how to use the graphblas.transpose op.

graphblas.transpose returns a new sparse matrix that’s the transpose of the input matrix. Note that the behavior of this op differs depending on the sparse encoding of the specified output tensor type.

The input/output behavior of graphblas.transpose is fairly simple. Our examples here aren’t intended to show anything interesting but to merely act as reproducible references.

The important thing to know about graphblas.transpose is how it is implemented.

When transposing a CSC matrix to a CSR matrix, we simply need to swap the dimension sizes and reverse the indexing. Thus, the only “real” work done here is changing metadata. The same goes for transposing a CSC matrix to a CSR matrix.

Here’s an example of transposing a CSR matrix to a CSC matrix.

dense_matrix = np.array(
    [
        [1.1, 0. , 0. , 0. ],
        [0. , 0. , 2.2, 0. ],
    ],
    dtype=np.float64,
)

csr_matrix = sparsify_array(dense_matrix, [False, True])

mlir_text = """
#CSR64 = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (i,j)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

#CSC64 = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (j,i)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

func @transpose_csr_to_csc(%sparse_tensor: tensor<?x?xf64, #CSR64>) -> tensor<?x?xf64, #CSC64> {
    %answer = graphblas.transpose %sparse_tensor : tensor<?x?xf64, #CSR64> to tensor<?x?xf64, #CSC64>
    return %answer : tensor<?x?xf64, #CSC64>
}
"""

engine.add(mlir_text, passes)

['transpose_csr_to_csc']

csc_matrix_transpose = engine.transpose_csr_to_csc(csr_matrix)

csc_matrix_transpose.toarray()

array([[1.1, 0. ],
       [0. , 0. ],
       [0. , 2.2],
       [0. , 0. ]])

Python Utilities for MLIR’s Sparse Tensors Vector Ops