tirx-kernels 0.0.0__py3-none-any.whl

tirx_kernels/__init__.py

@@ -0,0 +1,54 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""TIRX kernel library.
+
+Private or experimental kernels can be layered on top of the released package by
+setting ``TIRX_KERNELS_OVERLAY_PATHS`` to one or more ``tirx_kernels`` package
+directories, separated by ``os.pathsep``.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+_OVERLAY_ENV = "TIRX_KERNELS_OVERLAY_PATHS"
+
+
+def _iter_overlay_paths() -> list[Path]:
+    paths = []
+    raw = os.environ.get(_OVERLAY_ENV, "")
+    for item in raw.split(os.pathsep):
+        item = item.strip()
+        if not item:
+            continue
+        path = Path(item).expanduser()
+        if path.is_dir():
+            paths.append(path.resolve())
+    return paths
+
+
+def _append_overlay_paths() -> None:
+    seen = {str(Path(item).resolve()) for item in __path__}
+    for path in _iter_overlay_paths():
+        text = str(path)
+        if text not in seen:
+            __path__.append(text)
+            seen.add(text)
+
+
+_append_overlay_paths()

tirx_kernels/_protocol.py

@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Standard kernel interface protocol.
+
+Every kernel module under ``kernels/<category>/`` that wants to be
+discoverable by the registry must expose:
+
+Module-level constants
+----------------------
+KERNEL_META : dict
+    Required keys:
+    - "name" (str): unique kernel name used by CLI (e.g. "rmsnorm")
+    - "category" (str): one of gemm, gemm_comm, attention, normalization, activation, ssm, loss, moe
+    - "compute_capability" (int): minimum SM version (e.g. 10 for sm100a)
+
+CONFIGS : list[dict]
+    Each dict has a "label" key (str) plus arbitrary kernel-specific
+    parameters.  The union of tir test parametrize configs and bench-ci
+    benchmark configs.
+
+BENCH_CONFIGS : list[dict]   (optional)
+    Benchmark-only configs.  If present, ``python -m tirx_kernels.bench``
+    uses these instead of ``CONFIGS`` so expensive benchmark sweeps do not
+    automatically become pytest correctness cases.
+
+Functions
+---------
+get_kernel(**cfg) -> tvm.tirx.PrimFunc | list[tvm.tirx.PrimFunc]
+    Return the TIR PrimFunc(s) for this kernel.  Multi-kernel workloads
+    (e.g. split-k GEMM with a separate reduce kernel) return a list.
+
+prepare_data(**cfg) -> dict[str, Any]
+    Prepare input/output tensors.  Returns a dict mapping argument names
+    to tensors (torch.Tensor or numpy.ndarray).
+
+check_correctness(outputs: dict, **cfg) -> None
+    Validate kernel outputs against a reference.
+    Raise AssertionError on mismatch.
+
+get_baselines(**cfg) -> dict[str, Callable]   (optional)
+    Return {name: callable} for baseline implementations used in
+    benchmarking (e.g. cublas, flashinfer).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class KernelModule(Protocol):
+    """Structural type that a kernel module must satisfy."""
+
+    KERNEL_META: dict[str, Any]
+    CONFIGS: list[dict[str, Any]]
+
+    @staticmethod
+    def get_kernel(**kwargs: Any) -> Any: ...
+
+    @staticmethod
+    def prepare_data(**kwargs: Any) -> dict[str, Any]: ...
+
+    @staticmethod
+    def check_correctness(outputs: dict[str, Any], **kwargs: Any) -> None: ...