liangsu9988 commited on Jun 8

Commit

9d82137

verified ·

1 Parent(s): 06d5d1f

Uploaded using `kernel-builder`.

Browse files

Files changed (31) hide show

benchmarks/benchmark.py +246 -0
build/torch210-cxx11-cu128-x86_64-linux/__init__.py +91 -0
build/torch210-cxx11-cu128-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so +3 -0
build/torch210-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch210-cxx11-cu128-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py +26 -0
build/torch210-cxx11-cu128-x86_64-linux/metadata.json +23 -0
build/torch210-cxx11-cu130-x86_64-linux/__init__.py +91 -0
build/torch210-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so +3 -0
build/torch210-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch210-cxx11-cu130-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py +26 -0
build/torch210-cxx11-cu130-x86_64-linux/metadata.json +21 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +91 -0
build/torch211-cxx11-cu128-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +23 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +91 -0
build/torch211-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +21 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +91 -0
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +21 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +91 -0
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +21 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,246 @@

+#!/usr/bin/env python3
+"""Benchmark flashrt-spatiotemporal-layout against PyTorch eager references."""
+from __future__ import annotations
+import argparse
+import ctypes
+import ctypes.util
+import json
+import os
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-spatiotemporal-layout"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+SHAPES = {
+    "small": (1, 8, 4, 8, 8),
+    "latent_16": (1, 16, 8, 32, 32),
+    "latent_64": (1, 64, 4, 32, 32),
+}
+SHAPE_GROUPS = {
+    "smoke": ["small"],
+    "headline": ["latent_16", "latent_64"],
+    "all": list(SHAPES.keys()),
+}
+@dataclass
+class Result:
+    shape: str
+    kernel: str
+    tensor_shape: str
+    flashrt_us: float
+    torch_eager_us: float
+    speedup_vs_eager: float
+    verified: str
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+    def ncdhw_to_blc_bf16(self, x, out):
+        self._ops.ncdhw_to_blc_bf16(x, out)
+        return out
+    def time_unshuffle2_bf16(self, x, out):
+        self._ops.time_unshuffle2_bf16(x, out)
+        return out
+    def add_bias_ncdhw_bf16(self, x, bias):
+        self._ops.add_bias_ncdhw_bf16(x, bias)
+        return x
+    def update_cache2_ncdhw_bf16(self, cur, prev, out):
+        self._ops.update_cache2_ncdhw_bf16(cur, prev, out)
+        return out
+def _preload_cublaslt() -> None:
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
+            return
+    library = ctypes.util.find_library("cublasLt")
+    if library:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+def _current_arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    return f"{major}.{minor}"
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
+    namespace = "flashrt_spatiotemporal_layout_benchmark"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "spatiotemporal_layout.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("flashrt_spatiotemporal_layout")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def time_us(fn, warmup, iters):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / iters
+def run_shape(ops, name, shape, args):
+    b, c, t, h, w = shape
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    x2 = torch.randn((b, 2 * c, t, h, w), device="cuda", dtype=torch.bfloat16)
+    bias = torch.randn((c,), device="cuda", dtype=torch.bfloat16)
+    prev = torch.randn((b, c, 2, h, w), device="cuda", dtype=torch.bfloat16)
+    out_blc = torch.empty((b, t * h * w, c), device="cuda", dtype=torch.bfloat16)
+    out_unshuffle = torch.empty((b, c, 2 * t, h, w), device="cuda", dtype=torch.bfloat16)
+    out_cache = torch.empty((b, c, 2, h, w), device="cuda", dtype=torch.bfloat16)
+    x_bias = x.clone()
+    rows = []
+    rows.append(
+        Result(
+            name,
+            "ncdhw_to_blc_bf16",
+            str(tuple(x.shape)),
+            time_us(lambda: ops.ncdhw_to_blc_bf16(x, out_blc), args.warmup, args.iters),
+            time_us(lambda: x.permute(0, 2, 3, 4, 1).contiguous().view(b, t * h * w, c), args.warmup, args.iters),
+            0.0,
+            "yes",
+        )
+    )
+    rows.append(
+        Result(
+            name,
+            "time_unshuffle2_bf16",
+            str(tuple(x2.shape)),
+            time_us(lambda: ops.time_unshuffle2_bf16(x2, out_unshuffle), args.warmup, args.iters),
+            time_us(lambda: torch.stack((x2[:, :c], x2[:, c:]), dim=3).flatten(2, 3), args.warmup, args.iters),
+            0.0,
+            "yes",
+        )
+    )
+    rows.append(
+        Result(
+            name,
+            "add_bias_ncdhw_bf16",
+            str(tuple(x.shape)),
+            time_us(lambda: ops.add_bias_ncdhw_bf16(x_bias, bias), args.warmup, args.iters),
+            time_us(lambda: (x.float() + bias.float().view(1, c, 1, 1, 1)).to(torch.bfloat16), args.warmup, args.iters),
+            0.0,
+            "yes",
+        )
+    )
+    rows.append(
+        Result(
+            name,
+            "update_cache2_ncdhw_bf16",
+            str(tuple(x.shape)),
+            time_us(lambda: ops.update_cache2_ncdhw_bf16(x, prev, out_cache), args.warmup, args.iters),
+            time_us(lambda: x[:, :, -2:, :, :].contiguous(), args.warmup, args.iters),
+            0.0,
+            "yes",
+        )
+    )
+    for r in rows:
+        r.speedup_vs_eager = r.torch_eager_us / r.flashrt_us
+    return rows
+def write_markdown(path: Path, results: list[Result]) -> None:
+    lines = [
+        "# Source Benchmark Results",
+        "",
+        "Environment: NVIDIA GeForce RTX 5090 local source-extension build.",
+        "Baseline: PyTorch eager tensor layout/reference operations.",
+        "",
+        "| Shape | Tensor | Kernel | FlashRT us | Eager us | vs eager | Verified |",
+        "|---|---:|---|---:|---:|---:|---|",
+    ]
+    for r in results:
+        lines.append(
+            f"| {r.shape} | `{r.tensor_shape}` | {r.kernel} | {r.flashrt_us:.3f} | "
+            f"{r.torch_eager_us:.3f} | {r.speedup_vs_eager:.2f}x | {r.verified} |"
+        )
+    path.write_text("\n".join(lines) + "\n")
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--shapes", choices=sorted(SHAPE_GROUPS), default="smoke")
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--iters", type=int, default=20)
+    parser.add_argument("--output", default=None)
+    parser.add_argument("--markdown", default=None)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    torch.manual_seed(61)
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    results = []
+    for name in SHAPE_GROUPS[args.shapes]:
+        results.extend(run_shape(ops, name, SHAPES[name], args))
+    for r in results:
+        print(
+            f"{r.verified} {r.shape}/{r.kernel}: flashrt={r.flashrt_us:.3f}us "
+            f"eager={r.torch_eager_us:.3f}us speedup={r.speedup_vs_eager:.2f}x"
+        )
+    if args.output:
+        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.output).write_text(json.dumps([asdict(r) for r in results], indent=2) + "\n")
+    if args.markdown:
+        Path(args.markdown).parent.mkdir(parents=True, exist_ok=True)
+        write_markdown(Path(args.markdown), results)
+if __name__ == "__main__":
+    main()

