vramsuite 0.5.0__py3-none-any.whl

vramsuite/__init__.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import importlib.metadata
+
+from vramsuite.core.doctor import run_doctor
+from vramsuite.core.fingerprint import collect_fingerprint
+from vramsuite.core.vramcard import create_vramcard, load_vramcard, save_vramcard
+
+__version__ = importlib.metadata.version("vramsuite")
+
+__all__ = [
+    "__version__",
+    "run_doctor",
+    "collect_fingerprint",
+    "create_vramcard",
+    "save_vramcard",
+    "load_vramcard",
+]

vramsuite/__main__.py

@@ -0,0 +1,59 @@
+"""
+`python -m vramsuite` and the entry point of the Windows exe.
+
+Double-clicking the exe opens a console that closes as soon as the program
+exits. In that case the exe runs `doctor` and waits for Enter instead of
+printing the help and vanishing.
+"""
+
+from __future__ import annotations
+
+import os
+import platform
+import sys
+import traceback
+
+from vramsuite.cli.main import app
+
+
+def _started_by_double_click() -> bool:
+    """
+    True when no other program shares this console: Explorer created it for us.
+
+    From a terminal, the shell is attached too. The one-file exe runs as two
+    processes (the PyInstaller bootloader and its Python child, this one),
+    both attached to the console, so the bootloader parent does not count.
+    """
+    if platform.system() != "Windows":
+        return False
+    import ctypes
+
+    kernel32 = ctypes.windll.kernel32
+    kernel32.GetConsoleProcessList.argtypes = [ctypes.POINTER(ctypes.c_uint32), ctypes.c_uint32]
+    kernel32.GetConsoleProcessList.restype = ctypes.c_uint32
+    processes = (ctypes.c_uint32 * 16)()
+    count = kernel32.GetConsoleProcessList(processes, 16)
+
+    ours = {os.getpid()}
+    if getattr(sys, "frozen", False):
+        ours.add(os.getppid())
+    return 0 < count <= 16 and not set(processes[:count]) - ours
+
+
+def main() -> None:
+    if len(sys.argv) == 1 and _started_by_double_click():
+        sys.argv += ["doctor", "--no-write-card"]
+        try:
+            app()
+        except SystemExit:
+            pass
+        except Exception:
+            # Show the error instead of a window that closes before it can be read.
+            traceback.print_exc()
+        input("\nPress Enter to close...")
+        return
+    app()
+
+
+if __name__ == "__main__":
+    main()

