This page was generated from tools/cli/apply_passes_to_string_or_file.ipynb.

Applying Passes to MLIR Code

This tutorial will go over how to apply passes to MLIR code as a string or stored in files on disk using the mlir-opt CLI Wrapper.

Let’s first import some necessary modules and generate an instance of our mlir-opt CLI Wrapper.

import os
import tempfile
from mlir_graphblas import MlirOptCli

cli = MlirOptCli(executable=None, options=None)

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

Applying Passes to MLIR Strings

Let’s say we had the following MLIR code in Python as a string. The function scale_func scales each element of the given tensor by the given scalar.

mlir_string = """
#trait_1d_scalar = {
  indexing_maps = [
    affine_map<(i) -> (i)>,  // A
    affine_map<(i) -> (i)>   // X (out)
  ],
  iterator_types = ["parallel"],
  doc = "X(i) = A(i) OP Scalar"
}
func @scale_func(%input: tensor<?xf32>, %scale: f32) -> tensor<?xf32> {
  %0 = linalg.generic #trait_1d_scalar
     ins(%input: tensor<?xf32>)
     outs(%input: tensor<?xf32>) {
      ^bb(%a: f32, %s: f32):
        %0 = arith.mulf %a, %scale  : f32
        linalg.yield %0 : f32
  } -> tensor<?xf32>
  return %0 : tensor<?xf32>
}
"""

Let’s say we wanted to run some passes over this MLIR code.

passes = [
    "--linalg-bufferize",
    "--func-bufferize",
    "--finalizing-bufferize",
    "--convert-linalg-to-affine-loops",
    "--lower-affine",
    "--convert-scf-to-cf",
]

We’ll need our string as a bytes literal.

mlir_bytes = mlir_string.encode()

We can now run passes over our MLIR code like so:

result = cli.apply_passes(mlir_bytes, passes)
print(result)

module {
  func @scale_func(%arg0: memref<?xf32>, %arg1: f32) -> memref<?xf32> {
    %c0 = arith.constant 0 : index
    %0 = memref.dim %arg0, %c0 : memref<?xf32>
    %1 = memref.alloc(%0) : memref<?xf32>
    %2 = memref.dim %arg0, %c0 : memref<?xf32>
    %c0_0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    cf.br ^bb1(%c0_0 : index)
  ^bb1(%3: index):  // 2 preds: ^bb0, ^bb2
    %4 = arith.cmpi slt, %3, %2 : index
    cf.cond_br %4, ^bb2, ^bb3
  ^bb2:  // pred: ^bb1
    %5 = memref.load %arg0[%3] : memref<?xf32>
    %6 = arith.mulf %5, %arg1 : f32
    memref.store %6, %1[%3] : memref<?xf32>
    %7 = arith.addi %3, %c1 : index
    cf.br ^bb1(%7 : index)
  ^bb3:  // pred: ^bb1
    return %1 : memref<?xf32>
  }
}

Applying Passes to MLIR Files on Disk

Let’s say that we have some MLIR file on our machine we want to run passes over. We’ll create a temporary file for our example here and use the same MLIR code and passes as above.

temporary_directory = tempfile.TemporaryDirectory()
temporary_filename = os.path.join(temporary_directory.name, "example.mlir")
with open(temporary_filename, 'w') as f:
    f.write(mlir_string)

We can run these passes over our file using the apply_passes method of MlirOptCli to get a string containing the lowered IR:

result = cli.apply_passes(temporary_filename, passes)
print(result)

module {
  func @scale_func(%arg0: memref<?xf32>, %arg1: f32) -> memref<?xf32> {
    %c0 = arith.constant 0 : index
    %0 = memref.dim %arg0, %c0 : memref<?xf32>
    %1 = memref.alloc(%0) : memref<?xf32>
    %2 = memref.dim %arg0, %c0 : memref<?xf32>
    %c0_0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    cf.br ^bb1(%c0_0 : index)
  ^bb1(%3: index):  // 2 preds: ^bb0, ^bb2
    %4 = arith.cmpi slt, %3, %2 : index
    cf.cond_br %4, ^bb2, ^bb3
  ^bb2:  // pred: ^bb1
    %5 = memref.load %arg0[%3] : memref<?xf32>
    %6 = arith.mulf %5, %arg1 : f32
    memref.store %6, %1[%3] : memref<?xf32>
    %7 = arith.addi %3, %c1 : index
    cf.br ^bb1(%7 : index)
  ^bb3:  // pred: ^bb1
    return %1 : memref<?xf32>
  }
}

Let’s make sure to clean up the temporary directory we created earlier.

temporary_directory.cleanup()

mlir-opt CLI Wrapper Using DebugResult