Darkweb007/cuda-kernels

1"""Generate CUDA source for a single FusionGroup.2 3Each fusion group becomes one `__global__` kernel with a grid-stride loop.4External inputs are read once from global memory into a register per5element; every intermediate value in the group lives in a local (register)6variable; only the group's single external output is written back to global7memory. That's the entire point of fusion: N ops -> 1 kernel launch, 1 read8per external input, 1 write.9"""10 11from __future__ import annotations12 13from typing import Dict, List14 15from .fuser import FusionGroup16from .ops import OP_REGISTRY17 18 19def _c_identifier(name: str) -> str:20    """Sanitize a graph value name into a valid C identifier."""21    safe = "".join(c if (c.isalnum() or c == "_") else "_" for c in name)22    if safe and safe[0].isdigit():23        safe = "_" + safe24    return safe25 26 27def generate_cuda_source(group: FusionGroup, kernel_name: str = "fused_kernel") -> str:28    inputs = group.inputs29    output = group.output30    produced_in_group = {n.output for n in group.nodes}31 32    def value_expr(name: str) -> str:33        """How to refer to `name` inside the loop body: a local var if it's34        an intermediate produced within this group, else an indexed global35        memory read."""36        if name in produced_in_group:37            return _c_identifier(name)38        return f"{_c_identifier(name)}[i]"39 40    lines: List[str] = []41    for node in group.nodes:42        spec = OP_REGISTRY[node.op]43        fmt_args: Dict[str, str] = {}44        arg_letters = ["a", "b", "c", "d"]45        for letter, input_name in zip(arg_letters, node.inputs):46            fmt_args[letter] = value_expr(input_name)47        for scalar_name, scalar_val in node.scalar_args.items():48            fmt_args[scalar_name] = str(scalar_val)49 50        expr = spec.template.format(**fmt_args)51        var_name = _c_identifier(node.output)52        lines.append(f"        float {var_name} = {expr};")53 54    body = "\n".join(lines)55    output_var = _c_identifier(output)56 57    params = ", ".join(f"const float* __restrict__ {_c_identifier(name)}" for name in inputs)58    if params:59        params += ", "60 61    source = f"""\62// Auto-generated by fusion_compiler.codegen -- do not hand-edit.63// Fuses {len(group.nodes)} elementwise op(s) into a single kernel:64//   {' ; '.join(repr(n) for n in group.nodes)}65extern "C" __global__ void {kernel_name}(66    {params}float* __restrict__ {output_var}_out, int n)67{{68    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {{69{body}70        {output_var}_out[i] = {output_var};71    }}72}}73"""74    return source75 76 77def generate_all(groups: List[FusionGroup], name_prefix: str = "fused_kernel") -> Dict[str, str]:78    """Generate CUDA source for every group, keyed by kernel name."""79    result = {}80    for idx, group in enumerate(groups):81        name = f"{name_prefix}_{idx}"82        result[name] = generate_cuda_source(group, kernel_name=name)83    return result84