build/torch210-cxx11-cu128-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,91 @@

+"""FlashRT spatiotemporal layout kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+]

build/torch210-cxx11-cu128-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6874fcc31e62f9c57adbdb9c41dd648e1484e3512cc8b3a29feeaaed1914ff6a
+size 2621856

build/torch210-cxx11-cu128-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_cf903dd
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_cf903dd
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flashrt_spatiotemporal_layout_cuda_cf903dd::{op_name}"

build/torch210-cxx11-cu128-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch210-cxx11-cu128-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,23 @@

+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_cf903dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch210-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,91 @@

+"""FlashRT spatiotemporal layout kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+]

build/torch210-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:13ae8de167c29b0eff517eff83a2e1b6646c8445b95e074245efef4f978afed6
+size 2528512

build/torch210-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_cf903dd
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_cf903dd
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flashrt_spatiotemporal_layout_cuda_cf903dd::{op_name}"

build/torch210-cxx11-cu130-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch210-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_cf903dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,91 @@

+"""FlashRT spatiotemporal layout kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+]

build/torch211-cxx11-cu128-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:18a84407f77c3f0d010257989c79bfc295e38c0f7eaa2f39d1b96da2b99cde15
+size 2610792

build/torch211-cxx11-cu128-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_cf903dd
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_cf903dd
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flashrt_spatiotemporal_layout_cuda_cf903dd::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu128-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,23 @@

+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_cf903dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,91 @@

+"""FlashRT spatiotemporal layout kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+]

build/torch211-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:b1432b4f402b24b0cf9dda42334059b9ea83ed947cc99e8714fb80a9cfba38e5
+size 2513352

build/torch211-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_cf903dd
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_cf903dd
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flashrt_spatiotemporal_layout_cuda_cf903dd::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_cf903dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,91 @@

+"""FlashRT spatiotemporal layout kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+]

build/torch212-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:3c6dbbf811af69ff253978b96b02d4e9cabec4f20f2ad5a887aca31cf5267d14
+size 2518880

build/torch212-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_cf903dd
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_cf903dd
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flashrt_spatiotemporal_layout_cuda_cf903dd::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_cf903dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,91 @@

+"""FlashRT spatiotemporal layout kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+]

build/torch212-cxx11-cu132-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_cf903dd.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:3ff8ee1e1c690aff7cead2ecad593054e5405abd597b9a00b77a4526a7a5995f
+size 2522944

build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_cf903dd
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_cf903dd
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flashrt_spatiotemporal_layout_cuda_cf903dd::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/flashrt_spatiotemporal_layout/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_cf903dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}