1"""Common utilities that are useful for all the benchmarks.""" 2import numpy as np 3 4from mlir import ir 5from mlir.dialects import arith 6from mlir.dialects import func 7from mlir.dialects import memref 8from mlir.dialects import scf 9from mlir.passmanager import PassManager 10 11 12def 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"sparsifier{{{opt}}}" 19 PassManager.parse(pipeline).run(mlir_module) 20 21 22def 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 40def 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 57def 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 67def 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 106