Add virtual memory support: PE_MMU, VA allocator, fabric MmuMapMsg
Implement VA/MMU layer (ADR-0011 Phase 1) enabling Triton kernels to use contiguous virtual addresses on sharded tensors. Key changes: - PE_MMU component: hybrid inbox (MmuMapMsg) + sync translate() for PE_DMA - VirtualAllocator + PEMemAllocator: free-list with coalescing - MmuMapMsg/MmuUnmapMsg fabric path with SIP-level routing - DPPolicy-based mapping: replicate=local, sharded=broadcast - Tensor lifecycle: del + weakref cleanup, context manager - Rename: TensorHandle.pa→addr, DmaReadCmd.src_pa→src_addr, ctx→torch Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -69,6 +69,7 @@ class TensorArgShard:
|
||||
class TensorArg:
|
||||
shards: tuple[TensorArgShard, ...]
|
||||
arg_kind: Literal["tensor"] = "tensor"
|
||||
va_base: int = 0 # VA base address for the entire tensor
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -121,3 +122,33 @@ class PeDmaMsg:
|
||||
nbytes: int
|
||||
is_write: bool = False
|
||||
msg_type: Literal["pe_dma"] = "pe_dma"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmuMapMsg:
|
||||
"""MMU mapping install: broadcast VA→PA entries to target PEs.
|
||||
|
||||
Sent via fabric: Host → PCIE_EP → IO_CPU → M_CPU → NOC → PE_MMU.
|
||||
target_sips controls which SIPs receive the message.
|
||||
"""
|
||||
|
||||
correlation_id: str
|
||||
request_id: str
|
||||
entries: tuple[dict, ...] # ({"va": int, "pa": int, "size": int}, ...)
|
||||
target_sips: tuple[int, ...] | Literal["all"] = "all"
|
||||
target_cubes: tuple[int, ...] | Literal["all"] = "all"
|
||||
target_pe: int | Literal["all"] = "all"
|
||||
msg_type: Literal["mmu_map"] = "mmu_map"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmuUnmapMsg:
|
||||
"""MMU mapping removal: broadcast VA ranges to unmap from all PEs."""
|
||||
|
||||
correlation_id: str
|
||||
request_id: str
|
||||
entries: tuple[dict, ...] # ({"va": int, "size": int}, ...)
|
||||
target_sips: tuple[int, ...] | Literal["all"] = "all"
|
||||
target_cubes: tuple[int, ...] | Literal["all"] = "all"
|
||||
target_pe: int | Literal["all"] = "all"
|
||||
msg_type: Literal["mmu_unmap"] = "mmu_unmap"
|
||||
|
||||
Reference in New Issue
Block a user