This page was generated from tools/engine/sparse_tensor_sum.ipynb.

JIT Engine: Sparse Tensor Summation

This example will go over how to compile MLIR code for taking the sum over a sparse tensor.

Previous tutorials covered how to handle functions on sparse tensors and how to take the sum of dense tensors. Accomplishing this task is mostly applying the knowledge from these tutorials. Thus, this will be more of a demonstration or example than it will be a tutorial.

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

import mlir_graphblas
import mlir_graphblas.sparse_utils
import numpy as np

engine = mlir_graphblas.MlirJitEngine()

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

This is the code we’ll use to take the sum over a sparse tensor.

mlir_text = """
#trait_sum_reduction = {
  indexing_maps = [
    affine_map<(i,j,k) -> (i,j,k)>,
    affine_map<(i,j,k) -> ()>
  ],
  sparse = [
    [ "S", "S", "S" ],
    [ ]
  ],
  iterator_types = ["reduction", "reduction", "reduction"],
  doc = "Sparse Tensor Sum"
}

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

func @sparse_sum(%argA: tensor<10x20x30xf32, #SparseEncoding>) -> f32 {
  %output_storage = linalg.init_tensor [] : tensor<f32>
  %reduction = linalg.generic #trait_sum_reduction
    ins(%argA: tensor<10x20x30xf32, #SparseEncoding>)
    outs(%output_storage: tensor<f32>) {
      ^bb(%a: f32, %x: f32):
        %0 = arith.addf %x, %a : f32
        linalg.yield %0 : f32
  } -> tensor<f32>
  %answer = tensor.extract %reduction[] : tensor<f32>
  return %answer : f32
}
"""

These are the passes we’ll use.

passes = [
    "--graphblas-structuralize",
    "--graphblas-optimize",
    "--graphblas-lower",
    "--sparsification",
    "--sparse-tensor-conversion",
    "--linalg-bufferize",
    "--func-bufferize",
    "--tensor-bufferize",
    "--finalizing-bufferize",
    "--convert-linalg-to-loops",
    "--convert-scf-to-cf",
    "--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"
]

Let’s generate our Python function.

engine.add(mlir_text, passes)
sparse_sum = engine.sparse_sum

Let’s generate our inputs.

indices = np.array([
    [0, 0, 0],
    [1, 1, 3],
    [2, 1, 6],
    [3, 9, 9],
    [4, 9, 12],
    [9, 9, 15],
    [9, 9, 18],
    [9, 15, 21],
    [9, 15, 24],
    [9, 19, 27],
], dtype=np.uint64)
values = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], dtype=np.float32)
sizes = np.array([10, 20, 30], dtype=np.uint64)
sparsity = np.array([True, True, True], dtype=np.bool8)

sparse_tensor = mlir_graphblas.sparse_utils.MLIRSparseTensor(indices, values, sizes, sparsity)

Let’s grab our result.

answer = sparse_sum(sparse_tensor)

Let’s see if our results match what we would expect.

answer == np.sum(values)

True

JIT Engine: Sparse Vector x Sparse Vector