Add vmem chunked allocator

iris/drivers/__init__.py

@@ -2,22 +2,22 @@
 # Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
 
 """
-Shared driver package types for fabric backends.
+Shared driver package types for memory backends.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
 from typing import Optional
 
-from iris.drivers.base import BaseFabricDriver
+from iris.drivers.base import BaseDriver
 
 __all__ = ["DriverStack"]
 
 
 @dataclass
 class DriverStack:
-    """Fabric drivers available for a rank."""
+    """Driver available for a rank."""
 
     vendor: str
-    fabric: Optional[BaseFabricDriver]
+    driver: Optional[BaseDriver]

iris/drivers/base.py

@@ -2,21 +2,21 @@
 # Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
 
 """
-Abstract base classes, shared dataclasses, and exceptions for fabric drivers.
+Abstract base classes, shared dataclasses, and exceptions for memory drivers.
 """
 
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 
 from iris.host.distributed.topology import InterconnectLevel
 
 __all__ = [
     "PeerMapping",
     "LocalAllocation",
-    "BaseFabricDriver",
+    "BaseDriver",
     "DriverError",
     "DriverNotSupported",
 ]
@@ -40,6 +40,7 @@ class LocalAllocation:
     va: int
     size: int
     handle: Any
+    _va_owned: bool = True
 
 
 class DriverError(RuntimeError):
@@ -50,23 +51,23 @@ class DriverNotSupported(DriverError):
     """The current hardware or software stack does not support this driver."""
 
 
-class BaseFabricDriver(ABC):
-    """Cross-node fabric memory sharing (for example NVSwitch or xGMI)."""
+class BaseDriver(ABC):
+    """Generic base class for local and fabric memory drivers."""
 
     @abstractmethod
     def initialize(self, device_ordinal: int) -> None:
         """Prepare the driver for a specific local GPU."""
 
     @abstractmethod
-    def allocate_exportable(self, size: int) -> LocalAllocation:
-        """Allocate memory that can be shared through the fabric transport."""
+    def allocate_exportable(self, size: int, va: Optional[int] = None) -> LocalAllocation:
+        """Allocate exportable memory, optionally mapping it at a caller-reserved VA."""
 
     @abstractmethod
     def export_handle(self, allocation: LocalAllocation) -> bytes:
         """Export a transport-specific handle for a local allocation."""
 
     @abstractmethod
-    def import_and_map(self, peer_rank: int, handle_bytes: bytes, size: int) -> PeerMapping:
+    def import_and_map(self, peer_rank: int, handle_bytes: bytes, size: int, va: Optional[int] = None) -> PeerMapping:
         """Import a peer handle and map it into the local virtual address space."""
 
     @abstractmethod
@@ -76,3 +77,27 @@ def cleanup_import(self, mapping: PeerMapping) -> None:
     @abstractmethod
     def cleanup_local(self, allocation: LocalAllocation) -> None:
         """Release a locally-exported allocation."""
+
+    @abstractmethod
+    def get_minimum_granularity(self) -> int:
+        """Minimum allocation granularity in bytes for this driver+device."""
+
+    @abstractmethod
+    def reserve_va(self, size: int, alignment: int = 0) -> int:
+        """Reserve a virtual address range without backing physical memory."""
+
+    @abstractmethod
+    def free_va(self, va: int, size: int) -> None:
+        """Free a VA range previously returned by reserve_va."""
+
+    def get_address_range(self, ptr: int) -> tuple[int, int]:
+        """Return the base VA and size of the allocation containing ptr."""
+        raise DriverNotSupported(
+            f"{type(self).__name__} does not support get_address_range"
+        )
+
+    def export_pointer_handle(self, ptr: int, size: int) -> bytes:
+        """Export a peer handle for an arbitrary device pointer."""
+        raise DriverNotSupported(
+            f"{type(self).__name__} does not support export_pointer_handle"
+        )

iris/drivers/fabric/amd.py

@@ -7,30 +7,50 @@
 
 from __future__ import annotations
 
-from iris.drivers.base import BaseFabricDriver, DriverNotSupported, LocalAllocation, PeerMapping
+from typing import Optional
+
+from iris.drivers.base import (
+    BaseDriver,
+    DriverNotSupported,
+    LocalAllocation,
+    PeerMapping,
+)
 
 __all__ = ["AmdFabricDriver"]
 
 _NOT_IMPLEMENTED_MESSAGE = "AMD fabric driver not yet implemented"
 
 
-class AmdFabricDriver(BaseFabricDriver):
+class AmdFabricDriver(BaseDriver):
     """AMD fabric driver placeholder."""
 
     def initialize(self, device_ordinal: int) -> None:
         raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
 
-    def allocate_exportable(self, size: int) -> LocalAllocation:
+    def allocate_exportable(
+        self, size: int, va: Optional[int] = None
+    ) -> LocalAllocation:
         raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
 
     def export_handle(self, allocation: LocalAllocation) -> bytes:
         raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
 
-    def import_and_map(self, peer_rank: int, handle_bytes: bytes, size: int) -> PeerMapping:
+    def import_and_map(
+        self, peer_rank: int, handle_bytes: bytes, size: int, va: Optional[int] = None
+    ) -> PeerMapping:
         raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
 
     def cleanup_import(self, mapping: PeerMapping) -> None:
         raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
 
     def cleanup_local(self, allocation: LocalAllocation) -> None:
         raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
+
+    def get_minimum_granularity(self) -> int:
+        raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
+
+    def reserve_va(self, size: int, alignment: int = 0) -> int:
+        raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)
+
+    def free_va(self, va: int, size: int) -> None:
+        raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE)