src/pqc_gpu_driver/backends/base.py
| 1 | """GPU backend base class - wraps the host's GPU runtime API.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from abc import ABC, abstractmethod |
| 6 | |
| 7 | from pqc_gpu_driver.tensor import EncryptedTensor |
| 8 | |
| 9 | |
| 10 | class GPUBackend(ABC): |
| 11 | """Abstract base for GPU runtime backends. |
| 12 | |
| 13 | A backend's job is to move EncryptedTensors between CPU and GPU memory. |
| 14 | The framework does NOT decrypt on the backend side - decryption is the |
| 15 | session's responsibility. Backends only move bytes. |
| 16 | """ |
| 17 | |
| 18 | name: str = "" |
| 19 | device_type: str = "" # "cuda" | "rocm" | "in-memory" | ... |
| 20 | |
| 21 | @abstractmethod |
| 22 | def upload(self, tensor: EncryptedTensor) -> str: |
| 23 | """Copy an encrypted tensor to device memory. Returns device_handle.""" |
| 24 | |
| 25 | @abstractmethod |
| 26 | def download(self, device_handle: str) -> EncryptedTensor: |
| 27 | """Copy an encrypted tensor from device memory back to host.""" |
| 28 | |
| 29 | @abstractmethod |
| 30 | def free(self, device_handle: str) -> None: |
| 31 | """Free device memory associated with the handle.""" |
| 32 | |
| 33 | @abstractmethod |
| 34 | def device_info(self) -> dict: |
| 35 | """Return a dict describing the device (name, compute_capability, memory).""" |
| 36 | |