code/tt_diffusion_planner/ttaw/ops/__init__.py

# SPDX-License-Identifier: Apache-2.0
"""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