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:
@@ -140,16 +140,65 @@ class ComponentRegistry:
|
||||
|
||||
Resolution order for ComponentRegistry.create(node, overrides, ctx):
|
||||
1. overrides[node.impl] — caller-injected override
|
||||
2. _registry[node.impl] — globally registered impl
|
||||
2. _registry[node.impl] — globally registered impl (lazy import)
|
||||
3. Error — no fallback; every node must have an impl
|
||||
|
||||
Registry is populated from components.yaml via load_components_yaml().
|
||||
Manual register() is still supported for tests and overrides.
|
||||
"""
|
||||
|
||||
_registry: dict[str, type[ComponentBase]] = {}
|
||||
_lazy: dict[str, str] = {} # impl → "module.path:ClassName"
|
||||
_loaded: bool = False
|
||||
|
||||
@classmethod
|
||||
def register(cls, impl: str, component_cls: type[ComponentBase]) -> None:
|
||||
cls._registry[impl] = component_cls
|
||||
|
||||
@classmethod
|
||||
def load_components_yaml(cls, path: str | None = None) -> None:
|
||||
"""Load impl→class mappings from components.yaml. Lazy imports on first use."""
|
||||
if cls._loaded:
|
||||
return
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
|
||||
if path is None:
|
||||
# Search: project root (cwd), then relative to this file
|
||||
candidates = [
|
||||
Path.cwd() / "components.yaml",
|
||||
Path(__file__).parent.parent.parent.parent / "components.yaml",
|
||||
]
|
||||
for p in candidates:
|
||||
if p.exists():
|
||||
path = str(p)
|
||||
break
|
||||
if path is None:
|
||||
return
|
||||
|
||||
with open(path) as f:
|
||||
spec = yaml.safe_load(f)
|
||||
for impl, class_path in (spec.get("components") or {}).items():
|
||||
cls._lazy[impl] = class_path
|
||||
cls._loaded = True
|
||||
|
||||
@classmethod
|
||||
def _resolve(cls, impl: str) -> type[ComponentBase] | None:
|
||||
"""Resolve impl name: check _registry first, then lazy import from _lazy."""
|
||||
if impl in cls._registry:
|
||||
return cls._registry[impl]
|
||||
if not cls._loaded:
|
||||
cls.load_components_yaml()
|
||||
class_path = cls._lazy.get(impl)
|
||||
if class_path is None:
|
||||
return None
|
||||
import importlib
|
||||
module_path, class_name = class_path.rsplit(":", 1)
|
||||
mod = importlib.import_module(module_path)
|
||||
component_cls = getattr(mod, class_name)
|
||||
cls._registry[impl] = component_cls # cache for next lookup
|
||||
return component_cls
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
@@ -159,9 +208,10 @@ class ComponentRegistry:
|
||||
) -> ComponentBase:
|
||||
if overrides and node.impl in overrides:
|
||||
return overrides[node.impl](node, ctx)
|
||||
if node.impl in cls._registry:
|
||||
return cls._registry[node.impl](node, ctx)
|
||||
component_cls = cls._resolve(node.impl)
|
||||
if component_cls is not None:
|
||||
return component_cls(node, ctx)
|
||||
raise ValueError(
|
||||
f"No component registered for impl '{node.impl}' (node: {node.id}). "
|
||||
f"Register it in kernbench.components.impls.__init__."
|
||||
f"Add it to components.yaml or call ComponentRegistry.register()."
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user