def _load_compiled_function(dtype: str, device: str):
from tvm.module import load # this can be the small runtime python library, and doesn't need to be the whole thing
filename = DiagonaledMM._get_lib_filename(dtype, device)
current_dir = os.path.dirname(os.path.abspath(__file__))
potential_dirs = ['../../', '../', './', f'{current_dir}/', f'{current_dir}/../']
for potential_dir in potential_dirs:
filepath = '{}{}'.format(potential_dir, filename)
if os.path.isfile(filepath):
print('Loading tvm binary from: {}'.format(filepath))
return load(filepath)
return None
def _load_compiled_function(dtype: str, device: str):
filename = DiagonaledMM._get_lib_filename(dtype, device)
current_dir = os.path.dirname(os.path.abspath(__file__))
potential_dirs = [
'../../', '../', './', f'{current_dir}/', f'{current_dir}/../'
]
for potential_dir in potential_dirs:
filepath = '{}{}'.format(potential_dir, filename)
if os.path.isfile(filepath):
print('Loading tvm binary from: {}'.format(filepath))
return load(filepath)
return None