mlir/benchmark/python/common.py

   1 """Common utilities that are useful for all the benchmarks."""
   2 import numpy as np
   3
   4 from mlir import ir
   5 from mlir.dialects import arith
   6 from mlir.dialects import func
   7 from mlir.dialects import memref
   8 from mlir.dialects import scf
   9 from mlir.passmanager import PassManager
  10
  11
  12 def setup_passes(mlir_module):
  13     """Setup pass pipeline parameters for benchmark functions."""
  14     opt = (
  15         "parallelization-strategy=none"
  16         " vectorization-strategy=none vl=1 enable-simd-index32=False"
  17     )
  18     pipeline = f"sparse-compiler{{{opt}}}"
  19     PassManager.parse(pipeline).run(mlir_module)
  20
  21
  22 def create_sparse_np_tensor(dimensions, number_of_elements):
  23     """Constructs a numpy tensor of dimensions `dimensions` that has only a
  24     specific number of nonzero elements, specified by the `number_of_elements`
  25     argument.
  26     """
  27     tensor = np.zeros(dimensions, np.float64)
  28     tensor_indices_list = [
  29         [np.random.randint(0, dimension) for dimension in dimensions]
  30         for _ in range(number_of_elements)
  31     ]
  32     for tensor_indices in tensor_indices_list:
  33         current_tensor = tensor
  34         for tensor_index in tensor_indices[:-1]:
  35             current_tensor = current_tensor[tensor_index]
  36         current_tensor[tensor_indices[-1]] = np.random.uniform(1, 100)
  37     return tensor
  38
  39
  40 def get_kernel_func_from_module(module: ir.Module) -> func.FuncOp:
  41     """Takes an mlir module object and extracts the function object out of it.
  42     This function only works for a module with one region, one block, and one
  43     operation.
  44     """
  45     assert (
  46         len(module.operation.regions) == 1
  47     ), "Expected kernel module to have only one region"
  48     assert (
  49         len(module.operation.regions[0].blocks) == 1
  50     ), "Expected kernel module to have only one block"
  51     assert (
  52         len(module.operation.regions[0].blocks[0].operations) == 1
  53     ), "Expected kernel module to have only one operation"
  54     return module.operation.regions[0].blocks[0].operations[0]
  55
  56
  57 def emit_timer_func() -> func.FuncOp:
  58     """Returns the declaration of nanoTime function. If nanoTime function is
  59     used, the `MLIR_RUNNER_UTILS` and `MLIR_C_RUNNER_UTILS` must be included.
  60     """
  61     i64_type = ir.IntegerType.get_signless(64)
  62     nanoTime = func.FuncOp("nanoTime", ([], [i64_type]), visibility="private")
  63     nanoTime.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
  64     return nanoTime
  65
  66
  67 def emit_benchmark_wrapped_main_func(kernel_func, timer_func):
  68     """Takes a function and a timer function, both represented as FuncOp
  69     objects, and returns a new function. This new function wraps the call to
  70     the original function between calls to the timer_func and this wrapping
  71     in turn is executed inside a loop. The loop is executed
  72     len(kernel_func.type.results) times. This function can be used to
  73     create a "time measuring" variant of a function.
  74     """
  75     i64_type = ir.IntegerType.get_signless(64)
  76     memref_of_i64_type = ir.MemRefType.get([-1], i64_type)
  77     wrapped_func = func.FuncOp(
  78         # Same signature and an extra buffer of indices to save timings.
  79         "main",
  80         (kernel_func.arguments.types + [memref_of_i64_type], kernel_func.type.results),
  81         visibility="public",
  82     )
  83     wrapped_func.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
  84
  85     num_results = len(kernel_func.type.results)
  86     with ir.InsertionPoint(wrapped_func.add_entry_block()):
  87         timer_buffer = wrapped_func.arguments[-1]
  88         zero = arith.ConstantOp.create_index(0)
  89         n_iterations = memref.DimOp(ir.IndexType.get(), timer_buffer, zero)
  90         one = arith.ConstantOp.create_index(1)
  91         iter_args = list(wrapped_func.arguments[-num_results - 1 : -1])
  92         loop = scf.ForOp(zero, n_iterations, one, iter_args)
  93         with ir.InsertionPoint(loop.body):
  94             start = func.CallOp(timer_func, [])
  95             call = func.CallOp(
  96                 kernel_func,
  97                 wrapped_func.arguments[: -num_results - 1] + loop.inner_iter_args,
  98             )
  99             end = func.CallOp(timer_func, [])
 100             time_taken = arith.SubIOp(end, start)
 101             memref.StoreOp(time_taken, timer_buffer, [loop.induction_variable])
 102             scf.YieldOp(list(call.results))
 103         func.ReturnOp(loop)
 104
 105     return wrapped_func