Add SIP-level tensor parallelism, component registry YAML, VA offset verification
- DPPolicy: 3-level (sip/cube/pe), unified naming (column_wise/row_wise) - PE_CPU: auto num_programs from cube shard count - context.launch(): per-SIP KernelLaunchMsg with local va_base + auto local shape - deploy_tensor: removed mmus param, MMU mapping is context-only responsibility - ComponentRegistry: YAML-based lazy loading (components.yaml), impls→builtin rename - VA offset bench + tests: 2D/1D, standard Triton kernel pattern Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
"""PE_MMU component: address translation unit.
|
||||
|
||||
Component role: receives MmuMapMsg/MmuUnmapMsg via inbox (independent of PE_CPU).
|
||||
Utility role: PE_DMA/PE_GEMM call mmu.translate() directly (no SimPy overhead).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import simpy
|
||||
|
||||
from kernbench.components.base import ComponentBase, ComponentRegistry
|
||||
from kernbench.policy.address.pe_mmu import PeMMU
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from kernbench.components.context import ComponentContext
|
||||
from kernbench.topology.types import Node
|
||||
|
||||
|
||||
class PeMmuComponent(ComponentBase):
|
||||
"""PE_MMU: per-PE virtual-to-physical address translation.
|
||||
|
||||
Receives MmuMapMsg/MmuUnmapMsg via inbox and updates the internal
|
||||
page table. PE_DMA and PE_GEMM access the underlying PeMMU object
|
||||
via the ``mmu`` property for synchronous VA→PA translation.
|
||||
"""
|
||||
|
||||
def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
|
||||
super().__init__(node, ctx)
|
||||
page_size = int(node.attrs.get("page_size", 2 * 1024 * 1024))
|
||||
overhead_ns = float(node.attrs.get("tlb_overhead_ns", 0.0))
|
||||
self._mmu = PeMMU(page_size=page_size, overhead_ns=overhead_ns)
|
||||
|
||||
@property
|
||||
def mmu(self) -> PeMMU:
|
||||
"""The underlying PeMMU utility object for direct translate() calls."""
|
||||
return self._mmu
|
||||
|
||||
def run(self, env: simpy.Environment, nbytes: int) -> Generator:
|
||||
yield env.timeout(0)
|
||||
|
||||
def _worker(self, env: simpy.Environment) -> Generator:
|
||||
"""Process MmuMapMsg/MmuUnmapMsg from inbox."""
|
||||
from kernbench.runtime_api.kernel import MmuMapMsg, MmuUnmapMsg
|
||||
|
||||
while True:
|
||||
txn: Any = yield self._inbox.get()
|
||||
|
||||
if hasattr(txn, "request"):
|
||||
request = txn.request
|
||||
if isinstance(request, MmuMapMsg):
|
||||
for entry in request.entries:
|
||||
self._mmu.map(
|
||||
va=entry["va"], pa=entry["pa"], size=entry["size"],
|
||||
)
|
||||
txn.done.succeed()
|
||||
elif isinstance(request, MmuUnmapMsg):
|
||||
for entry in request.entries:
|
||||
self._mmu.unmap(va=entry["va"], size=entry["size"])
|
||||
txn.done.succeed()
|
||||
else:
|
||||
# Forward non-MMU transactions normally
|
||||
yield from self._forward_txn(env, txn)
|
||||
else:
|
||||
yield from self._forward_txn(env, txn)
|
||||
Reference in New Issue
Block a user