vramsuite/backends/cublas.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+import ctypes
+import platform
+from dataclasses import asdict, dataclass
+from types import TracebackType
+from typing import Any
+
+from vramsuite.backends.cublaslt.errors import decode_cublas_status
+from vramsuite.backends.native_loader import load_native_library
+
+
+CUBLAS_STATUS_SUCCESS = 0
+
+CUBLAS_OP_N = 0
+CUBLAS_OP_T = 1
+CUBLAS_OP_C = 2
+
+@dataclass(frozen=True)
+class CublasInfo:
+    available: bool
+    error: str | None
+    library_path: str | None
+    handle_created: bool
+    version: int | None
+    version_text: str | None
+
+
+def _candidate_library_names() -> list[str]:
+    system = platform.system()
+
+    if system == "Windows":
+        return [
+            "cublas64_13.dll",
+            "cublas64_12.dll",
+            "cublas64_11.dll",
+        ]
+    
+    if system == "Linux":
+        return [
+            "libcublas.so",
+            "libcublas.so.13",
+            "libcublas.so.12",
+            "libcublas.so.11",
+        ]
+    
+    return ["libcublas.so"]
+
+
+def _check(status: int, operation: str) -> str | None:
+    if status == CUBLAS_STATUS_SUCCESS:
+        return None
+
+    return f"{operation} failed: {decode_cublas_status(status)}"
+
+def _format_cublas_version(version: int | None) -> str | None:
+    if version is None:
+        return None
+    
+    major = version // 10000
+    minor = (version % 10000) // 100
+    patch = version % 100
+
+    return f"{major}.{minor}.{patch}"
+
+def _configure_signatures(cublas: ctypes.CDLL) -> None:
+    cublas.cublasCreate_v2.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
+    cublas.cublasCreate_v2.restype = ctypes.c_int
+
+    cublas.cublasDestroy_v2.argtypes = [ctypes.c_void_p]
+    cublas.cublasDestroy_v2.restype = ctypes.c_int
+
+    cublas.cublasGetVersion_v2.argtypes = [
+        ctypes.c_void_p,
+        ctypes.POINTER(ctypes.c_int),
+    ]
+    cublas.cublasGetVersion_v2.restype = ctypes.c_int
+
+    cublas.cublasSgemm_v2.argtypes = [
+        ctypes.c_void_p,                    # handle
+        ctypes.c_int,                       # transa
+        ctypes.c_int,                       # transb
+        ctypes.c_int,                       # m
+        ctypes.c_int,                       # n
+        ctypes.c_int,                       # k
+        ctypes.POINTER(ctypes.c_float),     # alpha
+        ctypes.c_void_p,                    # A
+        ctypes.c_int,                       # lda
+        ctypes.c_void_p,                    # B
+        ctypes.c_int,                       # ldb
+        ctypes.POINTER(ctypes.c_float),     # beta
+        ctypes.c_void_p,                    # c
+        ctypes.c_int,                       # ldc
+    ]
+    cublas.cublasSgemm_v2.restype = ctypes.c_int
+
+
+class CublasRuntime:
+    def __init__(self) -> None:
+        self._native_library = None
+        self.cublas: ctypes.CDLL | None = None
+        self.library_path: str | None = None
+        self.handle: ctypes.c_void_p | None = None
+
+    def __enter__(self) -> "CublasRuntime":
+        venv_patterns = (
+            ["cublas*.dll"]
+            if platform.system() == "Windows"
+            else ["libcublas.so*"]
+        )
+
+        self._native_library = load_native_library(
+            logical_name="cublas",
+            candidates=_candidate_library_names(),
+            venv_patterns=venv_patterns,
+        )
+
+        try:
+            loaded = self._native_library.__enter__()
+
+            self.cublas = loaded.require_library()
+            self.library_path = loaded.library_path
+
+            _configure_signatures(self.cublas)
+
+            self.handle = self._create_handle()
+
+            return self
+        
+        except Exception:
+            if self._native_library is not None:
+                self._native_library.__exit__(None, None, None)
+            
+            self._native_library = None
+            self.cublas = None
+            self.library_path = None
+            self.handle = None
+
+            raise
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if self.handle is not None:
+            try:
+                self._destroy_handle(self.handle)
+            finally:
+                self.handle = None
+
+        self.cublas = None
+        self.library_path = None
+
+        if self._native_library is not None:
+            self._native_library.__exit__(exc_type, exc, traceback)
+
+        self._native_library = None
+
+    def _require_cublas(self) -> ctypes.CDLL:
+        if self.cublas is None:
+            raise RuntimeError("cuBLAS runtime is not loaded")
+        
+        return self.cublas
+    
+    def _require_handle(self) -> ctypes.c_void_p:
+        if self.handle is None:
+            raise RuntimeError("cuBLAS handle is not created")
+        
+        return self.handle
+    
+    def _check_status(self, status: int, operation: str) -> None:
+        error = _check(status, operation)
+        if error is not None:
+            raise RuntimeError(error)
+    
+    
+    def _create_handle(self) -> ctypes.c_void_p:
+        cublas = self._require_cublas()
+
+        handle = ctypes.c_void_p()
+        status = cublas.cublasCreate_v2(ctypes.byref(handle))
+        self._check_status(status, "cublasCreate_v2")
+
+        if handle.value is None:
+            raise RuntimeError("cublasCreate_v2 returned NULL handle")
+        
+        
+        return handle
+    
+    def _destroy_handle(self, handle: ctypes.c_void_p) -> None:
+        cublas = self._require_cublas()
+
+        status = cublas.cublasDestroy_v2(handle)
+        self._check_status(status, "cublasDestroy_v2")
+
+    def get_version(self) -> int:
+        cublas = self._require_cublas()
+        handle = self._require_handle()
+
+        version = ctypes.c_int()
+        status = cublas.cublasGetVersion_v2(handle, ctypes.byref(version))
+        self._check_status(status, "cublasGetVersion_v2")
+
+        return int(version.value)
+    
+    def sgemm(
+            self,
+            *,
+            m: int,
+            n: int,
+            k: int,
+            alpha: float,
+            a_ptr: ctypes.c_void_p,
+            lda: int,
+            b_ptr: ctypes.c_void_p,
+            ldb: int,
+            beta: float,
+            c_ptr: ctypes.c_void_p,
+            ldc: int,
+            transa: int = CUBLAS_OP_N,
+            transb: int = CUBLAS_OP_N,
+    ) -> None:
+        cublas = self._require_cublas()
+        handle = self._require_handle()
+
+        alpha_value = ctypes.c_float(alpha)
+        beta_value = ctypes.c_float(beta)
+
+        status = cublas.cublasSgemm_v2(
+            handle,
+            transa,
+            transb,
+            m,
+            n,
+            k,
+            ctypes.byref(alpha_value),
+            a_ptr,
+            lda,
+            b_ptr,
+            ldb,
+            ctypes.byref(beta_value),
+            c_ptr,
+            ldc,
+        )
+
+        self._check_status(status, "cublasSgemm_v2")
+
+
+    def collect_info(self) -> CublasInfo:
+        version = self.get_version()
+
+        return CublasInfo(
+            available=True,
+            error=None,
+            library_path=self.library_path,
+            handle_created=self.handle is not None and self.handle.value is not None,
+            version=version,
+            version_text=_format_cublas_version(version),
+        )
+        
+
+
+def collect_cublas_info() -> CublasInfo:
+    try:
+        with CublasRuntime() as runtime:
+            return runtime.collect_info()
+        
+    except Exception as exc:
+        return CublasInfo(
+            available=False,
+            error=f"Unexpected cuBLAS runtime backend error: {exc}",
+            library_path=None,
+            handle_created=False,
+            version=None,
+            version_text=None
+        )
+    
+def collect_cublas_info_dict() -> dict[str, Any]:
+    return asdict(collect_cublas_info())

