torch-mlir/backend_test/iree/Sample/simple_invoke.py

38 lines
812 B
Python
Raw Normal View History

# RUN: %PYTHON %s
from npcomp.compiler.backend import iree
from npcomp.compiler.frontend import *
from npcomp.compiler import logging
from npcomp.compiler.target import *
# TODO: This should all exist in a high level API somewhere.
from _npcomp import mlir
logging.enable()
def compile_function(f):
fe = ImportFrontend(target_factory=GenericTarget32)
fe.import_global_function(f)
compiler = iree.CompilerBackend()
vm_blob = compiler.compile(fe.ir_module)
loaded_m = compiler.load(vm_blob)
return loaded_m[f.__name__]
@compile_function
def int_add(a: int, b: int):
return a + b
result = int_add(5, 6)
assert result == 11
@compile_function
def simple_control_flow(a: int, b: int):
return (a * b) and (a - b)
assert simple_control_flow(5, 6) == -1
assert simple_control_flow(-1, 0) == 0