| |
| """Custom-op helpers. Kernel ``.cpp`` / ``.hpp`` files are package data under ``ops/kernels/`` (shipped through |
| ``source.extra_code`` and ``[tool.setuptools.package-data]``), and are always referenced by absolute path, because |
| tt-metal resolves a relative kernel path against the CWD, ``TT_METAL_KERNEL_PATH`` and ``TT_METAL_HOME`` first |
| (TT_PLATFORM.md section 5.1). The C17-C23 op modules and the K1-K11 kernels of PLAN.md section 1.3 are added here |
| by the first port that needs them, with a numpy oracle test and the hang protocol of PLAN.md section 4.4. |
| |
| Op modules (import them explicitly; this package imports none of them): ``conv`` (C17 conv builders), ``upsample`` |
| (C18 up-sampling and interpolation), ``attention`` (C20 SDPA wrapper), ``gather`` (C19 row gathers), ``topk`` (C21 |
| top-k selection), ``heatmap`` (C22 heatmap local max), ``deform`` (C23 grid_sample helpers), ``segment`` (K1 |
| ``segment_reduce``: the first ``generic_op`` kernel, ``kernels/segment_reduce_dm.cpp``). |
| """ |
| from __future__ import annotations |
|
|
| import os |
|
|
| __all__ = ["KERNEL_DIR", "kernel_path"] |
|
|
| KERNEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels") |
|
|
|
|
| def kernel_path(name: str, base: str = KERNEL_DIR) -> str: |
| """Absolute path of a kernel source shipped with the package (``kernel_path("segment_max_reader.cpp")``). |
| Raises ``FileNotFoundError`` early instead of tt-metal's "doesn't exist in any of the searched paths".""" |
| path = os.path.join(base, name) |
| if not os.path.isfile(path): |
| raise FileNotFoundError(f"kernel source {path} is missing (not shipped as package data?)") |
| return path |
|
|
