liangsu9988 commited on Jul 20

Commit

ff2eba6

verified ·

1 Parent(s): 662c1da

Uploaded using `kernel-builder`.

Browse files

Files changed (31) hide show

benchmarks/benchmark.py +172 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +49 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +36 -0
build/torch211-cxx11-cu128-x86_64-linux/small_matrix_cholesky/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +49 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +36 -0
build/torch211-cxx11-cu130-x86_64-linux/small_matrix_cholesky/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +49 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +36 -0
build/torch212-cxx11-cu130-x86_64-linux/small_matrix_cholesky/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +49 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +36 -0
build/torch212-cxx11-cu132-x86_64-linux/small_matrix_cholesky/__init__.py +26 -0
build/torch213-cxx11-cu130-x86_64-linux/__init__.py +49 -0
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so +3 -0
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +36 -0
build/torch213-cxx11-cu130-x86_64-linux/small_matrix_cholesky/__init__.py +26 -0
build/torch213-cxx11-cu132-x86_64-linux/__init__.py +49 -0
build/torch213-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu132-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so +3 -0
build/torch213-cxx11-cu132-x86_64-linux/metadata.json +36 -0
build/torch213-cxx11-cu132-x86_64-linux/small_matrix_cholesky/__init__.py +26 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,172 @@

+#!/usr/bin/env python3
+"""Benchmark small FP32 Cholesky against preallocated PyTorch POTRF."""
+from __future__ import annotations
+import argparse
+import importlib
+import math
+import statistics
+import sys
+from pathlib import Path
+import torch
+TESTS = Path(__file__).resolve().parents[1] / "tests"
+sys.path.insert(0, str(TESTS))
+from _source_loader import load_source_ops  # noqa: E402
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("small_matrix_cholesky")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def make_spd(batch: int, n: int) -> torch.Tensor:
+    generator = torch.Generator(device="cuda").manual_seed(41000 + n + batch)
+    x = torch.randn(
+        batch,
+        n,
+        n,
+        device="cuda",
+        dtype=torch.float32,
+        generator=generator,
+    ) / n**0.5
+    return (
+        x @ x.transpose(-1, -2)
+        + 0.5 * torch.eye(n, device="cuda", dtype=torch.float32)
+    ).contiguous()
+def median_ms(fn, warmup: int, iterations: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    samples: list[float] = []
+    for _ in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end))
+    return statistics.median(samples)
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--backend", choices=["source", "installed"], default="source"
+    )
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--registration-include", default=None)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iterations", type=int, default=50)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = (
+        load_source_ops(args.registration_include)
+        if args.backend == "source"
+        else load_installed_ops(args.artifact)
+    )
+    shapes = [(4096, 32), (1024, 64), (256, 128)]
+    candidate_times: list[float] = []
+    baseline_times: list[float] = []
+    compiled_times: list[float] = []
+    def pytorch_reference(
+        input: torch.Tensor,
+        output: torch.Tensor,
+        info: torch.Tensor,
+    ) -> torch.Tensor:
+        torch.linalg.cholesky_ex(
+            input,
+            check_errors=False,
+            out=(output, info),
+        )
+        return output
+    compiled_reference = torch.compile(pytorch_reference, fullgraph=True)
+    print(
+        "batch,n,candidate_ms,pytorch_eager_ms,pytorch_compile_ms,"
+        "speedup_eager,speedup_compile,candidate_tflops,io_gbps"
+    )
+    for batch, n in shapes:
+        input = make_spd(batch, n)
+        candidate_output = torch.empty_like(input)
+        baseline_output = torch.empty_like(input)
+        info = torch.empty(batch, device="cuda", dtype=torch.int32)
+        def candidate() -> None:
+            ops.cholesky_small_fp32(input, out=candidate_output)
+        def baseline() -> None:
+            pytorch_reference(input, baseline_output, info)
+        def baseline_compiled() -> None:
+            compiled_reference(input, baseline_output, info)
+        candidate()
+        baseline()
+        baseline_compiled()
+        torch.cuda.synchronize()
+        torch.testing.assert_close(
+            candidate_output,
+            baseline_output,
+            rtol=5e-4,
+            atol=2e-4,
+        )
+        candidate_ms = median_ms(candidate, args.warmup, args.iterations)
+        baseline_ms = median_ms(baseline, args.warmup, args.iterations)
+        compiled_ms = median_ms(
+            baseline_compiled, args.warmup, args.iterations
+        )
+        flops = batch * n**3 / 3.0
+        tflops = flops / (candidate_ms * 1e-3) / 1e12
+        io_bytes = 2 * batch * n * n * 4
+        io_gbps = io_bytes / (candidate_ms * 1e-3) / 1e9
+        candidate_times.append(candidate_ms)
+        baseline_times.append(baseline_ms)
+        compiled_times.append(compiled_ms)
+        print(
+            f"{batch},{n},{candidate_ms:.6f},{baseline_ms:.6f},"
+            f"{compiled_ms:.6f},{baseline_ms / candidate_ms:.3f},"
+            f"{compiled_ms / candidate_ms:.3f},{tflops:.3f},{io_gbps:.3f}"
+        )
+    candidate_geomean = math.exp(
+        sum(math.log(value) for value in candidate_times)
+        / len(candidate_times)
+    )
+    baseline_geomean = math.exp(
+        sum(math.log(value) for value in baseline_times)
+        / len(baseline_times)
+    )
+    compiled_geomean = math.exp(
+        sum(math.log(value) for value in compiled_times)
+        / len(compiled_times)
+    )
+    print(f"candidate_geomean_ms={candidate_geomean:.6f}")
+    print(f"pytorch_geomean_ms={baseline_geomean:.6f}")
+    print(f"pytorch_compile_geomean_ms={compiled_geomean:.6f}")
+    print(
+        f"geomean_speedup_eager={baseline_geomean / candidate_geomean:.3f}"
+    )
+    print(
+        f"geomean_speedup_compile={compiled_geomean / candidate_geomean:.3f}"
+    )
+    return 0
+if __name__ == "__main__":
+    raise SystemExit(main())

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

	@@ -0,0 +1,49 @@

