Files
hao-ai-lab-FastVideo/fastvideo-kernel/wan22_vae_decode/tools/loader.py
T
William Lin cf3ea82e74 [perf] PAN2 sampling: Wan2.2 VAE decode kernel, split decode, fused DiT kernels
- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on
  first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames
  per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output.
- WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32
  accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it);
  decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the
  video, with the ranks agreeing on whether every one of them could build the kernel.
- PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions,
  set_fused_kernels).
- PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels).
- VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame.
- The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.
2026-10-28 19:43:04 +00:00

543 lines
27 KiB
Python

"""Builds and loads a solution exactly as it is deployed, so that what runs here runs there.
A solution is a directory holding exactly the files that are deployed, plus `solution.json`, the deployment's build
spec, with its field names and meaning:
{"languages": ["cuda_cpp"], "entry_point": "binding.cpp::run", "destination_passing_style": true,
"compile_options": {"cflags": [...], "cuda_cflags": [...], "ld_flags": [...]}, "dependencies": []}
- `languages`: always ["cuda_cpp"].
- `entry_point`: "<file>::<function>", the function the deployment calls, in a C++/CUDA file.
- `destination_passing_style`: required, and true for this task: run() writes its output into `out`, its last
argument, and returns nothing. The calling convention must never be left to inference.
- `compile_options` (optional): extra flags; see the build below. `dependencies` (optional): a list of names, which
does not change the build. `binding` (optional): null or "torch".
Any other field is refused: the deployment would ignore it, so a misspelled field would silently fall back to its
default. Every other file in the directory is a source: non-empty UTF-8 text at a relative path, a plain file (no
links; `__pycache__` and `.pyc` files are ignored).
The module exports the entry, called positionally as the deployment calls it (argument names play no part):
run(z, weights, rank, world, pg, out) every call; writes out on rank 0
prepare(weights, rank, world, pg) -> list[Tensor] optional, under that name: once per process, before run
(spec.py says what each argument is.) Load one solution per process.
stage(solution_dir) writes the sources into a fresh working directory and compiles every .cpp/.cc/.cxx/.c/.cu file at
its top into one extension module, in a separate process, with torch.utils.cpp_extension.load, build_directory = the
working directory, include paths [working directory, $CUTLASS_DIR/include, $CUTLASS_DIR/tools/util/include] ($CUTLASS_DIR
defaults to fastvideo-kernel/include/cutlass, the CUTLASS submodule), and these
flags on top of torch's own:
- no compile_options: cuda_cflags [], cflags [], ld_flags ["-lcuda"];
- compile_options given: its lists, where a missing one defaults to cflags [], cuda_cflags ["-O3",
"--use_fast_math"], ld_flags ["-lcuda"];
- unless a cuda_cflag names an architecture (-gencode, -arch), a -gencode flag for the GPU's own architecture with
its arch-specific features goes first (sm_100a on compute capability 10.0), as the deployment does for its GPU.
Files in subdirectories are not compiled, but can be included. import_built(working directory) then makes it the
current directory, as the deployment does, imports the module and takes run (and prepare) from it; load() does both.
The module's name is the build's, not the solution's: PYBIND11_MODULE must be given TORCH_EXTENSION_NAME, as the
deployment builds the module under a name of its own. bench.py stages a build once and its ranks import that build.
Development builds are kept on each node, so that a job does not compile again what an earlier job on its node built:
a content-addressed cache in a directory of the node that setup/node.sh makes for the agent user (CACHE_DIR, or
$GUILD_BUILD_CACHE), keyed by everything that determines the build: the sources, solution.json, the flags, the target
architecture, this loader, and the toolchain (the machine, Python, torch with its CUDA, nvcc, the host compiler, and
the variables they read). A build found there is copied into the working directory in place of compiling it; one made
is put there with its sources and module, without what the compiler left beside them. Only the commands of a
development copy of the harness use it (tools/ in a repo: `loader.py check` and `build`, `bench.py`), unless they are
given --no-cache. The copy in the task's checks/, which CI, the hooks and the task's commands run, never reads or
writes it and builds fresh: the jobs of agents write that directory, so what it holds is theirs to forge. Past a
quarter of its filesystem, a limit that only stops a runaway, the least recently used builds are removed.
python tools/loader.py check [solution] [--build-dir DIR] build and load it, print what was loaded
python tools/loader.py build [solution] [DIR] only write and build it in DIR (default build/): with
--arch (e.g. sm_100a), on a machine without a GPU,
such as a CPU builder, where the module is not loaded
"""
import argparse
import atexit
import functools
import hashlib
import importlib.util
import json
import os
import platform
import re
import shutil
import stat
import subprocess
import sys
import tempfile
import time
SPEC_FILE = "solution.json"
# The name the build gives the module; solutions never name it themselves (see the module docstring). It is this
# kernel's own: another kernel's loader may build an extension of a generic name (solution_ext), and in one process the
# second extension of a name loads as the first one.
MODULE = "wan22_vae_decode_ext"
# The optional function a solution exports beside its entry, called once per process before any call of the entry.
PREPARE = "prepare"
BINDINGS = ("torch",)
CPP_ENTRY_SUFFIXES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh")
# The fields solution.json may hold.
SPEC_FIELDS = ("languages", "entry_point", "dependencies", "destination_passing_style", "binding", "compile_options")
OPTION_FIELDS = ("cflags", "cuda_cflags", "ld_flags")
# The languages a solution may declare, and the files compiled.
LANGUAGES = ("cuda_cpp",)
CPP_SOURCES = (".cu", ".cpp", ".cc", ".cxx", ".c")
DEFAULT_OPTIONS = {"cflags": [], "cuda_cflags": ["-O3", "--use_fast_math"], "ld_flags": ["-lcuda"]}
BUILD_TIMEOUT_SEC = 1800
# The first argument of a PYBIND11_MODULE in a source.
MODULE_NAME = re.compile(r"\bPYBIND11_MODULE\s*\(\s*([A-Za-z_]\w*)")
HERE = os.path.dirname(os.path.abspath(__file__))
# The node's cache of development builds (see the module docstring), and what each build there holds beside its sources.
CACHE_DIR = "/scratch/kdev-builds"
STAMP, SIZE = ".stamp", ".size"
# The share of its filesystem past which the cache removes its least recently used builds: a limit that stops a runaway.
CACHE_SHARE = 0.25
# What a job that stopped while it stored a build left in the cache is removed after this long.
STALE_SEC = 3600
# The variables the compilers read, which change what they make.
BUILD_VARIABLES = ("CC", "CXX", "CUDA_HOME", "CUDA_PATH", "CUTLASS_DIR", "NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS",
"CPATH", "C_INCLUDE_PATH", "CPLUS_INCLUDE_PATH", "LIBRARY_PATH")
class SolutionError(Exception):
"""The solution is malformed, or fails to build or load."""
def _nonempty_str(x):
return isinstance(x, str) and len(x) >= 1
def _str_list(x):
return isinstance(x, list) and all(isinstance(s, str) for s in x)
def spec_problems(spec):
"""Problems of a build spec, as a list of strings (empty: valid); see the module docstring.
Stricter than the deployment only where it would coerce a value (a "true" string as a boolean) or fall back to a
default (an unknown field, a missing destination_passing_style)."""
if not isinstance(spec, dict):
return [f"{SPEC_FILE} must hold an object"]
out = []
unknown = sorted(set(spec) - set(SPEC_FIELDS))
if unknown:
out.append(f"unknown fields {unknown} (known: {list(SPEC_FIELDS)})")
langs = spec.get("languages")
if not isinstance(langs, list) or not langs or any(lang not in LANGUAGES for lang in langs):
out.append(f"languages: required, {list(LANGUAGES)}")
entry = spec.get("entry_point")
if not _nonempty_str(entry):
out.append("entry_point: required, a non-empty string")
elif entry.count("::") != 1:
out.append(f'entry_point: {entry!r} is not "<file_path>::<function_name>"')
elif os.path.splitext(entry.split("::")[0])[1] not in CPP_ENTRY_SUFFIXES:
out.append(f"entry_point: needs a C++/CUDA entry file, not {entry.split('::')[0]!r}")
deps = spec.get("dependencies", [])
if not isinstance(deps, list) or not all(_nonempty_str(d) for d in deps):
out.append("dependencies: a list of non-empty strings")
if "destination_passing_style" not in spec:
out.append("destination_passing_style: must be given, and true here (run writes out, its last argument)")
elif spec["destination_passing_style"] is not True:
out.append("destination_passing_style: true: this task's run() writes the video into out, its last argument, "
"and returns nothing")
if spec.get("binding") not in (None,) + BINDINGS:
out.append(f"binding: null or one of {list(BINDINGS)}")
options = spec.get("compile_options")
if options is not None:
if not isinstance(options, dict):
out.append("compile_options: null or an object")
else:
for k in OPTION_FIELDS:
if k in options and not _str_list(options[k]):
out.append(f"compile_options.{k}: a list of strings")
return out
def read_spec(solution_dir):
"""solution.json, checked against the build spec's rules (see the module docstring)."""
path = os.path.join(solution_dir, SPEC_FILE)
try:
with open(path) as f:
spec = json.load(f)
except (OSError, ValueError) as e:
raise SolutionError(f"cannot read {path}: {e}") from e
problems = spec_problems(spec)
if problems:
raise SolutionError(f"{SPEC_FILE}: " + "; ".join(problems))
return spec
def sources(solution_dir):
"""{relative path: text} of every file of the solution except solution.json."""
out = {}
for root, dirs, files in os.walk(solution_dir):
dirs[:] = sorted(d for d in dirs if d != "__pycache__")
for name in sorted(files):
path = os.path.join(root, name)
rel = os.path.relpath(path, solution_dir).replace(os.sep, "/")
if rel == SPEC_FILE or name.endswith(".pyc"):
continue
if os.path.islink(path):
raise SolutionError(f"{rel}: sources must be plain files, not links")
try:
with open(path, encoding="utf-8") as f:
text = f.read()
except (OSError, UnicodeDecodeError) as e:
raise SolutionError(f"{rel}: sources must be UTF-8 text ({e})") from e
if not text:
raise SolutionError(f"{rel}: sources must not be empty")
out[rel] = text
if not out:
raise SolutionError(f"{solution_dir} holds no sources")
named = sorted({m for text in out.values() for m in MODULE_NAME.findall(text) if m != "TORCH_EXTENSION_NAME"})
if named:
raise SolutionError(f"PYBIND11_MODULE names its module {', '.join(named)}: give it TORCH_EXTENSION_NAME, the "
f"name the build chooses (the deployment builds the module under a name of its own)")
return out
def read(solution_dir):
"""(spec, {relative path: text}) of a solution, checked as stage() checks it before building: solution.json against
the build spec's rules, the sources as the deployment takes them, the entry file among them."""
spec = read_spec(solution_dir)
files = sources(solution_dir)
entry = spec["entry_point"].split("::")[0]
if entry not in files:
raise SolutionError(f"entry file {entry} is not among the sources")
return spec, files
def gpu_arch():
"""The architecture solutions are built for, e.g. sm_100a: the GPU's own, with its arch-specific features. It asks
the GPU's properties only, which makes no CUDA context."""
import torch
if not torch.cuda.is_available():
raise SolutionError("there is no GPU here to take the architecture from: name it (--arch sm_100a, or the "
"architecture of the GPUs the build is for)")
major, minor = torch.cuda.get_device_capability()
return f"sm_{major}{minor}" + ("a" if major >= 9 else "")
def compile_flags(spec, arch):
"""(cflags, cuda_cflags, ld_flags) added to torch's own flags (see the module docstring)."""
given = spec.get("compile_options")
options = {"cflags": [], "cuda_cflags": [], "ld_flags": ["-lcuda"]} if given is None else {
k: list(given.get(k, v)) for k, v in DEFAULT_OPTIONS.items()}
if not any("-gencode" in f or "-arch" in f for f in options["cuda_cflags"]):
options["cuda_cflags"].insert(0, f"-gencode=arch=compute_{arch[3:]},code={arch}")
return options["cflags"], options["cuda_cflags"], options["ld_flags"]
def cache_dir():
"""The node's cache of development builds, or None where a build is made fresh: in the copy of the harness in the
task's checks/ (CI, hooks and the task's commands run it, which must not load what a job could have put there), and
on a node without a cache that this user may write."""
checks = os.environ.get("CHECKS")
if checks and os.path.realpath(checks) == os.path.realpath(HERE):
return None
cache = os.environ.get("GUILD_BUILD_CACHE", CACHE_DIR)
return cache if os.path.isdir(cache) and os.access(cache, os.W_OK | os.X_OK) else None
@functools.lru_cache(maxsize=None)
def toolchain():
"""What determines a build beside the solution, its flags and the architecture: this loader, the machine, Python,
torch with its CUDA, nvcc, the host compiler, and the variables the compilers read."""
import torch
from torch.utils import cpp_extension
def version(cmd):
try:
p = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
return (p.stdout + p.stderr).strip()
except (OSError, subprocess.TimeoutExpired) as e:
return f"unavailable: {e}"
with open(os.path.abspath(__file__), "rb") as f:
loader = hashlib.sha256(f.read()).hexdigest()
return {"loader": loader, "machine": platform.machine(), "python": sys.version,
"torch": [torch.__version__, torch.version.cuda, os.path.dirname(torch.__file__)],
"nvcc": version([os.path.join(cpp_extension.CUDA_HOME or "/usr/local/cuda", "bin", "nvcc"), "--version"]),
"cxx": version([os.environ.get("CXX", "c++"), "--version"]),
"variables": {k: os.environ[k] for k in BUILD_VARIABLES if k in os.environ}}
def cache_key(spec, files, flags, arch, tools):
"""The key of a build in the cache: a hash of everything that determines it (see the module docstring)."""
what = {"spec": spec, "files": sorted(files.items()), "flags": flags, "arch": arch, "toolchain": tools}
return hashlib.sha256(json.dumps(what, sort_keys=True).encode()).hexdigest()
def disk_usage(path):
"""The bytes a directory takes on disk."""
total = 0
for root, _, names in os.walk(path):
for name in names:
try:
st = os.lstat(os.path.join(root, name))
except OSError:
continue
total += getattr(st, "st_blocks", 0) * 512 or st.st_size
return total
def fetch(cache, key, build_dir):
"""Copies the build of a key from the cache into build_dir, in its place, and marks the build used; False when the
cache has none, or it went while it was copied."""
entry = os.path.join(cache, key)
if not os.path.exists(os.path.join(entry, MODULE + ".so")):
return False
shutil.rmtree(build_dir, ignore_errors=True)
try:
shutil.copytree(entry, build_dir, symlinks=True, ignore=lambda d, _: [SIZE] if d == entry else [])
os.utime(entry)
except OSError:
shutil.rmtree(build_dir, ignore_errors=True)
return False
return True
def store(cache, key, build_dir, files):
"""Puts the build in build_dir into the cache under its key: its sources, module and stamp, without what the
compiler left beside them. A build of the key that another job put there first stays. Then the cache is kept within
its limit. A cache that cannot take it (a full disk, say) is left as it is: the build stands without it."""
entry = os.path.join(cache, key)
if os.path.exists(entry):
return
try:
tmp = tempfile.mkdtemp(prefix=f".{key[:16]}-", dir=cache)
except OSError:
return
try:
for rel in [*files, MODULE + ".so", STAMP]:
dest = os.path.join(tmp, rel)
os.makedirs(os.path.dirname(dest), exist_ok=True)
shutil.copy2(os.path.join(build_dir, rel), dest)
with open(os.path.join(tmp, SIZE), "w") as f:
f.write(str(disk_usage(tmp)))
os.rename(tmp, entry)
except OSError:
shutil.rmtree(tmp, ignore_errors=True)
return
try:
st = os.statvfs(cache)
evict(cache, CACHE_SHARE * st.f_blocks * st.f_frsize)
except OSError:
pass
def evict(cache, limit):
"""Removes the least recently used builds of the cache until it holds at most `limit` bytes, and what jobs that
stopped while they stored a build left there (after STALE_SEC). Anything but a directory is left alone."""
now = time.time()
builds = []
for name in os.listdir(cache):
path = os.path.join(cache, name)
try:
st = os.lstat(path)
except OSError:
continue
if not stat.S_ISDIR(st.st_mode):
continue
if name.startswith("."):
if now - st.st_mtime > STALE_SEC:
shutil.rmtree(path, ignore_errors=True)
continue
try:
with open(os.path.join(path, SIZE)) as f:
size = int(f.read())
except (OSError, ValueError):
size = disk_usage(path)
builds.append((st.st_mtime, path, size))
total = sum(size for _, _, size in builds)
for _, path, size in sorted(builds):
if total <= limit:
break
shutil.rmtree(path, ignore_errors=True)
total -= size
def stage(solution_dir, build_dir=None, arch=None, extra_cuda_cflags=(), cache=None):
"""Writes the sources into build_dir (a fresh temporary directory if None) and builds them there.
Returns (workdir, spec, build seconds or None when an earlier build was reused). A build_dir that already holds
a build of the same sources and flags is reused, so CI builds each solution once for its rounds. With a cache
(cache_dir(): a development build), a build of the same key is copied from it, and one made is put there."""
spec, files = read(solution_dir)
arch = arch or gpu_arch()
flags = compile_flags(spec, arch)
# Flags a tool adds that do not change the code, such as -lineinfo for source-level profiles.
flags[1].extend(extra_cuda_cflags)
digest = hashlib.sha256(json.dumps([spec, flags, sorted(files.items())]).encode()).hexdigest()
if build_dir is None:
build_dir = tempfile.mkdtemp(prefix="solution-")
atexit.register(shutil.rmtree, build_dir, True)
build_dir = os.path.abspath(build_dir)
stamp = os.path.join(build_dir, STAMP)
try:
with open(stamp) as f:
if f.read() == digest and os.path.exists(os.path.join(build_dir, MODULE + ".so")):
return build_dir, spec, None
except OSError:
pass
key = cache and cache_key(spec, files, flags, arch, toolchain())
if key and fetch(cache, key, build_dir):
return build_dir, spec, None
shutil.rmtree(build_dir, ignore_errors=True)
os.makedirs(build_dir)
for rel, text in files.items():
dest = os.path.join(build_dir, rel)
os.makedirs(os.path.dirname(dest), exist_ok=True)
with open(dest, "w", encoding="utf-8") as f:
f.write(text)
start = time.monotonic()
env = dict(os.environ, PYTORCH_ALLOC_CONF="expandable_segments:True")
cmd = [sys.executable, "-s", os.path.abspath(__file__), "_build", json.dumps(flags)]
try:
p = subprocess.run(cmd, cwd=build_dir, env=env, capture_output=True, text=True, timeout=BUILD_TIMEOUT_SEC)
except subprocess.TimeoutExpired as e:
raise SolutionError(f"the build timed out after {BUILD_TIMEOUT_SEC} s") from e
seconds = time.monotonic() - start
if p.returncode != 0 or not os.path.exists(os.path.join(build_dir, MODULE + ".so")):
log = (p.stdout + p.stderr).strip()
raise SolutionError(f"the build failed after {seconds:.0f} s:\n{_errors(log)}")
with open(stamp, "w") as f:
f.write(digest)
if key:
store(cache, key, build_dir, files)
return build_dir, spec, seconds
def _errors(log, limit=6000):
"""The compiler's error lines with some context, else the log's tail."""
lines = log.splitlines()
keep = [i for i, line in enumerate(lines) if "error" in line.lower()]
if not keep:
return log[-limit:]
picked = sorted({j for i in keep for j in range(max(0, i - 2), min(len(lines), i + 3))})
return "\n".join(lines[j] for j in picked)[-limit:]
# fastvideo-kernel's CUTLASS submodule (git submodule update --init fastvideo-kernel/include/cutlass).
_CUTLASS_SUBMODULE = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "include",
"cutlass")
def _build(flags):
"""The compile step, run in the working directory by stage() in its own process."""
import torch.utils.cpp_extension as ext
cflags, cuda_cflags, ld_flags = json.loads(flags)
here = os.getcwd()
srcs = sorted(os.path.join(here, p) for p in os.listdir(here)
if p.endswith(CPP_SOURCES) and os.path.isfile(os.path.join(here, p)))
if not srcs:
raise RuntimeError("No CUDA/C++ source files found in working directory")
cutlass = os.environ.get("CUTLASS_DIR") or _CUTLASS_SUBMODULE
try:
ext.load(name=MODULE, sources=srcs, extra_cuda_cflags=cuda_cflags, extra_cflags=cflags, extra_ldflags=ld_flags,
extra_include_paths=[here, f"{cutlass}/include", f"{cutlass}/tools/util/include"],
build_directory=here, verbose=True)
except ImportError as e:
# Built, and not loadable here: a machine without a GPU has no driver to load it against, and whoever loads it
# says whether it loads where it runs. Where a GPU is, a module that does not load (an undefined symbol, a
# missing library) is a build that failed, and is neither stamped nor kept in the cache.
import torch
if torch.cuda.is_available():
raise
print(f"built; not loaded on this machine: {e}", flush=True)
built = [p for p in os.listdir(here) if p.startswith(MODULE) and p.endswith(".so") and p != MODULE + ".so"]
if built:
os.rename(built[0], MODULE + ".so")
class Solution:
"""A loaded solution: fn, the entry (run), called as the deployment calls it, fn(*inputs, *outputs); prepare, the
module's prepare function or None; spec (solution.json), entry, languages, language, module, workdir, and
build_seconds (None when an earlier build was reused, or when the build was staged elsewhere)."""
def __init__(self, spec, fn, prepare, module, workdir, build_seconds):
self.spec = spec
self.languages = list(spec["languages"])
self.language = self.languages[0]
self.entry = spec["entry_point"]
self.fn, self.prepare, self.module = fn, prepare, module
self.workdir, self.build_seconds = workdir, build_seconds
def import_built(workdir, spec, chdir=True, build_seconds=None):
"""Imports the module that stage() built in workdir and takes its entry and prepare from it. chdir: make the
working directory the current one, as in deployment."""
import torch # noqa: F401 (the module needs libtorch loaded first)
if chdir:
os.chdir(workdir)
file, func = spec["entry_point"].split("::")
module_spec = importlib.util.spec_from_file_location(MODULE, os.path.join(workdir, MODULE + ".so"))
module = importlib.util.module_from_spec(module_spec)
try:
module_spec.loader.exec_module(module)
except Exception as e:
raise SolutionError(f"the built module does not load: {e}") from e
fn = getattr(module, func, None)
if not callable(fn):
raise SolutionError(f"{file} defines no function {func} in its PYBIND11_MODULE")
prepare = getattr(module, PREPARE, None)
if prepare is not None and not callable(prepare):
raise SolutionError(f"the module's {PREPARE} is not a function")
return Solution(spec, fn, prepare, module, workdir, build_seconds)
def load(solution_dir, build_dir=None, chdir=True, extra_cuda_cflags=(), cache=None):
"""Builds and loads a solution as it is deployed; see the module docstring.
build_dir: where to write and build it (a fresh temporary directory if None; reused if it holds the same build).
chdir: make the working directory the current one, as in deployment. extra_cuda_cflags: flags a tool adds to the
deployed ones, that must not change the code (e.g. -lineinfo). cache: the node's cache of development builds
(cache_dir()), or None to build fresh."""
workdir, spec, seconds = stage(solution_dir, build_dir, extra_cuda_cflags=extra_cuda_cflags, cache=cache)
return import_built(workdir, spec, chdir=chdir, build_seconds=seconds)
def main():
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
sub = p.add_subparsers(dest="cmd", required=True)
c = sub.add_parser("check", help="build and load a solution")
c.add_argument("solution", nargs="?", default="solution")
c.add_argument("--build-dir")
d = sub.add_parser("build", help="write and build a solution in a directory, without loading it")
d.add_argument("solution", nargs="?", default="solution")
d.add_argument("build_dir", nargs="?", default="build")
d.add_argument("--arch", help="the architecture to build for, e.g. sm_100a (default: this machine's GPU's)")
for q in (c, d):
q.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds")
b = sub.add_parser("_build")
b.add_argument("flags")
a = p.parse_args()
if a.cmd == "_build":
_build(a.flags)
return
cache = None if a.no_cache else cache_dir()
try:
if a.cmd == "build":
where, _, seconds = stage(os.path.abspath(a.solution), os.path.abspath(a.build_dir), a.arch, cache=cache)
module = os.path.join(where, MODULE + ".so")
print(f"built in {seconds:.1f} s: {module}" if seconds is not None else f"reused a build: {module}")
return
sol = load(os.path.abspath(a.solution), a.build_dir and os.path.abspath(a.build_dir), cache=cache)
except SolutionError as e:
sys.exit(f"FAILED: {e}")
built = f", built in {sol.build_seconds:.1f} s" if sol.build_seconds is not None else ", a build reused"
prepare = f", with {PREPARE}" if sol.prepare is not None else f", without {PREPARE}"
print(f"loaded {sol.entry} ({', '.join(sol.languages)}){prepare}{built}")
if __name__ == "__main__":
main()