Darkweb007/cuda-kernels

1"""Registry of supported elementwise ops and their CUDA expression templates.2 3Every op is "elementwise" in the fusion sense: given already-computed scalar4inputs for a single element, it produces a single scalar output with no5cross-element dependency. That's exactly the property that makes fusion6valid without any data-flow analysis beyond "is every op elementwise".7 8Each entry maps op name -> (arity, expr_template). expr_template is a Python9str.format() template using placeholder names {a}, {b} (positional inputs)10and any declared scalar args (e.g. {alpha}).11"""12 13from __future__ import annotations14 15from dataclasses import dataclass16from typing import Callable, Dict17 18 19@dataclass(frozen=True)20class OpSpec:21    arity: int                 # number of tensor inputs22    template: str               # C expression template, uses {a}, {b}, ... and scalar arg names23    scalar_args: tuple = ()     # names of required scalar args, e.g. ("alpha",)24    reference: Callable = None  # optional pure-python reference impl for testing, fn(*args, **kwargs)25 26 27OP_REGISTRY: Dict[str, OpSpec] = {28    "add": OpSpec(2, "({a} + {b})"),29    "sub": OpSpec(2, "({a} - {b})"),30    "mul": OpSpec(2, "({a} * {b})"),31    "div": OpSpec(2, "({a} / {b})"),32    "neg": OpSpec(1, "(-{a})"),33    "relu": OpSpec(1, "fmaxf({a}, 0.0f)"),34    "sigmoid": OpSpec(1, "(1.0f / (1.0f + expf(-{a})))"),35    "tanh": OpSpec(1, "tanhf({a})"),36    "exp": OpSpec(1, "expf({a})"),37    "sqrt": OpSpec(1, "sqrtf({a})"),38    "scalar_mul": OpSpec(1, "({a} * {alpha}f)", scalar_args=("alpha",)),39    "scalar_add": OpSpec(1, "({a} + {alpha}f)", scalar_args=("alpha",)),40    # GELU (tanh approximation), same formula used in the LayerNorm+GELU41    # kernel in cuda-ml-kernels -- expressed here as a single fused op so the42    # fuser can treat it as one node, or you could decompose it into43    # mul/add/tanh nodes and let the fuser re-derive the same kernel.44    "gelu": OpSpec(45        1,46        "(0.5f * {a} * (1.0f + tanhf(0.7978845608f * ({a} + 0.044715f * {a} * {a} * {a}))))",47    ),48}49 50 51def is_elementwise(op_name: str) -> bool:52    return op_name in OP_REGISTRY53 54 55def op_arity(op_name: str) -> int:56    return OP_REGISTRY[op_name].arity57