This page was generated from dialect/graphblas_dialect_tutorials/graphblas_structuralize/lower_reduce_to_scalar_rewrite.ipynb.

Lowering graphblas.reduce_to_scalar to Generic Form

This example will go over how to use the --graphblas-structuralize pass from graphblas-opt to lower graphblas.reduce_to_scalar ops to graphblas.reduce_to_scalar_generic ops.

graphblas.reduce_to_scalar ops specify behavior via an aggregator attributes. The --graphblas-structuralize pass will lower those ops into equivalent graphblas.reduce_to_scalar_generic ops with blocks specifying the behavior indicated by those semirings.

Let’s first import some necessary libraries.

import tempfile
from mlir_graphblas.cli import GRAPHBLAS_OPT_EXE

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

Since sparse tensor encodings can be very verbose in MLIR, let’s import some helpers to make the MLIR code more readable.

from mlir_graphblas.tools import tersify_mlir

Example

Here’s some example graphblas.apply code using the “plus” aggregator.

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

module {
    func @matrix_reduce_to_scalar_plus(%sparse_tensor: tensor<2x3xi64, #CSR64>) -> i64 {
        %answer = graphblas.reduce_to_scalar %sparse_tensor { aggregator = "plus" } : tensor<2x3xi64, #CSR64> to i64
        return %answer : i64
    }
}
"""

Let’s see what code we get when we run it through graphblas-opt with the --graphblas-structuralize pass.

with tempfile.NamedTemporaryFile() as temp:
    temp_file_name = temp.name
    with open(temp_file_name, 'w') as f:
        f.write(mlir_text)
    temp.flush()

    output_mlir = ! cat $temp_file_name | $GRAPHBLAS_OPT_EXE --graphblas-structuralize
    output_mlir = "\n".join(output_mlir)
    output_mlir = tersify_mlir(output_mlir)

print(output_mlir)

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

module {
  func @matrix_reduce_to_scalar_plus(%arg0: tensor<2x3xi64, #CSR64>) -> i64 {
    %c0_i64 = arith.constant 0 : i64
    %0 = graphblas.reduce_to_scalar_generic %arg0 : tensor<2x3xi64, #CSR64> to i64 {
      graphblas.yield agg_identity %c0_i64 : i64
    }, {
    ^bb0(%arg1: i64, %arg2: i64):
      %1 = arith.addi %arg1, %arg2 : i64
      graphblas.yield agg %1 : i64
    }
    return %0 : i64
  }
}

As shown above, --graphblas-structuralize expanded the “plus” aggregator into blocks performing that exact behavior.

We’ll leave exploring how the other aggregators expand as an exercise for the reader.

Lowering graphblas.apply to Generic Form GraphBLAS Optimizing Pass