Uploaded using `kernel-builder`.

Browse files

Files changed (5) hide show

build/torch210-cxx11-cu128-x86_64-linux/__init__.py +138 -0
build/torch210-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch210-cxx11-cu128-x86_64-linux/_torch_harmonics_kernels_cuda_19700101000000.abi3.so +3 -0
build/torch210-cxx11-cu128-x86_64-linux/metadata.json +20 -0
build/torch210-cxx11-cu128-x86_64-linux/torch_harmonics_kernels/__init__.py +26 -0

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

	@@ -0,0 +1,138 @@

+from __future__ import annotations
+from typing import Tuple
+import torch
+from ._ops import ops
+def _stable_argsort(x: torch.Tensor) -> torch.Tensor:
+    try:
+        return torch.argsort(x, stable=True)
+    except TypeError:
+        return torch.argsort(x)
+def preprocess_psi(
+    kernel_size: int,
+    nlat_out: int,
+    ker_idx: torch.Tensor,
+    row_idx: torch.Tensor,
+    col_idx: torch.Tensor,
+    vals: torch.Tensor,
+) -> torch.Tensor:
+    device = ker_idx.device
+    order = _stable_argsort(ker_idx.to("cpu"))
+    sorted_ker = ker_idx.to("cpu").index_select(0, order)
+    sorted_row = row_idx.to("cpu").index_select(0, order)
+    sorted_col = col_idx.to("cpu").index_select(0, order)
+    sorted_vals = vals.to("cpu").index_select(0, order)
+    ker_idx.copy_(sorted_ker.to(device))
+    row_idx.copy_(sorted_row.to(device))
+    col_idx.copy_(sorted_col.to(device))
+    vals.copy_(sorted_vals.to(device))
+    if sorted_row.numel() == 0:
+        return torch.zeros(1, dtype=row_idx.dtype, device=device)
+    row_breaks = torch.nonzero(sorted_row[1:] != sorted_row[:-1], as_tuple=False).flatten() + 1
+    roff_idx = torch.empty(row_breaks.numel() + 2, dtype=row_idx.dtype, device=device)
+    roff_idx[0] = 0
+    if row_breaks.numel():
+        roff_idx[1:-1] = row_breaks.to(device=device, dtype=row_idx.dtype)
+    roff_idx[-1] = sorted_row.numel()
+    max_rows = kernel_size * nlat_out
+    if roff_idx.numel() - 1 > max_rows:
+        raise RuntimeError(
+            f"found more rows in the K COOs than Ho*K ({max_rows})"
+        )
+    return roff_idx
+def cuda_kernels_is_available() -> bool:
+    return True
+def optimized_kernels_is_available() -> bool:
+    return True
+def disco_forward(
+    inp: torch.Tensor,
+    roff_idx: torch.Tensor,
+    ker_idx: torch.Tensor,
+    row_idx: torch.Tensor,
+    col_idx: torch.Tensor,
+    vals: torch.Tensor,
+    kernel_size: int,
+    nlat_out: int,
+    nlon_out: int,
+) -> torch.Tensor:
+    return ops.disco_forward(
+        inp, roff_idx, ker_idx, row_idx, col_idx, vals, kernel_size, nlat_out, nlon_out
+    )
+def disco_backward(
+    inp: torch.Tensor,
+    roff_idx: torch.Tensor,
+    ker_idx: torch.Tensor,
+    row_idx: torch.Tensor,
+    col_idx: torch.Tensor,
+    vals: torch.Tensor,
+    kernel_size: int,
+    nlat_out: int,
+    nlon_out: int,
+) -> torch.Tensor:
+    return ops.disco_backward(
+        inp, roff_idx, ker_idx, row_idx, col_idx, vals, kernel_size, nlat_out, nlon_out
+    )
+def attention_forward(
+    kx: torch.Tensor,
+    vx: torch.Tensor,
+    qy: torch.Tensor,
+    quad_weights: torch.Tensor,
+    col_idx: torch.Tensor,
+    row_off: torch.Tensor,
+    nlon_in: int,
+    nlat_out: int,
+    nlon_out: int,
+) -> torch.Tensor:
+    return ops.attention_forward(
+        kx, vx, qy, quad_weights, col_idx, row_off, nlon_in, nlat_out, nlon_out
+    )
+def attention_backward(
+    kx: torch.Tensor,
+    vx: torch.Tensor,
+    qy: torch.Tensor,
+    dy: torch.Tensor,
+    quad_weights: torch.Tensor,
+    col_idx: torch.Tensor,
+    row_off: torch.Tensor,
+    nlon_in: int,
+    nlat_out: int,
+    nlon_out: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return ops.attention_backward(
+        kx, vx, qy, dy, quad_weights, col_idx, row_off, nlon_in, nlat_out, nlon_out
+    )
+__all__ = [
+    "attention_backward",
+    "attention_forward",
+    "cuda_kernels_is_available",
+    "disco_backward",
+    "disco_forward",
+    "ops",
+    "optimized_kernels_is_available",
+    "preprocess_psi",
+]

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

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

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

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:45f06869044fe89b321c8cc772623452a4f8701cae21632afd77487ca99d3158
+size 144599960

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

	@@ -0,0 +1,20 @@

+{
+  "version": 1,
+  "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-cu128-x86_64-linux/torch_harmonics_kernels/__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")))