File size: 1,556 Bytes
0c6aadc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """Load metakernel natively on this machine (Windows/SAURON included).
The published Hugging Face kernel carries `*-linux` variants only, so on
Windows `get_kernel` has no matching build. The source is right here:
JIT-build it with `torch.utils.cpp_extension` and expose the identical
API.
import load_local
mk = load_local.load()
"""
import os
import sys
import types
_cached = None
_OPS = ["mk_triad", "mk_triad_passes", "mk_read", "mk_write", "mk_gather",
"mk_fma_f64", "mk_fma_dep", "mk_chase_global", "mk_chase_shared",
"mk_fma_f32", "mk_mma", "mk_atomics", "mk_spin", "mk_empty",
"mk_barrier", "mk_occupancy"]
def load(verbose=False):
"""Return the metakernel module, JIT-built from the local source."""
global _cached
if _cached is not None:
return _cached
from torch.utils.cpp_extension import load as _jit
root = os.path.dirname(os.path.abspath(__file__))
cuda_dir = os.path.join(root, "metakernel_cuda")
ext = _jit(name="metakernel_ext",
sources=[os.path.join(root, "local_bind.cpp"),
os.path.join(cuda_dir, "probes.cu")],
extra_include_paths=[cuda_dir],
verbose=verbose)
ops_mod = types.ModuleType("metakernel._ops")
ops_mod.ops = type("_Ops", (), {
name: staticmethod(getattr(ext, name)) for name in _OPS
})()
sys.modules["metakernel._ops"] = ops_mod
sys.path.insert(0, os.path.join(root, "torch-ext"))
import metakernel
_cached = metakernel
return metakernel
|