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

Vector Ops¶

This tutorial will go over how to use the --graphblas-lower pass from graphblas-opt to lower some ops from the GraphBLAS dialect that deal with vectors into executable Python code. There are many such ops. We’ll focus on graphblas.reduce_to_scalar whhen applied to vectors.

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
import numpy as np
from mlir_graphblas.tools.utils import sparsify_array

engine = mlir_graphblas.MlirJitEngine()

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-openmp-to-llvm",
    "--convert-math-to-llvm",
    "--convert-std-to-llvm",
    "--reconcile-unrealized-casts",
]

Overview of graphblas.reduce_to_scalar¶

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

graphblas.reduce_to_scalar reduces a sparse tensor (CSR matrix, CSC matrix, or sparse vector) to a scalar according to the given aggregator. We’ll use the “argmin” operator for this example. If there are multiple values that can be the argmin, an arbitrary one is chosen from them.

Let’s create an example input sparse vector.

dense_vector = np.array(
    [0, 1, 0, -999, 2, 0, 3, 0],
    dtype=np.int64,
)

sparse_vector = sparsify_array(dense_vector, [True])

Here’s what some MLIR code using graphblas.vector_argmin looks like.

mlir_text = """
#CV64 = #sparse_tensor.encoding<{
  dimLevelType = [ "compressed" ],
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

module {
    func @vector_argmin_wrapper(%vec: tensor<?xi64, #CV64>) -> i64 {
        %answer = graphblas.reduce_to_scalar %vec { aggregator = "argmin" } : tensor<?xi64, #CV64> to i64
        return %answer : i64
    }
}
"""

Let’s compile it and demonstrate its use.

engine.add(mlir_text, passes)

['vector_argmin_wrapper']

argmin = engine.vector_argmin_wrapper(sparse_vector)
argmin

Let’s verify that the behavior is the same as that of NumPy.

argmin == np.argmin(dense_vector)

True

Overview of graphblas.equal¶

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

graphblas.equal performs an equality check.

Here’s what some MLIR code using graphblas.equal looks like when used with vectors.

mlir_text = """
#CV64 = #sparse_tensor.encoding<{
  dimLevelType = [ "compressed" ],
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

module {
    func @vector_eq(%u: tensor<?xi64, #CV64>, %v: tensor<?xi64, #CV64>) -> i1 {
        %answer = graphblas.equal %u, %v : tensor<?xi64, #CV64>, tensor<?xi64, #CV64>
        return %answer : i1
    }
}
"""

Let’s compile it and demonstrate its use.

engine.add(mlir_text, passes)

['vector_eq']

Let’s create some example inputs.

a_dense = np.array([0, 1, 0], dtype=np.int64)
b_dense = np.array([9, 9, 9], dtype=np.int64)
c_dense = np.array([0, 1, 0], dtype=np.int64)

a = sparsify_array(a_dense, [True])
b = sparsify_array(b_dense, [True])
c = sparsify_array(c_dense, [True])

engine.vector_eq(a, b)

False

engine.vector_eq(a, c)

True

Let’s verify that the behavior is the same as that of NumPy.

Working with Sparse Layouts in the GraphBLAS Dialect

Matrix Ops