vramsuite/backends/cublaslt/__init__.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from vramsuite.backends.cublaslt.descriptors import (
+    MatmulDesc,
+    MatrixLayout,
+    MatmulPreference,
+)
+from vramsuite.backends.cublaslt.runtime import (
+    CublasLtInfo,
+    CublasLtRuntime,
+    collect_cublaslt_info,
+    collect_cublaslt_info_dict,
+)
+
+from vramsuite.backends.cublaslt.layouts import (
+    GemmShape,
+    MatrixLayoutSpec,
+    MatmulLayoutSpec,
+)
+from vramsuite.backends.cublaslt.workspace import DeviceWorkspace
+
+from vramsuite.backends.cublaslt.heuristics import (
+    RawMatmulHeuristic,
+    select_matmul_algorithm,
+)
+
+from vramsuite.backends.cublaslt.matmul import run_cublaslt_matmul
+
+__all__ = [
+    "CublasLtInfo",
+    "CublasLtRuntime",
+    "collect_cublaslt_info",
+    "collect_cublaslt_info_dict",
+    "MatmulDesc",
+    "MatrixLayout",
+    "MatmulPreference",
+    "GemmShape",
+    "MatrixLayoutSpec",
+    "MatmulLayoutSpec",
+    "DeviceWorkspace",
+    "RawMatmulHeuristic",
+    "select_matmul_algorithm",
+    "run_cublaslt_matmul",
+]

vramsuite/backends/cublaslt/descriptors.py

