ywkang
63669f82cb
- 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>
67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
"""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)
|