diff --git a/CMakeLists.txt b/CMakeLists.txt index a8ab048..83fabbc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,54 @@ endif() set(GGML_BUILD_TESTS OFF) set(GGML_BUILD_EXAMPLES OFF) +# tinyBLAS (llamafile) cache-blocked GEMM for the CPU backend. ggml ships it OFF +# by default, leaving mul_mat on the per-row ggml_vec_dot fallback; turning it on +# gives a cache-blocked F16/F32 GEMM that runs the transformer Linears ~1.6x +# faster (53%->32% of CPU time at S=8192). It is numerically a no-op (same f32 +# accumulation) and is parity-verified layer-by-layer against the f64 reference +# at the strict CPU-f32 gate. Set as the ggml *default* (not forced) so it covers +# every preset yet still yields to an explicit -DGGML_LLAMAFILE=OFF. +set(GGML_LLAMAFILE_DEFAULT ON) + +# Portable, fast CPU build by DEFAULT -- including a preset-less `cmake -B build`. +# Build every ggml CPU ISA variant (SSE..AVX-512) as runtime-dispatched dynamic +# backends and pick the best at load time, instead of a single -march=native +# target. Native is fragile across machines and, under Nix, NIX_ENFORCE_NO_NATIVE +# strips it -- silently leaving a ~13x slower SCALAR build. The presets used to be +# the only portable path; this makes the bare build "just work" too. +# +# Scope: only a plain CPU Release/RelWithDebInfo (or the typeless bare build). +# Skipped for GPU/fuzz/debug builds (kept simple and quick to compile), and each +# knob is guarded so an explicit -D or a preset cacheVariable still wins: +# - release-native (GGML_NATIVE=ON) keeps its single native target -- the +# `AND NOT GGML_NATIVE` guards avoid ggml's NATIVE-vs-BACKEND_DL conflict; +# - release / release-portable already set these ON (DEFINED -> left as-is). +# (GGML_BACKEND_DL's BUILD_SHARED_LIBS requirement is met by ggml's own +# default-ON shared-libs option on Linux.) +if ((NOT CMAKE_BUILD_TYPE OR CMAKE_BUILD_TYPE MATCHES "Release|RelWithDebInfo") + AND NOT FREE_SPLATTER_VULKAN AND NOT FREE_SPLATTER_CUDA AND NOT FREE_SPLATTER_FUZZ) + # CACHE (no FORCE): visible/overridable, and an explicit -D or preset value + # already in the cache wins. ggml's own option() then honors these. + set(GGML_NATIVE OFF CACHE BOOL "ggml: optimize for the build CPU (off => portable multi-variant build)") + if (NOT GGML_NATIVE) + set(GGML_BACKEND_DL ON CACHE BOOL "ggml: build backends as dynamic libs (enables runtime CPU ISA dispatch)") + set(GGML_CPU_ALL_VARIANTS ON CACHE BOOL "ggml: build every CPU ISA variant; best is picked at load time") + endif() +endif() + +# Co-locate our executables with ggml's dynamically-loaded backend .so files. A +# GGML_BACKEND_DL build emits libggml-*.so to /bin and loads them by path +# at runtime; ggml sets CMAKE_RUNTIME_OUTPUT_DIRECTORY for that, but only in its +# own (subdirectory) scope -- so without this our cli/bench land in the build +# root, away from the backends, and startup fails with "CPU backend init failed". +# The presets set this explicitly; an explicit value still wins (if-not-set). +# Gated on GGML_BACKEND_DL: only a dynamic-backend build needs the co-location, +# so static builds (vulkan/debug/native presets) keep their default binary +# location and existing script paths (e.g. bench.sh's build/vulkan/...) hold. +if (GGML_BACKEND_DL AND NOT CMAKE_RUNTIME_OUTPUT_DIRECTORY) + set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) +endif() + # --- Optional custom-op patch slot (LocalVQE pattern) ----------------------- # If pieces 1-3 ever need a ggml op not upstream, drop a patch in patches/ and it # is applied to the submodule working tree at configure time, idempotently. diff --git a/CMakePresets.json b/CMakePresets.json index f48d7f4..d2c977b 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -15,20 +15,30 @@ { "name": "release", "binaryDir": "${sourceDir}/build/release", - "cacheVariables": { "CMAKE_BUILD_TYPE": "Release" } - }, - { - "name": "release-portable", - "binaryDir": "${sourceDir}/build/release-portable", - "description": "Portable + fast CPU: build every ggml-cpu ISA variant and pick the best at runtime. Avoids -march=native (fragile, and stripped by Nix's NIX_ENFORCE_NO_NATIVE).", + "description": "Default release build. Portable + fast CPU: builds every ggml-cpu ISA variant (SSE..AVX-512) and selects the best at runtime. We deliberately do NOT default to -march=native: it is fragile across machines, and Nix's NIX_ENFORCE_NO_NATIVE strips it, leaving a SCALAR build ~13x slower. For a single-target native binary on non-Nix, use `release-native`.", "cacheVariables": { "CMAKE_BUILD_TYPE": "Release", "GGML_NATIVE": "OFF", "GGML_BACKEND_DL": "ON", "GGML_CPU_ALL_VARIANTS": "ON", + "CMAKE_RUNTIME_OUTPUT_DIRECTORY": "${sourceDir}/build/release/bin" + } + }, + { + "name": "release-portable", + "inherits": "release", + "binaryDir": "${sourceDir}/build/release-portable", + "description": "Alias of `release` (the portable default), kept for scripts/muscle-memory; identical flags in a separate build dir.", + "cacheVariables": { "CMAKE_RUNTIME_OUTPUT_DIRECTORY": "${sourceDir}/build/release-portable/bin" } }, + { + "name": "release-native", + "binaryDir": "${sourceDir}/build/release-native", + "description": "Single-target -march=native build (one self-contained binary, no runtime ISA dispatch). Fastest on the build machine, but NOT portable, and scalar under Nix (native is stripped) -- use only off-Nix.", + "cacheVariables": { "CMAKE_BUILD_TYPE": "Release", "GGML_NATIVE": "ON" } + }, { "name": "vulkan", "binaryDir": "${sourceDir}/build/vulkan", @@ -58,6 +68,7 @@ { "name": "debug", "configurePreset": "debug" }, { "name": "release", "configurePreset": "release" }, { "name": "release-portable", "configurePreset": "release-portable" }, + { "name": "release-native", "configurePreset": "release-native" }, { "name": "vulkan", "configurePreset": "vulkan" }, { "name": "profile", "configurePreset": "profile" }, { "name": "fuzz", "configurePreset": "fuzz" } diff --git a/README.md b/README.md index 2463327..3b57e21 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,23 @@ network on your CPU (or a Vulkan GPU). You give it N images; it gives you, for every pixel, a 3D Gaussian (position, colour, opacity, size, orientation) that a Gaussian-splatting viewer can render. +## Speed + +One scene — 2 photos at 512×512 — on a desktop (AMD Ryzen 9 7900, NVIDIA +RTX 5070 Ti), median time for a single forward pass: + +| run on | free-splatter.cpp | reference PyTorch¹ | +|---|--:|--:| +| **GPU** | **0.22 s** · Vulkan, f16 | 1.37 s · CUDA, fp16 | +| **CPU** | **14 s** · 12 threads, f16 | 54 s · fp32 | + +¹ the upstream FreeSplatter transformer in eager PyTorch — the same forward pass on +the same machine (reproduce with `scripts/bench.sh`). + +A GPU turns a handful of photos into a splat in well under a second; even the +no-GPU, no-Python CPU path beats the reference by ~4×. Run time scales with the +number of input views. + ## Build ```sh @@ -20,6 +37,12 @@ cmake --preset release cmake --build --preset release ``` +The `release` preset builds a **portable** CPU binary — every ggml CPU ISA +variant (up to AVX-512), with the best picked at runtime. (It avoids +`-march=native`, which Nix strips, leaving a ~13× slower scalar build; off-Nix +you can use `--preset release-native` for a single-target binary.) The +executables land in `build/release/bin/`. + No Nix? You just need CMake, a C++17 compiler, and the bundled `ggml` submodule (`git submodule update --init`). diff --git a/bench/free_splatter-bench.cpp b/bench/free_splatter-bench.cpp index 21d9eb8..8f3b165 100644 --- a/bench/free_splatter-bench.cpp +++ b/bench/free_splatter-bench.cpp @@ -1,36 +1,130 @@ -// free_splatter-bench: load a model and time end-to-end forward passes. -// (Forward is built incrementally across M1-M3; until then this reports load + -// geometry so the target links and the harness exists.) +// free_splatter-bench: time end-to-end forward passes of the engine. +// +// Loads a model, builds an N-view input (a deterministic synthetic batch, or a +// raw view-major NCHW .f32 file via --input), then runs free_splatter_run a few +// times and reports per-iteration latency (min/median/mean/max) and throughput. +// Load time is reported separately and excluded from the forward timings. +// +// free_splatter-bench [--device DEV] [--views N] [--input FILE.f32] +// [--iters K] [--warmup W] [--threads T] MODEL.gguf +// +// The last line is machine-parseable (scripts/bench.sh reads it): +// RESULT engine device=DEV views=N gc=23 load_ms=.. min_ms=.. median_ms=.. \ +// mean_ms=.. max_ms=.. views_per_s=.. #include "free_splatter.h" +#include +#include +#include #include +#include +#include #include +#include + +using clk = std::chrono::steady_clock; +static double ms_since(clk::time_point t0) { + return std::chrono::duration(clk::now() - t0).count(); +} + +static int usage(const char * a0) { + std::fprintf(stderr, + "usage: %s [--device DEV] [--views N] [--input FILE.f32]\n" + " [--iters K] [--warmup W] [--threads T] MODEL.gguf\n", a0); + return 2; +} int main(int argc, char ** argv) { - if (argc < 2) { - std::fprintf(stderr, "usage: %s MODEL.gguf [--device DEV]\n", argv[0]); - return 2; - } - const char * device = nullptr; - for (int i = 2; i < argc; i++) { - if (std::string(argv[i]) == "--device" && i + 1 < argc) device = argv[++i]; + const char * device = nullptr, * model = nullptr, * input = nullptr; + int views = 2, iters = 10, warmup = 2, threads = 0; + + for (int i = 1; i < argc; i++) { + std::string a = argv[i]; + if (a == "--device" && i+1 < argc) device = argv[++i]; + else if (a == "--views" && i+1 < argc) views = std::atoi(argv[++i]); + else if (a == "--input" && i+1 < argc) input = argv[++i]; + else if (a == "--iters" && i+1 < argc) iters = std::atoi(argv[++i]); + else if (a == "--warmup" && i+1 < argc) warmup = std::atoi(argv[++i]); + else if (a == "--threads" && i+1 < argc) threads = std::atoi(argv[++i]); + else if (a == "-h" || a == "--help") return usage(argv[0]); + else if (!model) model = argv[i]; + else return usage(argv[0]); } + if (!model || iters < 1 || views < 1) return usage(argv[0]); free_splatter_options * opts = free_splatter_options_new(); - if (device) free_splatter_options_set_device(opts, device); - free_splatter_ctx * ctx = free_splatter_load(argv[1], opts); - free_splatter_options_free(opts); + if (device) free_splatter_options_set_device(opts, device); + if (threads > 0) free_splatter_options_set_threads(opts, threads); - if (!ctx || free_splatter_last_error(ctx)) { - std::fprintf(stderr, "load failed: %s\n", - ctx ? free_splatter_last_error(ctx) : "oom"); - free_splatter_free(ctx); - return 1; + const auto t_load0 = clk::now(); + free_splatter_ctx * ctx = free_splatter_load(model, opts); + const double load_ms = ms_since(t_load0); + free_splatter_options_free(opts); + if (!ctx) { std::fprintf(stderr, "load: out of memory\n"); return 1; } + if (const char * err = free_splatter_last_error(ctx)) { + std::fprintf(stderr, "load failed: %s\n", err); free_splatter_free(ctx); return 1; } + free_splatter_geometry geo; free_splatter_geometry_of(ctx, &geo); - std::printf("loaded: %dx%d gaussian_channels=%d\n", - geo.image_width, geo.image_height, geo.gaussian_channels); + const int64_t per_view = (int64_t) geo.in_channels * geo.image_height * geo.image_width; + + // Build the input: a raw .f32 (its size fixes the view count), or a + // deterministic synthetic batch in [0,1]. Values don't affect timing; an + // LCG keeps it reproducible and free of denormals. + std::vector buf; + if (input) { + std::ifstream f(input, std::ios::binary | std::ios::ate); + if (!f) { std::fprintf(stderr, "cannot open %s\n", input); free_splatter_free(ctx); return 1; } + const std::streamsize bytes = f.tellg(); f.seekg(0); + buf.resize(bytes / sizeof(float)); f.read((char *) buf.data(), bytes); + if (per_view == 0 || buf.size() % per_view != 0) { + std::fprintf(stderr, "input size is not a whole number of views\n"); free_splatter_free(ctx); return 1; } + views = (int) (buf.size() / per_view); + } else { + buf.resize((size_t) views * per_view); + uint32_t s = 0x9e3779b9u; + for (float & v : buf) { s = s * 1664525u + 1013904223u; v = (s >> 8) * (1.0f / 16777216.0f); } + } + + std::printf("model: %s device=%s %dx%d in=%d gc=%d views=%d S=%lld load=%.1f ms\n", + model, device ? device : "cpu", geo.image_width, geo.image_height, + geo.in_channels, geo.gaussian_channels, views, + (long long) ((int64_t) views * (geo.image_height/8) * (geo.image_width/8)), load_ms); + + auto once = [&](double * out_ms) -> bool { + float * out = nullptr; size_t n_out = 0; + const auto t0 = clk::now(); + const int rc = free_splatter_run(ctx, buf.data(), (int32_t) views, + geo.image_height, geo.image_width, &out, &n_out); + if (out_ms) *out_ms = ms_since(t0); + if (rc != 0) { std::fprintf(stderr, "run failed: %s\n", free_splatter_last_error(ctx)); return false; } + free_splatter_buf_free(out); + return true; + }; + + for (int i = 0; i < warmup; i++) { if (!once(nullptr)) { free_splatter_free(ctx); return 1; } } + + std::vector t(iters); + for (int i = 0; i < iters; i++) { + if (!once(&t[i])) { free_splatter_free(ctx); return 1; } + std::printf(" iter %2d: %8.1f ms\n", i, t[i]); + } free_splatter_free(ctx); + + std::vector sorted = t; + std::sort(sorted.begin(), sorted.end()); + const double mn = sorted.front(), mx = sorted.back(); + const double med = sorted[iters / 2]; + double sum = 0; for (double v : t) sum += v; + const double mean = sum / iters; + const double vps = views / (med / 1000.0); + + std::printf("\nsummary: min=%.1f median=%.1f mean=%.1f max=%.1f ms " + "(%.2f views/s, %.2f scenes/s)\n", mn, med, mean, mx, vps, 1000.0 / med); + std::printf("RESULT engine device=%s views=%d gc=%d load_ms=%.1f min_ms=%.1f " + "median_ms=%.1f mean_ms=%.1f max_ms=%.1f views_per_s=%.3f\n", + device ? device : "cpu", views, geo.gaussian_channels, load_ms, + mn, med, mean, mx, vps); return 0; } diff --git a/scripts/bench.sh b/scripts/bench.sh new file mode 100755 index 0000000..bfff642 --- /dev/null +++ b/scripts/bench.sh @@ -0,0 +1,86 @@ +#!/usr/bin/env bash +# Benchmark free-splatter.cpp against the upstream FreeSplatter PyTorch model. +# +# Times one forward pass (N views -> the [N,512,512,23] transformer logits) and +# compares the ggml engine (CPU with AVX/AVX512, and Vulkan GPU) against PyTorch +# (CPU, and CUDA GPU), at matching precisions, then prints a table. +# +# nix develop -c scripts/bench.sh # default: 2 views +# VIEWS=4 ITERS=10 nix develop -c scripts/bench.sh +# +# Notes: +# - CPU uses the `release-portable` build (every ggml CPU ISA variant, best +# picked at runtime). The plain `release` preset is SCALAR under Nix +# (NIX_ENFORCE_NO_NATIVE strips -march=native) and ~13x slower -- do not use +# it for CPU timing. +# - PyTorch-on-CUDA needs a cu1xx torch build (.venv-torch-cu128) and the host +# NVIDIA driver libs (/run/opengl-driver/lib) on LD_LIBRARY_PATH; both are set +# up below. The CPU-only .venv-torch cannot use the GPU. +set -euo pipefail +cd "$(dirname "$0")/.." +ROOT=$(pwd) + +VIEWS=${VIEWS:-2} +ITERS=${ITERS:-8} +CPU_ITERS=${CPU_ITERS:-3} +GGUF_F16=${GGUF_F16:-$ROOT/.cache/freesplatter-scene-f16.gguf} +GGUF_F32=${GGUF_F32:-$ROOT/.cache/freesplatter-scene-f32.gguf} +CKPT=${CKPT:-$ROOT/.cache/freesplatter-scene.safetensors} +TORCH=${TORCH_PY:-$ROOT/.venv-torch-cu128/bin/python} + +# pip-CUDA torch finds libcuda.so.1 from the host driver here (NixOS). +export LD_LIBRARY_PATH="/run/opengl-driver/lib:${LD_LIBRARY_PATH:-}" + +CPU_BIN=$ROOT/build/release-portable/bin/free_splatter-bench +VK_BIN=$ROOT/build/vulkan/free_splatter-bench + +echo "building bench binaries (if needed) ..." +[ -x "$CPU_BIN" ] || { cmake --preset release-portable >/dev/null && \ + cmake --build --preset release-portable --target free_splatter-bench -j >/dev/null; } +[ -x "$VK_BIN" ] || { cmake --preset vulkan >/dev/null && \ + cmake --build build/vulkan --target free_splatter-bench -j >/dev/null; } + +RESULTS=$(mktemp) +trap 'rm -f "$RESULTS"' EXIT + +run() { # label, command... + local label=$1; shift + printf '>>> %s\n' "$label" >&2 + local line + line=$("$@" 2>/dev/null | grep '^RESULT' || true) + if [ -n "$line" ]; then echo "$label|$line" >> "$RESULTS" + else printf ' (no result -- skipped or failed)\n' >&2; fi +} + +# --- engine (ggml) --- +[ -f "$GGUF_F32" ] && run "engine | cpu | f32 " "$CPU_BIN" --device cpu --views "$VIEWS" --iters "$CPU_ITERS" --warmup 1 "$GGUF_F32" +[ -f "$GGUF_F16" ] && run "engine | cpu | f16 " "$CPU_BIN" --device cpu --views "$VIEWS" --iters "$CPU_ITERS" --warmup 1 "$GGUF_F16" +[ -f "$GGUF_F16" ] && run "engine | vulkan | f16 " "$VK_BIN" --device vulkan --views "$VIEWS" --iters "$ITERS" --warmup 2 "$GGUF_F16" +[ -f "$GGUF_F32" ] && run "engine | vulkan | f32 " "$VK_BIN" --device vulkan --views "$VIEWS" --iters "$ITERS" --warmup 2 "$GGUF_F32" + +# --- PyTorch reference --- +if [ -x "$TORCH" ] && [ -f "$CKPT" ]; then + run "torch | cpu | fp32" "$TORCH" scripts/bench_torch.py --ckpt "$CKPT" --device cpu --dtype fp32 --views "$VIEWS" --iters "$CPU_ITERS" --warmup 1 + if "$TORCH" -c 'import torch,sys; sys.exit(0 if torch.cuda.is_available() else 1)' 2>/dev/null; then + run "torch | cuda | fp16" "$TORCH" scripts/bench_torch.py --ckpt "$CKPT" --device cuda --dtype fp16 --views "$VIEWS" --iters "$ITERS" --warmup 3 + run "torch | cuda | fp32" "$TORCH" scripts/bench_torch.py --ckpt "$CKPT" --device cuda --dtype fp32 --views "$VIEWS" --iters "$ITERS" --warmup 3 + run "torch | cuda | bf16" "$TORCH" scripts/bench_torch.py --ckpt "$CKPT" --device cuda --dtype bf16 --views "$VIEWS" --iters "$ITERS" --warmup 3 + else + echo ">>> torch CUDA unavailable (need .venv-torch-cu128 + /run/opengl-driver/lib)" >&2 + fi +else + echo ">>> torch skipped (no $TORCH or $CKPT)" >&2 +fi + +# --- table --- +echo +printf '%-26s | %12s | %10s\n' "engine | device | dtype" "median (ms)" "scenes/s" +printf '%-26s-+-%12s-+-%10s\n' "$(printf '%.0s-' {1..26})" "------------" "----------" +sort "$RESULTS" | while IFS='|' read -r who dev dt rest; do + med=$(echo "$rest" | grep -o 'median_ms=[0-9.]*' | cut -d= -f2) + vps=$(echo "$rest" | grep -o 'views_per_s=[0-9.]*' | cut -d= -f2) + sps=$(awk -v v="$vps" -v n="$VIEWS" 'BEGIN{ if (v>0) printf "%.3f", v/n; else print "-" }') + printf '%-7s|%-8s|%-5s | %12s | %10s\n' "$who" "$dev" "$dt" "$med" "$sps" +done +echo +echo "(N=$VIEWS views, S=$((VIEWS*4096)) tokens; median of timed iters; lower ms / higher scenes/s is better)" diff --git a/scripts/bench_torch.py b/scripts/bench_torch.py new file mode 100644 index 0000000..2514440 --- /dev/null +++ b/scripts/bench_torch.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +"""Time the upstream FreeSplatter Transformer front-half in PyTorch. + +This is the *performance* reference (cf. hf_dump.py, which is the float64 +*numerical* reference and is deliberately slow). It runs the real transformer at +native precision with an efficient attention kernel (torch SDPA, the same +flash/mem-efficient path our engine's ggml_flash_attn_ext targets), with proper +warmup and CUDA synchronization, and reports per-forward latency so it can be +compared head-to-head with free_splatter-bench. + +What is timed matches the engine's forward: patchify -> +pos/view embed -> +24 transformer blocks -> final norm -> unpatchify (the [N,512,512,23] raw +logits). The engine additionally applies the SH residual + activations, a +sub-millisecond elementwise tail on top of this; it is not included here (the +upstream forward stops at the raw logits). + + nix develop -c bash -c 'export LD_LIBRARY_PATH=/run/opengl-driver/lib:$LD_LIBRARY_PATH; \ + .venv-torch-cu128/bin/python scripts/bench_torch.py \ + --ckpt .cache/freesplatter-scene.safetensors --device cuda --dtype fp16 --views 2' + +The last line is machine-parseable (scripts/bench.sh reads it): + RESULT torch device=cuda dtype=fp16 views=N ... median_ms=.. views_per_s=.. +""" +from __future__ import annotations + +import argparse +import statistics +import sys +import time +import types + +import torch + +OUTPUT_DIM = 23 + + +def install_sdpa_stub() -> None: + """Provide xformers.ops.memory_efficient_attention backed by torch SDPA. + + The upstream CrossAttention calls xops.memory_efficient_attention(q,k,v) with + q,k,v shaped ((b h), n, d); SDPA on that 3D layout computes exactly + softmax(q káµ€ / sqrt(d)) v per (batch*head), the fast flash/mem-efficient path. + Original FreeSplatter uses the real xformers flash kernel; SDPA's flash + backend is the equivalent. attention is ~57% of the FLOPs at S=8192, so the + kernel choice matters -- bench_torch's --sdpa pins it for a fair comparison. + """ + def mea(q, k, v, attn_bias=None, scale=None, **kw): + return torch.nn.functional.scaled_dot_product_attention(q, k, v, scale=scale) + xf = types.ModuleType("xformers") + ops = types.ModuleType("xformers.ops") + ops.memory_efficient_attention = mea + xf.ops = ops + sys.modules["xformers"], sys.modules["xformers.ops"] = xf, ops + + +def sdpa_context(which: str): + """Pin the SDPA backend (flash/mem/math) so attention kernel choice is + explicit rather than the heuristic default. Returns a context manager.""" + import contextlib + if which == "auto": + return contextlib.nullcontext() + from torch.nn.attention import SDPBackend, sdpa_kernel + backends = { + "flash": [SDPBackend.FLASH_ATTENTION], + "mem": [SDPBackend.EFFICIENT_ATTENTION], + "math": [SDPBackend.MATH], + }[which] + return sdpa_kernel(backends) + + +def load_model(ckpt: str, upstream_dir: str, dtype, device): + install_sdpa_stub() + sys.path.insert(0, upstream_dir) + from transformer import Transformer # type: ignore + from safetensors import safe_open + + model = Transformer(image_size=512, patch_size=8, input_dim=3, inner_dim=1024, + output_dim=OUTPUT_DIM, n_heads=16, depth=24) + sd = {} + with safe_open(ckpt, framework="pt") as f: + for k in f.keys(): + sd[k[len("transformer."):] if k.startswith("transformer.") else k] = f.get_tensor(k) + missing, unexpected = model.load_state_dict(sd, strict=False) + assert not missing and not unexpected, (missing[:3], unexpected[:3]) + return model.eval().to(device=device, dtype=dtype) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--ckpt", required=True) + ap.add_argument("--upstream", default=".cache/upstream") + ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"]) + ap.add_argument("--dtype", default=None, choices=["fp32", "fp16", "bf16"], + help="default: fp32 on cpu, fp16 on cuda") + ap.add_argument("--views", type=int, default=2) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument("--warmup", type=int, default=2) + ap.add_argument("--threads", type=int, default=0, help="cpu intra-op threads (0=torch default)") + ap.add_argument("--compile", action="store_true", help="wrap the model in torch.compile (first warmup absorbs compilation)") + ap.add_argument("--sdpa", default="auto", choices=["auto", "flash", "mem", "math"], + help="pin the SDPA attention backend (default: auto/heuristic)") + ap.add_argument("--seed", type=int, default=20260624) + args = ap.parse_args() + + if args.device == "cuda" and not torch.cuda.is_available(): + print("error: CUDA requested but torch.cuda.is_available() is False " + "(need a cu1xx torch build + /run/opengl-driver/lib on LD_LIBRARY_PATH)", file=sys.stderr) + return 1 + if args.threads > 0: + torch.set_num_threads(args.threads) + dt = args.dtype or ("fp16" if args.device == "cuda" else "fp32") + dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[dt] + device = torch.device(args.device) + + t_load0 = time.perf_counter() + model = load_model(args.ckpt, args.upstream, dtype, device) + if args.compile: + model = torch.compile(model) + if args.device == "cuda": + torch.cuda.synchronize() + load_ms = (time.perf_counter() - t_load0) * 1e3 + + g = torch.Generator().manual_seed(args.seed) + images = torch.rand(1, args.views, 3, 512, 512, generator=g).to(device=device, dtype=dtype) + + S = args.views * (512 // 8) * (512 // 8) + name = torch.cuda.get_device_name(0) if args.device == "cuda" else "cpu" + print(f"torch {torch.__version__} device={args.device} ({name}) dtype={dt} " + f"views={args.views} S={S} threads={torch.get_num_threads()} load={load_ms:.1f} ms") + + def once() -> float: + if args.device == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + with torch.inference_mode(), sdpa_context(args.sdpa): # fresh ctx each call (one-shot) + _ = model(images) + if args.device == "cuda": + torch.cuda.synchronize() + return (time.perf_counter() - t0) * 1e3 + + for _ in range(args.warmup): + once() + ts = [] + for i in range(args.iters): + dt_ms = once() + ts.append(dt_ms) + print(f" iter {i:2d}: {dt_ms:8.1f} ms") + + ts_sorted = sorted(ts) + mn, mx = ts_sorted[0], ts_sorted[-1] + med = ts_sorted[len(ts_sorted) // 2] + mean = statistics.fmean(ts) + vps = args.views / (med / 1e3) + print(f"\nsummary: min={mn:.1f} median={med:.1f} mean={mean:.1f} max={mx:.1f} ms " + f"({vps:.2f} views/s, {1e3/med:.2f} scenes/s)") + print(f"RESULT torch device={args.device} dtype={dt} views={args.views} " + f"load_ms={load_ms:.1f} min_ms={mn:.1f} median_ms={med:.1f} mean_ms={mean:.1f} " + f"max_ms={mx:.1f} views_per_s={vps:.3f}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/model.cpp b/src/model.cpp index 32fb49a..fb16fd5 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -3,9 +3,30 @@ #include #include +#include #include +#include #include #include +#include +#include + +// Parallel-for over [0,n): splits into contiguous chunks, one per hardware +// thread, and joins. body(begin,end) must own disjoint output ranges (the +// post-processing loops below write one output row per index -> race-free). +template +static void parallel_for(int64_t n, F && body) { + if (n <= 0) return; + unsigned nt = std::max(1u, std::thread::hardware_concurrency()); + nt = (unsigned) std::min(nt, n); + if (nt <= 1) { body(int64_t{0}, n); return; } + const int64_t chunk = (n + nt - 1) / nt; + std::vector ths; + ths.reserve(nt); + for (int64_t a = 0; a < n; a += chunk) + ths.emplace_back([&body, a, b = std::min(n, a + chunk)] { body(a, b); }); + for (auto & t : ths) t.join(); +} namespace free_splatter { @@ -138,6 +159,17 @@ bool model::forward(const float * images, int32_t n_views, const float scale = 1.0f / std::sqrt((float) dh); // 1/8 const float eps = h.ln_eps; + // Optional host-phase timing (FREE_SPLATTER_PROFILE=1) to stderr. + const bool prof = std::getenv("FREE_SPLATTER_PROFILE") != nullptr; + auto t_clk = std::chrono::steady_clock::now(); + auto lap = [&](const char * name) { + if (!prof) return; + const auto now = std::chrono::steady_clock::now(); + std::fprintf(stderr, "[profile] %-12s %7.2f ms\n", name, + std::chrono::duration(now - t_clk).count()); + t_clk = now; + }; + const size_t graph_nodes = 8192; ggml_init_params gp = { ggml_tensor_overhead() * graph_nodes + ggml_graph_overhead_custom(graph_nodes, false), @@ -202,9 +234,22 @@ bool model::forward(const float * images, int32_t n_views, ggml_tensor * q3 = ggml_reshape_3d(ctx, q, dh, nh, S); ggml_tensor * k3 = ggml_reshape_3d(ctx, k, dh, nh, S); ggml_tensor * v3 = ggml_reshape_3d(ctx, v, dh, nh, S); - ggml_tensor * qp = ggml_permute(ctx, q3, 0, 2, 1, 3); // [dh, S, nh] + ggml_tensor * qp = ggml_permute(ctx, q3, 0, 2, 1, 3); // [dh, S, nh] (f32 query, required by FA) ggml_tensor * kp = ggml_permute(ctx, k3, 0, 2, 1, 3); - ggml_tensor * vp = ggml_cont(ctx, ggml_permute(ctx, v3, 0, 2, 1, 3)); + ggml_tensor * vp = ggml_permute(ctx, v3, 0, 2, 1, 3); + // On GPU, cast K/V to f16 so the Vulkan coopmat2 (tensor-core) flash-attn + // path is selected — the f32 K/V the projections produce falls off it + // (~4x slower FA). softmax still accumulates in f32 (GGML_PREC_F32 below): + // this is f16 *inputs*, not an f16 softmax. CPU has no tensor cores and + // its tiled FA handles f32 fine, so CPU keeps f32 K/V — leaving the + // CPU-f32 strict gate and the CPU-f16 head check byte-for-byte unchanged; + // only the (already --scale-widened) GPU path moves. + if (!be.is_cpu() && l.wk->type == GGML_TYPE_F16) { + kp = ggml_cast(ctx, kp, GGML_TYPE_F16); + vp = ggml_cast(ctx, vp, GGML_TYPE_F16); + } else { + vp = ggml_cont(ctx, vp); // FA requires V contiguous + } ggml_tensor * fa = ggml_flash_attn_ext(ctx, qp, kp, vp, nullptr, scale, 0.0f, 0.0f); ggml_flash_attn_ext_set_prec(fa, GGML_PREC_F32); ggml_tensor * attn = ggml_reshape_2d(ctx, fa, nh * dh, S); // [D, S] @@ -234,20 +279,25 @@ bool model::forward(const float * images, int32_t n_views, ggml_cgraph * gf = ggml_new_graph_custom(ctx, graph_nodes, false); ggml_build_forward_expand(gf, logits); for (ggml_tensor * t : tap_list) ggml_build_forward_expand(gf, t); + lap("graph_build"); if (!ggml_gallocr_alloc_graph(be.galloc, gf)) { error = "graph alloc failed"; ggml_free(ctx); return false; } + lap("alloc"); ggml_backend_tensor_set(img, images, 0, (size_t) N * C * IMG * IMG * sizeof(float)); + lap("upload"); if (ggml_backend_graph_compute(be.be, gf) != GGML_STATUS_SUCCESS) { error = "graph compute failed"; ggml_free(ctx); return false; } + lap("compute"); // Read head logits [U,S] (memory j + U*s) and unshuffle on the host. std::vector hl((size_t) U * S); ggml_backend_tensor_get(logits, hl.data(), 0, hl.size() * sizeof(float)); + lap("readback"); // out: render-ready gaussians, row-major [N*IMG*IMG, Gc], row = // ((n*IMG + hh)*IMG + ww). gaussians_raw is the same but pre-activation. @@ -257,36 +307,43 @@ bool model::forward(const float * images, int32_t n_views, return images[(((n * C) + ch) * IMG + hh) * IMG + ww]; }; const float C0 = 0.28209479177387814f; - for (int64_t s = 0; s < S; s++) { - const int64_t n = s / TPV, t = s % TPV, hp = t / G, wp = t % G; - for (int64_t j = 0; j < U; j++) { - const int64_t p = j / (P * Gc), q = (j / Gc) % P, c = j % Gc; - const int64_t hh = hp * P + p, ww = wp * P + q; - float val = hl[(size_t) j + (size_t) U * s]; - if (h.sh_residual && c >= 3 && c < 6) { // RGB2SH into SH-DC term - val += (img_at(n, c - 3, hh, ww) - 0.5f) / C0; + parallel_for(S, [&](int64_t s0, int64_t s1) { + for (int64_t s = s0; s < s1; s++) { + const int64_t n = s / TPV, t = s % TPV, hp = t / G, wp = t % G; + for (int64_t j = 0; j < U; j++) { + const int64_t p = j / (P * Gc), q = (j / Gc) % P, c = j % Gc; + const int64_t hh = hp * P + p, ww = wp * P + q; + float val = hl[(size_t) j + (size_t) U * s]; + if (h.sh_residual && c >= 3 && c < 6) { // RGB2SH into SH-DC term + val += (img_at(n, c - 3, hh, ww) - 0.5f) / C0; + } + raw[(size_t) ((n * IMG + hh) * IMG + ww) * Gc + c] = val; } - raw[(size_t) ((n * IMG + hh) * IMG + ww) * Gc + c] = val; } - } + }); + lap("unshuffle"); - // gaussians_raw tap (pre-activation), then activations -> gaussians. + // gaussians_raw tap (pre-activation), emitted before we mutate `raw`; then + // activate in place and move into `out` (no 48 MB copy). if (sink) sink("gaussians_raw", raw.data(), PIX, Gc); - out = raw; const float smin = h.scale_min_act, smax = h.scale_max_act; auto sigmoid = [](float v) { return 1.0f / (1.0f + std::exp(-v)); }; - for (int64_t r = 0; r < PIX; r++) { - float * g = out.data() + (size_t) r * Gc; - // xyz (0:3) and SH (3:15) pass through; opacity sigmoid; scale mapped - // sigmoid; rotation (19:23) L2-normalized (quat order w,x,y,z). - g[15] = sigmoid(g[15]); - for (int c = 16; c < 19; c++) g[c] = smin + (smax - smin) * sigmoid(g[c]); - float nrm = 0.0f; - for (int c = 19; c < 23; c++) nrm += g[c] * g[c]; - nrm = std::sqrt(nrm) + 1e-12f; - for (int c = 19; c < 23; c++) g[c] /= nrm; - } - if (sink) sink("gaussians", out.data(), PIX, Gc); + parallel_for(PIX, [&](int64_t r0, int64_t r1) { + for (int64_t r = r0; r < r1; r++) { + float * g = raw.data() + (size_t) r * Gc; + // xyz (0:3) and SH (3:15) pass through; opacity sigmoid; scale mapped + // sigmoid; rotation (19:23) L2-normalized (quat order w,x,y,z). + g[15] = sigmoid(g[15]); + for (int c = 16; c < 19; c++) g[c] = smin + (smax - smin) * sigmoid(g[c]); + float nrm = 0.0f; + for (int c = 19; c < 23; c++) nrm += g[c] * g[c]; + nrm = std::sqrt(nrm) + 1e-12f; + for (int c = 19; c < 23; c++) g[c] /= nrm; + } + }); + lap("activation"); + if (sink) sink("gaussians", raw.data(), PIX, Gc); + out = std::move(raw); // Stream graph taps to the sink one at a time (read into a reused temp, emit, // free) so host memory never holds all activations at once.