Update torch_vgpu.py

torch_vgpu.py

@@ -15,44 +15,38 @@ def init_vgpu_backend():
     global VGPU_BACKEND_INITIALIZED
     try:
         if not VGPU_BACKEND_INITIALIZED:
-            # Create an aten library implementation for our backend with PrivateUse1 dispatch key
-            lib = Library("aten", "IMPL", "PrivateUse1")
+            # First define our core library
+            lib = Library("vgpu", "DEF")
+            lib.define("custom_allocate(Device? device) -> Tensor")
+            lib.define("custom_to_cpu(Tensor self) -> Tensor")
+            lib.define("custom_from_cpu(Tensor self) -> Tensor")
             
-            # Register essential operations for the backend
-            @impl(lib, "empty.memory_format")
-            def empty_impl(size, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None):
-                # Create empty tensor on CPU first, will be moved to vGPU storage later
-                return torch.empty(size, dtype=dtype, device='cpu')
+            # Then implement the operations
+            impl_lib = Library("vgpu", "IMPL", "PrivateUse1")
             
-            @impl(lib, "copy.from")
-            def copy_impl(self, src, non_blocking=False):
-                # Handle tensor copying between devices
-                return torch.tensor(src.cpu().numpy(), device='cpu')
+            @impl(impl_lib, "custom_allocate")
+            def custom_allocate(device=None):
+                return torch.empty((), device='cpu')
             
-            @impl(lib, "_to_copy")
-            def to_impl(self, dtype=None, device=None, non_blocking=False, copy=False):
-                # Handle tensor device transfer
-                if device and str(device).startswith("privateuse1"):
-                    return self.cpu().clone()
-                return self.clone()
-                
-            # Finally rename privateuse1 to vgpu after implementing the ops
-            torch.utils.rename_privateuse1_backend("vgpu")
+            @impl(impl_lib, "custom_to_cpu")
+            def custom_to_cpu(tensor):
+                return tensor.clone()
             
-            # Generate all methods for our backend
-            torch.utils.generate_methods_for_privateuse1_backend(
-                for_tensor=True,
-                for_module=True,
-                for_packed_sequence=True,
-                for_storage=True
-            )
+            @impl(impl_lib, "custom_from_cpu")
+            def custom_from_cpu(tensor):
+                return tensor.clone()
+                
+            # Register our device type
+            torch._C._register_device_type("vgpu", 10)  # Use a custom type ID
             
+            # Mark initialization as complete
             VGPU_BACKEND_INITIALIZED = True
             
         return VGPU_BACKEND_INITIALIZED
     except Exception as e:
         print(f"Backend initialization warning: {e}")
         return False
+              
 
 class VGPUStorage(torch.Storage):
     """Custom storage class that uses our virtual VRAM"""
@@ -87,9 +81,8 @@ class VGPUDevice:
     def __init__(self, vram: Optional[VirtualVRAM] = None):
         self.vram = vram or VirtualVRAM()
         self.tensor_cores = None  # Will be initialized when needed
-        self.device_name = "privateuseone"  # Use original name first
+        self.device_name = "vgpu"  # Our registered device type
         self._register_device()
-        self.device_name = "vgpu"  # Then switch to renamed backend
 
     def _register_device(self):
         """Register vGPU device using PyTorch's device system"""
@@ -97,8 +90,8 @@ class VGPUDevice:
             if not VGPU_BACKEND_INITIALIZED:
                 raise RuntimeError("VGPU backend not properly initialized")
                 
-            # Create device with explicit index - use privateuse1 since that's the actual device type
-            self._device = torch.device("privateuse1")
+            # Create device using our registered device type
+            self._device = torch.device("vgpu")
             
             # Store this instance for reuse
             VGPUDevice._VGPU_INSTANCES[self.device_name] = self