| """Build fp64-emu from this checkout and return the package. |
| |
| The published variants are Linux x86_64. On Windows, or for a torch build |
| without a matching variant, `load_local.load()` JIT-compiles the same source |
| with `torch.utils.cpp_extension` and exposes the identical API: |
| |
| import load_local |
| emu = load_local.load() |
| C = emu.mm_nt(A, Bt) |
| |
| Requires the CUDA toolkit's `nvcc` and, on Windows, MSVC (run from a |
| Developer Command Prompt or after `vcvars64.bat`). |
| """ |
| import os |
| import sys |
| import types |
| from pathlib import Path |
|
|
| import torch |
| from torch.utils.cpp_extension import load as _load |
|
|
| ROOT = Path(__file__).resolve().parent |
| NAME = "fp64_emu_jit" |
|
|
|
|
| def load(verbose: bool = False) -> types.ModuleType: |
| """Compile the kernel, register its ops, and return the `fp64_emu` package.""" |
| win = os.name == "nt" |
| if "TORCH_CUDA_ARCH_LIST" not in os.environ and torch.cuda.is_available(): |
| cc = torch.cuda.get_device_capability(0) |
| os.environ["TORCH_CUDA_ARCH_LIST"] = f"{cc[0]}.{cc[1]}" |
| _load( |
| name=NAME, |
| sources=[str(ROOT / "torch-ext" / "torch_binding.cpp"), |
| str(ROOT / "fp64_emu_cuda" / "fp64_emu.cu")], |
| extra_include_paths=[str(ROOT / "torch-ext"), str(ROOT / "dev" / "include")], |
| extra_cflags=["-DCUDA_KERNEL", "/O2" if win else "-O3"], |
| extra_cuda_cflags=["-DCUDA_KERNEL", "-O3"], |
| extra_ldflags=["cublas.lib"] if win else ["-lcublas"], |
| is_python_module=False, |
| verbose=verbose, |
| ) |
|
|
| ops_mod = types.ModuleType("fp64_emu._ops") |
| ops_mod.ops = getattr(torch.ops, NAME) |
| sys.modules["fp64_emu._ops"] = ops_mod |
| pkg = types.ModuleType("fp64_emu") |
| pkg.__path__ = [str(ROOT / "torch-ext" / "fp64_emu")] |
| sys.modules["fp64_emu"] = pkg |
| src = (ROOT / "torch-ext" / "fp64_emu" / "__init__.py").read_text(encoding="utf-8") |
| exec(compile(src, str(ROOT / "torch-ext" / "fp64_emu" / "__init__.py"), "exec"), |
| pkg.__dict__) |
| return pkg |
|
|
