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])