55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Generator
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import simpy
|
|
|
|
from kernbench.components.base import PeEngineBase
|
|
|
|
if TYPE_CHECKING:
|
|
from kernbench.common.pe_commands import PeInternalTxn
|
|
from kernbench.components.context import ComponentContext
|
|
from kernbench.topology.types import Node
|
|
|
|
|
|
class PeMathComponent(PeEngineBase):
|
|
"""PE_MATH: element-wise computation engine sharing accel_slot (ADR-0014 D4).
|
|
|
|
Uses a shared compute resource (PE_ACCEL capacity=1) that is mutually
|
|
exclusive with PE_GEMM within the same PE.
|
|
"""
|
|
|
|
def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
|
|
super().__init__(node, ctx)
|
|
self._accel: simpy.Resource | None = None
|
|
|
|
def init_resources(self, env: simpy.Environment) -> None:
|
|
resource_name = self.node.attrs.get("shared_resource")
|
|
if resource_name and self.ctx:
|
|
self._accel = self.ctx.get_shared_resource(
|
|
env, f"{self._pe_prefix}.{resource_name}"
|
|
)
|
|
|
|
def run(self, env: simpy.Environment, nbytes: int) -> Generator:
|
|
overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0))
|
|
yield env.timeout(overhead_ns)
|
|
|
|
def handle_command(self, env: simpy.Environment, pe_txn: PeInternalTxn) -> Generator:
|
|
if self._accel:
|
|
with self._accel.request() as req:
|
|
yield req
|
|
yield from self.run(env, 0)
|
|
else:
|
|
yield from self.run(env, 0)
|
|
pe_txn.done.succeed()
|
|
|
|
def _forward_txn(self, env: simpy.Environment, txn: Any) -> Generator:
|
|
"""Transaction forwarding with accel_slot acquisition."""
|
|
if self._accel:
|
|
with self._accel.request() as req:
|
|
yield req
|
|
yield from super()._forward_txn(env, txn)
|
|
else:
|
|
yield from super()._forward_txn(env, txn)
|