mirror of https://github.com/llvm/torch-mlir
67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
# See https://llvm.org/LICENSE.txt for license information.
|
|
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
|
|
import os
|
|
import platform
|
|
|
|
_refjit = None
|
|
|
|
|
|
def get_refjit():
|
|
"""Dynamically resolves the refjit backend native module."""
|
|
global _refjit
|
|
if _refjit is not None:
|
|
return _refjit
|
|
from ...._mlir_libs import _npcomp as _cext
|
|
try:
|
|
imported_refjit = _cext.backend.refjit
|
|
except AttributeError:
|
|
raise ImportError(
|
|
"The npcomp native module was not compiled with refjit support")
|
|
_refjit = imported_refjit
|
|
return _refjit
|
|
|
|
|
|
def is_enabled() -> bool:
|
|
"""Returns whether the backend is enabled for the current build."""
|
|
try:
|
|
_get_refjit()
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def get_runtime_libs():
|
|
# The _refjit_resources directory is at the npcomp.compiler level.
|
|
resources_dir = os.path.join(os.path.dirname(__file__))
|
|
suffix = ".so"
|
|
if platform.system() == "Darwin":
|
|
suffix = ".dylib"
|
|
shlib_name = f"libNPCOMPCompilerRuntimeShlib{suffix}"
|
|
return [os.path.join(resources_dir, shlib_name)]
|
|
|
|
|
|
class JitModuleInvoker:
|
|
"""Wrapper around a native JitModule for calling functions."""
|
|
|
|
def __init__(self, jit_module):
|
|
super().__init__()
|
|
self._jit_module = jit_module
|
|
|
|
def __getattr__(self, function_name):
|
|
return self.__getitem__(function_name)
|
|
|
|
def __getitem__(self, function_name):
|
|
|
|
def invoke(*args):
|
|
results = self._jit_module.invoke(function_name, args)
|
|
if len(results) == 1:
|
|
# De-tuple.
|
|
return results[0]
|
|
else:
|
|
return tuple(results)
|
|
|
|
invoke.__isnpcomp__ = True
|
|
return invoke
|