Update torch_vgpu.py

torch_vgpu.py

@@ -1,90 +1,104 @@
-"""
-Custom PyTorch device implementation that routes operations through our virtual GPU.
-"""
-import torch
-from typing import Optional, Union, Tuple
-import numpy as np
-from virtual_vram import VirtualVRAM
-
-class VGPUStorage(torch.Storage):
-    """Custom storage class that uses our virtual VRAM"""
-    
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.vram = kwargs.get('vram')
-        if not self.vram:
-            from virtual_vram import VirtualVRAM
-            self.vram = VirtualVRAM()
-        self.tensor_id = kwargs.get('tensor_id', f"tensor_{id(self)}")
-    
-    def _new_shared(self, size):
-        return VGPUStorage(size, vram=self.vram)
-
-class VGPUTensor:
-    """Tensor implementation that uses vGPU for computations"""
-    @staticmethod
-    def __new__(cls, elem):
-        return torch.Tensor._make_subclass(cls, elem, elem.requires_grad)
-
-class VGPUDevice:
-    """
-    Custom PyTorch device implementation that routes operations through vGPU.
-    Usage:
-        vgpu = VGPUDevice()
-        tensor = torch.randn(2, 3).to(vgpu)
-    """
-    def __init__(self, vram: Optional[VirtualVRAM] = None):
-        self.vram = vram or VirtualVRAM()
-        self.tensor_cores = None  # Will be initialized when needed
-        
-    @property
-    def type(self):
-        return 'vgpu'
-        
-    def __str__(self):
-        return self.type
-        
-    def __repr__(self):
-        return self.type
-        
-    def _init_tensor_cores(self):
-        if self.tensor_cores is None:
-            from tensor_core import TensorCoreArray
-            self.tensor_cores = TensorCoreArray()
-            
-    def _to_vram(self, tensor: torch.Tensor) -> str:
-        """Store tensor data in virtual VRAM"""
-        tensor_id = f"tensor_{id(tensor)}"
-        data = tensor.detach().cpu().numpy()
-        self.vram.storage.store_tensor(tensor_id, data)
-        return tensor_id
-        
-    def _from_vram(self, tensor_id: str) -> torch.Tensor:
-        """Retrieve tensor data from virtual VRAM"""
-        data = self.vram.storage.load_tensor(tensor_id)
-        return torch.from_numpy(data)
-
-    def matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        """Matrix multiplication using tensor cores"""
-        self._init_tensor_cores()
-        
-        # Store inputs in VRAM
-        a_id = self._to_vram(a)
-        b_id = self._to_vram(b)
-        
-        # Perform matmul using tensor cores
-        result = self.tensor_cores.matmul(
-            self.vram.storage.load_tensor(a_id),
-            self.vram.storage.load_tensor(b_id)
-        )
-        
-        # Create new tensor with result
-        return torch.from_numpy(result)
-
-# Register vGPU device type with PyTorch
-torch.backends.register_custom_device("vgpu", VGPUDevice)
-
-def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
-    """Helper function to move tensors to vGPU"""
-    device = VGPUDevice(vram)
-    return tensor.to(device=device)
+"""
+Custom PyTorch device implementation that routes operations through our virtual GPU.
+"""
+import torch
+from typing import Optional, Union, Tuple
+import numpy as np
+from virtual_vram import VirtualVRAM
+
+class VGPUStorage(torch.Storage):
+    """Custom storage class that uses our virtual VRAM"""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.vram = kwargs.get('vram')
+        if not self.vram:
+            from virtual_vram import VirtualVRAM
+            self.vram = VirtualVRAM()
+        self.tensor_id = kwargs.get('tensor_id', f"tensor_{id(self)}")
+    
+    def _new_shared(self, size):
+        return VGPUStorage(size, vram=self.vram)
+
+class VGPUTensor:
+    """Tensor implementation that uses vGPU for computations"""
+    @staticmethod
+    def __new__(cls, elem):
+        return torch.Tensor._make_subclass(cls, elem, elem.requires_grad)
+
+class VGPUDevice:
+    """
+    Custom PyTorch device implementation that routes operations through vGPU.
+    Usage:
+        vgpu = VGPUDevice()
+        with vgpu.mode():
+            tensor = torch.randn(2, 3)  # Will be on vGPU
+    """
+    def __init__(self, vram: Optional[VirtualVRAM] = None):
+        self.vram = vram or VirtualVRAM()
+        self.tensor_cores = None  # Will be initialized when needed
+        self.device_name = "vgpu"
+        self._register_device()
+        
+    def _register_device(self):
+        """Register vGPU device using privateuse1 backend"""
+        try:
+            torch._C._dispatch._rename_privateuse1_backend(self.device_name)
+        except Exception as e:
+            raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
+    
+    @property
+    def type(self):
+        return self.device_name
+        
+    def __str__(self):
+        return f"{self.device_name}:0"
+        
+    def __repr__(self):
+        return f"{self.device_name}:0"
+        
+    def device(self):
+        """Get the PyTorch device object"""
+        return torch.device(str(self))
+    
+    def _init_tensor_cores(self):
+        if self.tensor_cores is None:
+            from tensor_core import TensorCoreArray
+            self.tensor_cores = TensorCoreArray()
+            
+    def _to_vram(self, tensor: torch.Tensor) -> str:
+        """Store tensor data in virtual VRAM"""
+        tensor_id = f"tensor_{id(tensor)}"
+        data = tensor.detach().cpu().numpy()
+        self.vram.storage.store_tensor(tensor_id, data)
+        return tensor_id
+        
+    def _from_vram(self, tensor_id: str) -> torch.Tensor:
+        """Retrieve tensor data from virtual VRAM"""
+        data = self.vram.storage.load_tensor(tensor_id)
+        return torch.from_numpy(data)
+
+    def matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Matrix multiplication using tensor cores"""
+        self._init_tensor_cores()
+        
+        # Store inputs in VRAM
+        a_id = self._to_vram(a)
+        b_id = self._to_vram(b)
+        
+        # Perform matmul using tensor cores
+        result = self.tensor_cores.matmul(
+            self.vram.storage.load_tensor(a_id),
+            self.vram.storage.load_tensor(b_id)
+        )
+        
+        # Create new tensor with result
+        return torch.from_numpy(result)
+
+# Register vGPU device type with PyTorch
+torch.backends.register_custom_device("vgpu", VGPUDevice)
+
+def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
+    """Helper function to move tensors to vGPU"""
+    device = VGPUDevice(vram)
+    return tensor.to(device=device)