def compile(self, compiler, support_lib: str): """Compile the ir.Module.""" assert self._module is not None, \ 'StressTest: must call build() before compile()' assert self._engine is None, \ 'StressTest: must not call compile() repeatedly' compiler(self._module) self._engine = ExecutionEngine( self._module, opt_level=0, shared_libs=[support_lib]) return self
def compiler(): with ir.Context(), ir.Location.unknown(): kernel_func = get_kernel_func_from_module(module) timer_func = emit_timer_func() wrapped_func = emit_benchmark_wrapped_main_func( kernel_func, timer_func) main_module_with_benchmark = ir.Module.parse( str(timer_func) + str(wrapped_func) + str(kernel_func)) setup_passes(main_module_with_benchmark) c_runner_utils = os.getenv("MLIR_C_RUNNER_UTILS", "") assert os.path.exists(c_runner_utils),\ f"{c_runner_utils} does not exist." \ f" Please pass a valid value for" \ f" MLIR_C_RUNNER_UTILS environment variable." runner_utils = os.getenv("MLIR_RUNNER_UTILS", "") assert os.path.exists(runner_utils),\ f"{runner_utils} does not exist." \ f" Please pass a valid value for MLIR_RUNNER_UTILS" \ f" environment variable." engine = ExecutionEngine( main_module_with_benchmark, 3, shared_libs=[c_runner_utils, runner_utils]) return engine.invoke
def __call__(self, module: ir.Module) -> ExecutionEngine: self._passmanager.run(module) return ExecutionEngine(module, opt_level=0, shared_libs=[self._support_lib])
class StressTest: def __init__(self, tyconv: TypeConverter): self._tyconv = tyconv self._roundtripTp = None self._module = None self._engine = None def _assertEqualsRoundtripTp(self, tp: ir.RankedTensorType): assert self._roundtripTp is not None, \ 'StressTest: uninitialized roundtrip type' if tp != self._roundtripTp: raise AssertionError( f"Type is not equal to the roundtrip type.\n" f"\tExpected: {self._roundtripTp}\n" f"\tFound: {tp}\n") def build(self, types: List[ir.Type]): """Builds the ir.Module. The module has only the @main function, which will convert the input through the list of types and then back to the initial type. The roundtrip type must be a dense tensor.""" assert self._module is None, 'StressTest: must not call build() repeatedly' self._module = ir.Module.create() with ir.InsertionPoint(self._module.body): tp0 = types.pop(0) self._roundtripTp = tp0 # TODO: assert dense? assert element type is recognised by the TypeConverter? types.append(tp0) funcTp = ir.FunctionType.get(inputs=[tp0], results=[tp0]) funcOp = builtin.FuncOp(name='main', type=funcTp) funcOp.attributes['llvm.emit_c_interface'] = ir.UnitAttr.get() with ir.InsertionPoint(funcOp.add_entry_block()): arg0 = funcOp.entry_block.arguments[0] self._assertEqualsRoundtripTp(arg0.type) v = st.ConvertOp(types.pop(0), arg0) for tp in types: w = st.ConvertOp(tp, v) # Release intermediate tensors before they fall out of scope. st.ReleaseOp(v.result) v = w self._assertEqualsRoundtripTp(v.result.type) std.ReturnOp(v) return self def writeTo(self, filename): """Write the ir.Module to the given file. If the file already exists, then raises an error. If the filename is None, then is a no-op.""" assert self._module is not None, \ 'StressTest: must call build() before writeTo()' if filename is None: # Silent no-op, for convenience. return self if os.path.exists(filename): raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), filename) with open(filename, 'w') as f: f.write(str(self._module)) return self def compile(self, compiler, support_lib: str): """Compile the ir.Module.""" assert self._module is not None, \ 'StressTest: must call build() before compile()' assert self._engine is None, \ 'StressTest: must not call compile() repeatedly' compiler(self._module) self._engine = ExecutionEngine( self._module, opt_level=0, shared_libs=[support_lib]) return self def run(self, np_arg0: np.ndarray) -> np.ndarray: """Runs the test on the given numpy array, and returns the resulting numpy array.""" assert self._engine is not None, \ 'StressTest: must call compile() before run()' self._assertEqualsRoundtripTp( self._tyconv.get_RankedTensorType_of_nparray(np_arg0)) np_out = np.zeros(np_arg0.shape, dtype=np_arg0.dtype) self._assertEqualsRoundtripTp( self._tyconv.get_RankedTensorType_of_nparray(np_out)) mem_arg0 = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_arg0))) mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_out))) self._engine.invoke('main', mem_out, mem_arg0) return rt.ranked_memref_to_numpy(mem_out[0])