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

Lowering graphblas.apply to Generic Form

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

graphblas.apply ops specify behavior via the apply_operator attribute. The --graphblas-structuralize pass will lower those ops into equivalent graphblas.apply_generic ops with blocks specifying the behavior indicated by those apply_operator attribute values.

Let’s first import some necessary libraries.


import tempfile
from mlir_graphblas.cli import GRAPHBLAS_OPT_EXE

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 additive inverse operator.


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

module {

    func @apply_vector_ainv(%sparse_tensor: tensor<3xi64, #CV64>) -> tensor<3xi64, #CV64> {
        %answer = graphblas.apply %sparse_tensor { apply_operator = "ainv" } : (tensor<3xi64, #CV64>) to tensor<3xi64, #CV64>
        return %answer : tensor<3xi64, #CV64>
    }


}
"""

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)
#CV64 = #sparse_tensor.encoding<{
    dimLevelType = [ "compressed" ],
    pointerBitWidth = 64,
    indexBitWidth = 64
}>

builtin.module  {
  builtin.func @apply_vector_ainv(%arg0: tensor<3xi64, #CV64>) -> tensor<3xi64, #CV64> {
    %c0_i64 = constant 0 : i64
    %0 = graphblas.apply_generic %arg0 : tensor<3xi64, #CV64> to tensor<3xi64, #CV64>  {
    ^bb0(%arg1: i64):  // no predecessors
      %1 = subi %c0_i64, %arg1 : i64
      graphblas.yield transform_out %1 : i64
    }
    return %0 : tensor<3xi64, #CV64>
  }
}


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

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