@@ -0,0 +1,243 @@
+from __future__ import annotations
+
+import ctypes
+from types import TracebackType
+
+from vramsuite.backends.cublaslt.enums import (
+    CUBLAS_COMPUTE_32F,
+    CUBLAS_OP_N,
+    CUBLASLT_MATMUL_DESC_TRANSA,
+    CUBLASLT_MATMUL_DESC_TRANSB,
+    CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+    CUDA_R_32F,
+)
+from vramsuite.backends.cublaslt.errors import check_cublas_status
+from vramsuite.backends.cublaslt.runtime import CublasLtRuntime
+
+
+
+class CublasLtDescriptor:
+    def __init__(self, runtime: CublasLtRuntime) -> None:
+        self.runtime = runtime
+        self.ptr: ctypes.c_void_p | None = None
+
+    def _require_cublaslt(self) -> ctypes.CDLL:
+        if self.runtime.cublaslt is None:
+            raise RuntimeError("cuBLASLt runtime is not loaded")
+        
+        return self.runtime.cublaslt
+    
+    def _require_ptr(self) -> ctypes.c_void_p:
+        if self.ptr is None or self.ptr.value is None:
+            raise RuntimeError(f"{self.__class__.__name__} is not created")
+        
+        return self.ptr
+    
+    def __enter__(self):
+        self.create()
+        return self
+    
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.destroy()
+
+
+class MatmulDesc(CublasLtDescriptor):
+    def __init__(
+        self,
+        runtime: CublasLtRuntime,
+        *,
+        compute_type: int = CUBLAS_COMPUTE_32F,
+        scale_type: int = CUDA_R_32F,
+        transa: int = CUBLAS_OP_N,
+        transb: int = CUBLAS_OP_N,
+    ) -> None:
+        super().__init__(runtime)
+        self.compute_type = compute_type
+        self.scale_type = scale_type
+        self.transa = transa
+        self.transb = transb
+
+    def create(self) -> None:
+        if self.ptr is not None:
+            return
+        
+        cublaslt = self._require_cublaslt()
+
+        desc = ctypes.c_void_p()
+        status = cublaslt.cublasLtMatmulDescCreate(
+            ctypes.byref(desc),
+            self.compute_type,
+            self.scale_type,
+        )
+        check_cublas_status(status, "cublasLtMatmulDescCreate")
+
+        if desc.value is None:
+            raise RuntimeError("cublasLtMatmulDescCreate returned NULL descriptor")
+
+        self.ptr = desc
+
+        try:
+            self.set_int32_attribute(CUBLASLT_MATMUL_DESC_TRANSA, self.transa)
+            self.set_int32_attribute(CUBLASLT_MATMUL_DESC_TRANSB, self.transb)
+        except Exception:
+            # __exit__ is not called when __enter__ raises, so destroy here.
+            self.destroy()
+            raise
+
+    def set_int32_attribute(self, attr: int, value: int) -> None:
+        cublaslt = self._require_cublaslt()
+        desc = self._require_ptr()
+
+        raw = ctypes.c_int32(value)
+        status = cublaslt.cublasLtMatmulDescSetAttribute(
+            desc,
+            attr,
+            ctypes.cast(ctypes.byref(raw), ctypes.c_void_p),
+            ctypes.sizeof(raw),
+        )
+        check_cublas_status(status, "cublasLtMatmulDescSetAttribute")
+
+    def destroy(self) -> None:
+        if self.ptr is None:
+            return
+        
+        cublaslt = self._require_cublaslt()
+
+        try:
+            status = cublaslt.cublasLtMatmulDescDestroy(self.ptr)
+            check_cublas_status(status, "cublasLtMatmulDescDestroy")
+        finally:
+            self.ptr = None
+
+class MatrixLayout(CublasLtDescriptor):
+    def __init__(
+        self,
+        runtime: CublasLtRuntime,
+        *,
+        data_type: int = CUDA_R_32F,
+        rows: int,
+        cols: int,
+        ld: int,
+    ) -> None:
+        super().__init__(runtime)
+        self.data_type = data_type
+        self.rows = rows
+        self.cols = cols
+        self.ld = ld
+
+    def create(self) -> None:
+        if self.ptr is not None:
+            return
+        
+        if self.rows <= 0:
+            raise ValueError("rows must be positive")
+        if self.cols <= 0:
+            raise ValueError("cols must be positive")
+        if self.ld <= 0:
+            raise ValueError("ld must be positive")
+        if self.ld < self.rows:
+            raise ValueError("ld must be >= rows for column-major layout")
+        
+        
+        cublaslt = self._require_cublaslt()
+
+        layout = ctypes.c_void_p()
+        status = cublaslt.cublasLtMatrixLayoutCreate(
+            ctypes.byref(layout),
+            self.data_type,
+            self.rows,
+            self.cols,
+            self.ld,
+        )
+        check_cublas_status(status, "cublasLtMatrixLayoutCreate")
+
+        if layout.value is None:
+            raise RuntimeError("cublasLtMatrixLayoutCreate returned NULL descriptor")
+        
+        self.ptr = layout
+
+    def destroy(self) -> None:
+        if self.ptr is None:
+            return
+        
+        cublaslt = self._require_cublaslt()
+
+        try:
+            status = cublaslt.cublasLtMatrixLayoutDestroy(self.ptr)
+            check_cublas_status(status, "cublasLtMatrixLayoutDestroy")
+        finally:
+            self.ptr = None
+
+
+class MatmulPreference(CublasLtDescriptor):
+    def __init__(
+        self,
+        runtime: CublasLtRuntime,
+        *,
+        max_workspace_bytes: int = 0,
+    ) -> None:
+        super().__init__(runtime)
+        self.max_workspace_bytes = max_workspace_bytes
+
+
+    def create(self) -> None:
+        if self.ptr is not None:
+            return
+        
+        if self.max_workspace_bytes < 0:
+            raise ValueError("max_workspace_bytes must be >= 0")
+        
+        cublaslt = self._require_cublaslt()
+
+        pref = ctypes.c_void_p()
+        status = cublaslt.cublasLtMatmulPreferenceCreate(ctypes.byref(pref))
+        check_cublas_status(status, "cublasLtMatmulPreferenceCreate")
+
+        if pref.value is None:
+            raise RuntimeError("cublasLtMatmulPreferenceCreate returned NULL descriptor")
+        
+        self.ptr = pref
+
+        try:
+            self.set_max_workspace_bytes(self.max_workspace_bytes)
+        except Exception:
+            # __exit__ is not called when __enter__ raises, so destroy here.
+            self.destroy()
+            raise
+
+    def destroy(self) -> None:
+        if self.ptr is None:
+            return
+        
+        cublaslt = self._require_cublaslt()
+
+        try:
+            status = cublaslt.cublasLtMatmulPreferenceDestroy(self.ptr)
+            check_cublas_status(status, "cublasLtMatmulPreferenceDestroy")
+        finally:
+            self.ptr = None
+
+    def set_max_workspace_bytes(self, value: int) -> None:
+        if value < 0:
+            raise ValueError("max workspace bytes must be >= 0")
+        
+        cublaslt = self._require_cublaslt()
+        pref = self._require_ptr()
+
+        workspace_value = ctypes.c_uint64(value)
+        workspace_ptr = ctypes.cast(ctypes.byref(workspace_value), ctypes.c_void_p)
+
+        status = cublaslt.cublasLtMatmulPreferenceSetAttribute(
+            pref,
+            CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+            workspace_ptr,
+            ctypes.sizeof(workspace_value),
+        )
+        check_cublas_status(status, "cublasLtMatmulPreferenceSetAttribute")
+
+        self.max_workspace_bytes = value