+"""Batched FP32 Cholesky kernels for small CUDA matrices."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(
+    add_op_namespace_prefix("cholesky_small_fp32_out")
+)
+def _cholesky_small_fp32_out_fake(
+    input: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if input.dim() < 2:
+        raise RuntimeError("input must have at least two dimensions")
+    if input.shape != output.shape:
+        raise RuntimeError("output must have the same shape as input")
+    if input.shape[-2] != input.shape[-1]:
+        raise RuntimeError("the last two dimensions must be square")
+    if input.shape[-1] not in (32, 64, 128):
+        raise RuntimeError("supported matrix orders are 32, 64, and 128")
+    return None
+def cholesky_small_fp32(
+    input: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute a lower-triangular Cholesky factor for small FP32 matrices.
+    ``input`` must be a contiguous CUDA FP32 tensor whose last two dimensions
+    are ``(n, n)`` with ``n`` equal to 32, 64, or 128. All leading dimensions
+    are flattened into a batch. The input matrices must be symmetric positive
+    definite. The output's upper triangle is explicitly zero.
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.cholesky_small_fp32_out(input, out)
+    return out
+__all__ = ["cholesky_small_fp32"]

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

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

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

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

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

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

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

	@@ -0,0 +1,36 @@

+{
+  "name": "small-matrix-cholesky",
+  "id": "_small_matrix_cholesky_cuda_f291092",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Sc7qcYjgjttvRwa0SznImIQhKoGy51oSJJjOznvvj18=",
+      "_ops.py": "wQ0b98E+GklbX93H03G+NzmGEGs2jD8ZzF6+DdmylWA=",
+      "_small_matrix_cholesky_cuda_f291092.abi3.so": "sy0ZRW0mJchIXFGhX9dafRLC0hVVT7ljlif/C9D0+Uc=",
+      "small_matrix_cholesky/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "f2910927e5fe824a1212a8147067b687e6b6d2ed",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/small_matrix_cholesky/__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/__init__.py ADDED Viewed

	@@ -0,0 +1,49 @@

+"""Batched FP32 Cholesky kernels for small CUDA matrices."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(
+    add_op_namespace_prefix("cholesky_small_fp32_out")
+)
+def _cholesky_small_fp32_out_fake(
+    input: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if input.dim() < 2:
+        raise RuntimeError("input must have at least two dimensions")
+    if input.shape != output.shape:
+        raise RuntimeError("output must have the same shape as input")
+    if input.shape[-2] != input.shape[-1]:
+        raise RuntimeError("the last two dimensions must be square")
+    if input.shape[-1] not in (32, 64, 128):
+        raise RuntimeError("supported matrix orders are 32, 64, and 128")
+    return None
+def cholesky_small_fp32(
+    input: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute a lower-triangular Cholesky factor for small FP32 matrices.
+    ``input`` must be a contiguous CUDA FP32 tensor whose last two dimensions
+    are ``(n, n)`` with ``n`` equal to 32, 64, or 128. All leading dimensions
+    are flattened into a batch. The input matrices must be symmetric positive
+    definite. The output's upper triangle is explicitly zero.
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.cholesky_small_fp32_out(input, out)
+    return out
+__all__ = ["cholesky_small_fp32"]

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

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

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

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:005558eb0ae2ae44d7eea858f3f726080d871063736fad8f85e61daa8cf2f380
+size 585392

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

	@@ -0,0 +1,36 @@

+{
+  "name": "small-matrix-cholesky",
+  "id": "_small_matrix_cholesky_cuda_f291092",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Sc7qcYjgjttvRwa0SznImIQhKoGy51oSJJjOznvvj18=",
+      "_ops.py": "wQ0b98E+GklbX93H03G+NzmGEGs2jD8ZzF6+DdmylWA=",
+      "_small_matrix_cholesky_cuda_f291092.abi3.so": "AFVY6wrirkTX7qhY8/cmCA2HEGNzb62PheYdqozy84A=",
+      "small_matrix_cholesky/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "f2910927e5fe824a1212a8147067b687e6b6d2ed",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/small_matrix_cholesky/__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/__init__.py ADDED Viewed

	@@ -0,0 +1,49 @@

+"""Batched FP32 Cholesky kernels for small CUDA matrices."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(
+    add_op_namespace_prefix("cholesky_small_fp32_out")
+)
+def _cholesky_small_fp32_out_fake(
+    input: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if input.dim() < 2:
+        raise RuntimeError("input must have at least two dimensions")
+    if input.shape != output.shape:
+        raise RuntimeError("output must have the same shape as input")
+    if input.shape[-2] != input.shape[-1]:
+        raise RuntimeError("the last two dimensions must be square")
+    if input.shape[-1] not in (32, 64, 128):
+        raise RuntimeError("supported matrix orders are 32, 64, and 128")
+    return None
+def cholesky_small_fp32(
+    input: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute a lower-triangular Cholesky factor for small FP32 matrices.
+    ``input`` must be a contiguous CUDA FP32 tensor whose last two dimensions
+    are ``(n, n)`` with ``n`` equal to 32, 64, or 128. All leading dimensions
+    are flattened into a batch. The input matrices must be symmetric positive
+    definite. The output's upper triangle is explicitly zero.
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.cholesky_small_fp32_out(input, out)
+    return out
+__all__ = ["cholesky_small_fp32"]

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

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

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

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:8c29f7d9ecb3691144a91cfa1f99ad42da913db83c495db5ddba59d95106aceb
+size 592168

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

	@@ -0,0 +1,36 @@

+{
+  "name": "small-matrix-cholesky",
+  "id": "_small_matrix_cholesky_cuda_f291092",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Sc7qcYjgjttvRwa0SznImIQhKoGy51oSJJjOznvvj18=",
+      "_ops.py": "wQ0b98E+GklbX93H03G+NzmGEGs2jD8ZzF6+DdmylWA=",
+      "_small_matrix_cholesky_cuda_f291092.abi3.so": "jCn32eyzaRFEqRz6H5mtQtqRPbg8SV213bpZ2VEGrOs=",
+      "small_matrix_cholesky/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "f2910927e5fe824a1212a8147067b687e6b6d2ed",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/small_matrix_cholesky/__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/__init__.py ADDED Viewed

	@@ -0,0 +1,49 @@

+"""Batched FP32 Cholesky kernels for small CUDA matrices."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(
+    add_op_namespace_prefix("cholesky_small_fp32_out")
+)
+def _cholesky_small_fp32_out_fake(
+    input: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if input.dim() < 2:
+        raise RuntimeError("input must have at least two dimensions")
+    if input.shape != output.shape:
+        raise RuntimeError("output must have the same shape as input")
+    if input.shape[-2] != input.shape[-1]:
+        raise RuntimeError("the last two dimensions must be square")
+    if input.shape[-1] not in (32, 64, 128):
+        raise RuntimeError("supported matrix orders are 32, 64, and 128")
+    return None
+def cholesky_small_fp32(
+    input: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute a lower-triangular Cholesky factor for small FP32 matrices.
+    ``input`` must be a contiguous CUDA FP32 tensor whose last two dimensions
+    are ``(n, n)`` with ``n`` equal to 32, 64, or 128. All leading dimensions
+    are flattened into a batch. The input matrices must be symmetric positive
+    definite. The output's upper triangle is explicitly zero.
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.cholesky_small_fp32_out(input, out)
+    return out
+__all__ = ["cholesky_small_fp32"]

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

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

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

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:6c66049b384edcff07a4bffd59d1ef409c259232b9873c6bd99b63c9a790d299
+size 596264

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

	@@ -0,0 +1,36 @@

+{
+  "name": "small-matrix-cholesky",
+  "id": "_small_matrix_cholesky_cuda_f291092",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Sc7qcYjgjttvRwa0SznImIQhKoGy51oSJJjOznvvj18=",
+      "_ops.py": "wQ0b98E+GklbX93H03G+NzmGEGs2jD8ZzF6+DdmylWA=",
+      "_small_matrix_cholesky_cuda_f291092.abi3.so": "bGYEmzhO3P8HpL/9WdHvQJwlkjK5hzxr2ZtjyaeQ0pk=",
+      "small_matrix_cholesky/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "f2910927e5fe824a1212a8147067b687e6b6d2ed",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/small_matrix_cholesky/__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/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,49 @@

+"""Batched FP32 Cholesky kernels for small CUDA matrices."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(
+    add_op_namespace_prefix("cholesky_small_fp32_out")
+)
+def _cholesky_small_fp32_out_fake(
+    input: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if input.dim() < 2:
+        raise RuntimeError("input must have at least two dimensions")
+    if input.shape != output.shape:
+        raise RuntimeError("output must have the same shape as input")
+    if input.shape[-2] != input.shape[-1]:
+        raise RuntimeError("the last two dimensions must be square")
+    if input.shape[-1] not in (32, 64, 128):
+        raise RuntimeError("supported matrix orders are 32, 64, and 128")
+    return None
+def cholesky_small_fp32(
+    input: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute a lower-triangular Cholesky factor for small FP32 matrices.
+    ``input`` must be a contiguous CUDA FP32 tensor whose last two dimensions
+    are ``(n, n)`` with ``n`` equal to 32, 64, or 128. All leading dimensions
+    are flattened into a batch. The input matrices must be symmetric positive
+    definite. The output's upper triangle is explicitly zero.
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.cholesky_small_fp32_out(input, out)
+    return out
+__all__ = ["cholesky_small_fp32"]

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

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

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

build/torch213-cxx11-cu130-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:10d82c9d01a56acfadadebf1456e4947136d80f1cbf3b9347f7a106fc13f4d51
+size 592008

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

	@@ -0,0 +1,36 @@

+{
+  "name": "small-matrix-cholesky",
+  "id": "_small_matrix_cholesky_cuda_f291092",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Sc7qcYjgjttvRwa0SznImIQhKoGy51oSJJjOznvvj18=",
+      "_ops.py": "wQ0b98E+GklbX93H03G+NzmGEGs2jD8ZzF6+DdmylWA=",
+      "_small_matrix_cholesky_cuda_f291092.abi3.so": "ENgsnQGlas+trevxRW5JRxNtgPHL87k0f3oQb8E/TVE=",
+      "small_matrix_cholesky/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "f2910927e5fe824a1212a8147067b687e6b6d2ed",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu130-x86_64-linux/small_matrix_cholesky/__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/torch213-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,49 @@

+"""Batched FP32 Cholesky kernels for small CUDA matrices."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(
+    add_op_namespace_prefix("cholesky_small_fp32_out")
+)
+def _cholesky_small_fp32_out_fake(
+    input: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if input.dim() < 2:
+        raise RuntimeError("input must have at least two dimensions")
+    if input.shape != output.shape:
+        raise RuntimeError("output must have the same shape as input")
+    if input.shape[-2] != input.shape[-1]:
+        raise RuntimeError("the last two dimensions must be square")
+    if input.shape[-1] not in (32, 64, 128):
+        raise RuntimeError("supported matrix orders are 32, 64, and 128")
+    return None
+def cholesky_small_fp32(
+    input: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute a lower-triangular Cholesky factor for small FP32 matrices.
+    ``input`` must be a contiguous CUDA FP32 tensor whose last two dimensions
+    are ``(n, n)`` with ``n`` equal to 32, 64, or 128. All leading dimensions
+    are flattened into a batch. The input matrices must be symmetric positive
+    definite. The output's upper triangle is explicitly zero.
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.cholesky_small_fp32_out(input, out)
+    return out
+__all__ = ["cholesky_small_fp32"]

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

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

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

build/torch213-cxx11-cu132-x86_64-linux/_small_matrix_cholesky_cuda_f291092.abi3.so ADDED Viewed

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

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

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

	@@ -0,0 +1,36 @@

+{
+  "name": "small-matrix-cholesky",
+  "id": "_small_matrix_cholesky_cuda_f291092",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Sc7qcYjgjttvRwa0SznImIQhKoGy51oSJJjOznvvj18=",
+      "_ops.py": "wQ0b98E+GklbX93H03G+NzmGEGs2jD8ZzF6+DdmylWA=",
+      "_small_matrix_cholesky_cuda_f291092.abi3.so": "s57SdruAItlWiLBWDrOuSs754PMeKwq0rCJPWk2hIJ8=",
+      "small_matrix_cholesky/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "f2910927e5fe824a1212a8147067b687e6b6d2ed",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/small_matrix_cholesky/__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")))