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