| |
| """C05: explicit compute-kernel configs and per-module precision policies. |
| |
| Silent defaults this module exists to avoid (TT_PLATFORM.md section 0 item 7, PLAN.md section 0.2): |
| |
| - ``ttnn.matmul`` / ``ttnn.linear`` fall back to **LoFi** when a ``program_config`` or ``core_grid`` is given without |
| a ``compute_kernel_config``; |
| - ``ttnn.WormholeComputeKernelConfig()`` built without ``math_fidelity`` carries ``MathFidelity.Invalid``. |
| |
| So every op that takes a compute config gets one built here, with the fidelity spelled out. The default precision |
| is HiFi2 + fp32 accumulation, no approximations (the reference bundles' safe start; LoFi failed their gates almost |
| everywhere, RP section 2.9). A :class:`PrecisionPolicy` maps module names (globs, first match wins) to a |
| :class:`Precision`; ``<PREFIX>_PRECISION`` overrides rules at build time (the A/B switch), e.g. |
| ``CENTERPOINT_PRECISION="backbone.*=HiFi4+fp32;head.*=LoFi:w=bfp8"``. |
| |
| Other silent defaults to override by hand (not compute configs): ``ttnn.layer_norm`` epsilon 1e-12, SDPA |
| ``is_causal=True``, ``ttnn.embedding`` PADDED returning the cached pad row, fused HARDSWISH skipped in conv2d. |
| """ |
| from __future__ import annotations |
|
|
| import fnmatch |
| import os |
| from dataclasses import asdict, dataclass, replace |
| from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union |
|
|
| from .tensors import dtype_name |
|
|
| __all__ = ["FIDELITIES", "Precision", "PRESETS", "compute_kernel_config", "PrecisionPolicy"] |
|
|
| FIDELITIES = ("LoFi", "HiFi2", "HiFi3", "HiFi4") |
| _FIDELITY_KEY = {f.lower(): f for f in FIDELITIES} |
|
|
|
|
| def _fidelity(name: str) -> str: |
| key = _FIDELITY_KEY.get(str(name).strip().lower()) |
| if key is None: |
| raise ValueError(f"math fidelity {name!r}: expected one of {FIDELITIES}") |
| return key |
|
|
|
|
| def compute_kernel_config(fidelity: str = "HiFi2", *, fp32_acc: bool = True, approx: bool = False, |
| packer_l1_acc: bool = False, dst_full_sync: bool = False): |
| """A ``ttnn.WormholeComputeKernelConfig`` (the same class as ``BlackholeComputeKernelConfig``) with every |
| field explicit. ``fp32_acc`` = ``fp32_dest_acc_en`` (halves DST capacity: 4 tiles in half-sync).""" |
| import ttnn |
|
|
| return ttnn.WormholeComputeKernelConfig(math_fidelity=getattr(ttnn.MathFidelity, _fidelity(fidelity)), |
| math_approx_mode=bool(approx), fp32_dest_acc_en=bool(fp32_acc), |
| packer_l1_acc=bool(packer_l1_acc), dst_full_sync_en=bool(dst_full_sync)) |
|
|
|
|
| @dataclass(frozen=True) |
| class Precision: |
| """Fidelity / accumulation / dtype choice of one module (an op or a group of ops).""" |
|
|
| fidelity: str = "HiFi2" |
| fp32_acc: bool = True |
| approx: bool = False |
| packer_l1_acc: bool = False |
| dst_full_sync: bool = False |
| weights: str = "bfloat16" |
| activations: str = "bfloat16" |
|
|
| def __post_init__(self) -> None: |
| object.__setattr__(self, "fidelity", _fidelity(self.fidelity)) |
| object.__setattr__(self, "weights", dtype_name(self.weights)) |
| object.__setattr__(self, "activations", dtype_name(self.activations)) |
|
|
| @classmethod |
| def parse(cls, spec: Union[str, "Precision"]) -> "Precision": |
| """``"HiFi4+fp32"``, ``"LoFi"``, ``"HiFi2+fp32+approx+l1acc:w=bfp8:a=bf16"`` or a preset name |
| (``accurate`` / ``balanced`` / ``fast``). Flags: ``fp32`` (fp32 accumulation; absent = bf16 DST), |
| ``approx``, ``l1acc``, ``fullsync``; ``w=`` / ``a=`` set weight / activation dtypes.""" |
| if isinstance(spec, Precision): |
| return spec |
| text = spec.strip() |
| if text.lower() in PRESETS: |
| return PRESETS[text.lower()] |
| head, *opts = text.split(":") |
| fid, *flags = [p.strip() for p in head.split("+")] |
| kw: Dict[str, Any] = {"fidelity": fid, "fp32_acc": False} |
| for flag in (f.lower() for f in flags): |
| if flag == "fp32": |
| kw["fp32_acc"] = True |
| elif flag == "approx": |
| kw["approx"] = True |
| elif flag == "l1acc": |
| kw["packer_l1_acc"] = True |
| elif flag == "fullsync": |
| kw["dst_full_sync"] = True |
| else: |
| raise ValueError(f"precision {spec!r}: unknown flag {flag!r}") |
| for opt in opts: |
| k, _, v = opt.partition("=") |
| k = k.strip().lower() |
| if k in ("w", "weights"): |
| kw["weights"] = v.strip() |
| elif k in ("a", "act", "activations"): |
| kw["activations"] = v.strip() |
| else: |
| raise ValueError(f"precision {spec!r}: unknown option {k!r}") |
| return cls(**kw) |
|
|
| @property |
| def label(self) -> str: |
| """Round-trips through :meth:`parse`.""" |
| flags = "".join(f"+{f}" for f, on in (("fp32", self.fp32_acc), ("approx", self.approx), |
| ("l1acc", self.packer_l1_acc), ("fullsync", self.dst_full_sync)) if on) |
| return f"{self.fidelity}{flags}:w={self.weights}:a={self.activations}" |
|
|
| def with_(self, **changes: Any) -> "Precision": |
| return replace(self, **changes) |
|
|
| def compute_kernel_config(self): |
| return compute_kernel_config(self.fidelity, fp32_acc=self.fp32_acc, approx=self.approx, |
| packer_l1_acc=self.packer_l1_acc, dst_full_sync=self.dst_full_sync) |
|
|
| def weights_dtype(self): |
| from .tensors import ttnn_dtype |
|
|
| return ttnn_dtype(self.weights) |
|
|
| def activations_dtype(self): |
| from .tensors import ttnn_dtype |
|
|
| return ttnn_dtype(self.activations) |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| return asdict(self) |
|
|
|
|
| PRESETS: Dict[str, Precision] = { |
| "accurate": Precision("HiFi4", fp32_acc=True), |
| "balanced": Precision("HiFi2", fp32_acc=True), |
| "fast": Precision("LoFi", fp32_acc=False, weights="bfloat8_b"), |
| } |
|
|
| RuleSpec = Union[Mapping[str, Union[str, Precision]], Sequence[Tuple[str, Union[str, Precision]]]] |
|
|
|
|
| class PrecisionPolicy: |
| """Ordered ``pattern -> Precision`` rules (``fnmatch`` globs on dotted module names, first match wins) plus a |
| default. Typical use, once at model build:: |
| |
| POLICY = PrecisionPolicy({"backbone.*": "balanced", "head.reg*": "accurate"}, default="balanced") |
| policy = POLICY.with_env("CENTERPOINT") # <PREFIX>_PRECISION overrides, read once |
| cfg = policy.compute_kernel_config("backbone.block3.conv2") |
| w_dtype = policy.resolve("backbone.block3.conv2").weights_dtype() |
| |
| ``resolve`` records each module it answered for, so ``describe()`` shows the policy that reached the ops.""" |
|
|
| def __init__(self, rules: Optional[RuleSpec] = None, default: Union[str, Precision] = "balanced"): |
| items = rules.items() if isinstance(rules, Mapping) else (rules or ()) |
| self.rules: List[Tuple[str, Precision]] = [(str(p), Precision.parse(v)) for p, v in items] |
| self.default = Precision.parse(default) |
| self.used: Dict[str, str] = {} |
| self._configs: Dict[Precision, Any] = {} |
|
|
| def resolve(self, module: str) -> Precision: |
| for pattern, prec in self.rules: |
| if fnmatch.fnmatchcase(module, pattern): |
| self.used[module] = prec.label |
| return prec |
| self.used[module] = self.default.label |
| return self.default |
|
|
| def compute_kernel_config(self, module: str): |
| prec = self.resolve(module) |
| if prec not in self._configs: |
| self._configs[prec] = prec.compute_kernel_config() |
| return self._configs[prec] |
|
|
| def override(self, spec: str) -> "PrecisionPolicy": |
| """A new policy with ``"pattern=precision;pattern=precision"`` rules placed before the existing ones |
| (``*=...`` effectively replaces the default for unmatched modules).""" |
| extra = [] |
| for part in filter(None, (p.strip() for p in spec.split(";"))): |
| pattern, sep, prec = part.partition("=") |
| if not sep: |
| raise ValueError(f"precision override {part!r}: expected pattern=precision") |
| extra.append((pattern.strip(), Precision.parse(prec))) |
| return PrecisionPolicy(extra + self.rules, self.default) |
|
|
| def with_env(self, prefix: str, env: Optional[Mapping[str, str]] = None) -> "PrecisionPolicy": |
| """Apply ``<PREFIX>_PRECISION`` if set (else return ``self``). Call once at build.""" |
| spec = (os.environ if env is None else env).get(f"{prefix}_PRECISION", "").strip() |
| return self.override(spec) if spec else self |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"default": self.default.label, "rules": [[p, prec.label] for p, prec in self.rules], |
| "used": dict(sorted(self.used.items()))} |
|
|