vramsuite/backends/cublaslt/enums.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+
+
+
+
+# cudaDataType_t
+CUDA_R_32F = 0
+CUDA_R_64F = 1
+CUDA_R_16F = 2
+CUDA_R_8I = 3
+CUDA_C_32F = 4
+CUDA_C_64F = 5
+CUDA_C_16F = 6
+CUDA_R_8U = 8
+CUDA_R_32I = 10
+CUDA_R_32U = 11
+CUDA_R_16BF = 14
+CUDA_R_8F_E4M3 = 28
+CUDA_R_8F_E5M2 = 29
+CUDA_R_8F_UE8M0 = 30
+CUDA_R_4F_E2M1 = 33
+
+
+# cublasComputeType_t
+CUBLAS_COMPUTE_16F = 64
+CUBLAS_COMPUTE_16F_PEDANTIC = 65
+CUBLAS_COMPUTE_32F = 68
+CUBLAS_COMPUTE_32F_PEDANTIC = 69
+CUBLAS_COMPUTE_64F = 70
+CUBLAS_COMPUTE_64F_PEDANTIC = 71
+CUBLAS_COMPUTE_32I = 72
+CUBLAS_COMPUTE_32I_PEDANTIC = 73
+CUBLAS_COMPUTE_32F_FAST_16F = 74
+CUBLAS_COMPUTE_32F_FAST_16BF = 75
+CUBLAS_COMPUTE_32F_FAST_TF32 = 77
+
+
+# cublasOperation_t
+CUBLAS_OP_N = 0
+CUBLAS_OP_T = 1
+CUBLAS_OP_C = 2
+
+
+# cublasLtMatmulDescAttributes_t
+CUBLASLT_MATMUL_DESC_TRANSA = 3
+CUBLASLT_MATMUL_DESC_TRANSB = 4
+
+
+# cublasLtMatmulPreferenceAttributes_t
+CUBLASLT_MATMUL_PREF_SEARCH_MODE = 0
+CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES = 1