metakernel v1: device dossier, throttle-rejected bench, ULP compare, fuzz, sweep, stamps
0c6aadc verified | """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 | |