<jupyter_start><jupyter_text>Fusion of graphblas.matrix_select OpsThis example will go over how to use the `--graphblas-optimize` pass from `graphblas-opt` to fuse `graphblas.matrix_select` ops.When fusing `graphblas.matrix_select` ops, `--graphblas-optimize` simply combines several sequential `graphblas.matrix_select` ops into a single use of `graphblas.matrix_select` with multiple `selector` attributes.Let's first import some necessary libraries.<jupyter_code>import tempfile
from mlir_graphblas.cli import GRAPHBLAS_OPT_EXE<jupyter_output><empty_output><jupyter_text>Since [sparse tensor encodings](https://mlir.llvm.org/docs/Dialects/SparseTensorOps/sparsetensorencodingattr) can be very verbose in MLIR, let's import some helpers to make the MLIR code more readable.<jupyter_code>from mlir_graphblas.tools import tersify_mlir<jupyter_output><empty_output><jupyter_text>Fusing graphblas.matrix_select Ops With Same Source TensorIf we have several uses of `graphblas.matrix_select`, then `--graphblas-optimize` fuses them into one call with many selectors.Here's some example code using 2 sequential `graphblas.matrix_select` ops with 4 total selectors.<jupyter_code>mlir_text = """
#CSR64 = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (i,j)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

func @select_fuse_multi(%sparse_tensor: tensor<?x?xf64, #CSR64>) -> (tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>) {
    %thunk_a = constant 1.2 : f64
    %thunk_b = constant 3.4 : f64
    %answer1, %answer2 = graphblas.matrix_select %sparse_tensor, %thunk_a { selectors = ["gt", "triu"] } : tensor<?x?xf64, #CSR64>, f64 to tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
    %answer3 = graphblas.matrix_select %sparse_tensor, %thunk_b { selectors = ["tril", "gt"] } : tensor<?x?xf64, #CSR64>, f64 to tensor<?x?xf64, #CSR64>
    return %answer1, %answer2, %answer3 : tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
}
"""<jupyter_output><empty_output><jupyter_text>Let's see what code we get when we run it through `graphblas-opt` with the `--graphblas-optimize` pass.<jupyter_code>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-optimize
    output_mlir = "\n".join(output_mlir)
    output_mlir = tersify_mlir(output_mlir)

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

builtin.module  {
  builtin.func @select_fuse_multi(%arg0: tensor<?x?xf64, #CSR64>) -> (tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>) {
    %cst = constant 1.200000e+00 : f64
    %cst_0 = constant 3.400000e+00 : f64
    %0:3 = graphblas.matrix_select %arg0, %cst_0, %cst {selectors = ["tril", "gt", "gt", "triu"]} : tensor<?x?xf64, #CSR64>, f64, f64 to tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
    return %0#1, %0#2, %0#0 : tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
  }
}<jupyter_text>As shown above, `--graphblas-optimize` combined the original 2 uses of `graphblas.matrix_select` into one!  Fusing graphblas.matrix_select Ops With Different Source TensorsOur previous examples fused ops that all selected from the same source tensor. `--graphblas-optimize` can also fuse calls that use different source tensors as shown here.<jupyter_code>mlir_text = """
#CSR64 = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (i,j)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>

func @select_fuse_separate(%sparse_tensor1: tensor<?x?xf64, #CSR64>, %sparse_tensor2: tensor<?x?xf64, #CSR64>) -> (tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>) {
    %c0_f64 = constant 0.0 : f64
    %answer1 = graphblas.matrix_select %sparse_tensor1, %c0_f64 { selectors = ["gt"] } : tensor<?x?xf64, #CSR64>, f64 to tensor<?x?xf64, #CSR64>
    %answer2 = graphblas.matrix_select %sparse_tensor2 { selectors = ["triu"] } : tensor<?x?xf64, #CSR64> to tensor<?x?xf64, #CSR64>
    %answer3 = graphblas.matrix_select %sparse_tensor1 { selectors = ["tril"] } : tensor<?x?xf64, #CSR64> to tensor<?x?xf64, #CSR64>
    return %answer1, %answer2, %answer3 : tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
}
"""
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 --graphblas-optimize
    output_mlir = "\n".join(output_mlir)
    output_mlir = tersify_mlir(output_mlir)

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

builtin.module  {
  builtin.func @select_fuse_separate(%arg0: tensor<?x?xf64, #CSR64>, %arg1: tensor<?x?xf64, #CSR64>) -> (tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>) {
    %cst = constant 0.000000e+00 : f64
    %0 = graphblas.matrix_select %arg1 {selectors = ["triu"]} : tensor<?x?xf64, #CSR64> to tensor<?x?xf64, #CSR64>
    %1:2 = graphblas.matrix_select %arg0, %cst {selectors = ["tril", "gt"]} : tensor<?x?xf64, #CSR64>, f64 to tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
    return %1#1, %0, %1#0 : tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>, tensor<?x?xf64, #CSR64>
  }
}
{'fraction_non_alphanumeric': 0.15838189921151868, 'fraction_numerical': 0.058107644840589646, 'mean_word_length': 5.384026258205689, 'pattern_counts': {'":': 0, '<': 77, '<?xml version=': 0, '>': 85, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '525336', 'license': 'Apache-2.0', 'n_tokens_mistral': 2494, 'n_tokens_neox': 2298, 'n_words': 483, 'path': 'docs/dialect/graphblas_dialect_tutorials/graphblas_optimize/fuse_select.ipynb', 'repo_name': 'paul-tqh-nguyen/mlir-graphblas', 'size': '5834'}