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