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()

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-optimize",
    "--graphblas-lower",
    "--sparsification",
    "--sparse-tensor-conversion",
    "--linalg-bufferize",
    "--func-bufferize",
    "--tensor-constant-bufferize",
    "--tensor-bufferize",
    "--finalizing-bufferize",
    "--convert-linalg-to-loops",
    "--convert-scf-to-std",
    "--convert-memref-to-llvm",
    "--convert-math-to-llvm",
    "--convert-openmp-to-llvm",
    "--convert-arith-to-llvm",
    "--convert-math-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)

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

graphblas.convert_layout (CSC->CSR)

Let’s convert the CSC matrix back to CSR format.

Let’s first get rid of our original csr_matrix so we don’t get correct results purely by accident.

del csr_matrix

Here’s the MLIR code to convert from CSC to CSR.

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 @csc_to_csr(%sparse_tensor: tensor<?x?xf64, #CSC64>) -> tensor<?x?xf64, #CSR64> {
    %answer = graphblas.convert_layout %sparse_tensor : tensor<?x?xf64, #CSC64> to tensor<?x?xf64, #CSR64>
    return %answer : tensor<?x?xf64, #CSR64>
}
"""

engine.add(mlir_text, passes)

['csc_to_csr']

csr_matrix = engine.csc_to_csr(csc_matrix)

csr_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 == csr_matrix.toarray())

True

graphblas.convert_layout (CSC->CSC, CSR->CSR)

For completeness, we’ll show how to convert to and from the same exact layouts.

The MLIR code to do so is shown below.

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 @csc_to_csc(%sparse_tensor: tensor<?x?xf64, #CSC64>) -> tensor<?x?xf64, #CSC64> {
    %answer = graphblas.convert_layout %sparse_tensor : tensor<?x?xf64, #CSC64> to tensor<?x?xf64, #CSC64>
    return %answer : tensor<?x?xf64, #CSC64>
}

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

engine.add(mlir_text, passes)

['csc_to_csc', 'csr_to_csr']

Let’s verify that converting to and from the same layout give correct results.

csc_result = engine.csc_to_csc(csc_matrix)

csc_result.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_result.toarray())

True

csr_result = engine.csr_to_csr(csr_matrix)

csr_result.toarray()

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

np.all(dense_matrix == csr_result.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. ]])

np.all(dense_matrix.T == csc_matrix_transpose.toarray())

True

However, when we’re transposing a CSR matrix and want to return a CSR matrix as well, there is “real” work that is done. This “real” work involves doing exactlly what graphblas.convert_layout does under the covers in addition to changing the metadata. The same goes for transposing a CSC matrix to a CSC matrix.

The example below shows how to transpose a CSC mmatrix to a CSC matrix.

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_csc_to_csc(%sparse_tensor: tensor<?x?xf64, #CSC64>) -> tensor<?x?xf64, #CSC64> {
    %answer = graphblas.transpose %sparse_tensor : tensor<?x?xf64, #CSC64> to tensor<?x?xf64, #CSC64>
    return %answer : tensor<?x?xf64, #CSC64>
}
"""

engine.add(mlir_text, passes)

['transpose_csc_to_csc']

csc_matrix = engine.transpose_csc_to_csc(csc_matrix_transpose)

csc_matrix.toarray()

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

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

True

previous

Python Utilities for MLIR’s Sparse Tensors

next

Vector Ops