Function build_kernel

python/little_kernel/benchmark/utils.py:97–111

Build a LittleKernel kernel with standard passes.

(kernel_func, grid, block, shared_mem_bytes=0, arch=None, verbose=False)

def build_kernel(kernel_func, grid, block, shared_mem_bytes=0, arch=None, verbose=False):
    """Build a LittleKernel kernel with standard passes."""
    passes = PASSES["cuda"]
    if arch is None:
        cap = torch.cuda.get_device_capability()
        arch = f"sm_{cap[0]}{cap[1]}"
    return kernel_func.build(
        passes,
        codegen_cuda,
        grid=grid,
        block=block,
        shared_mem_bytes=shared_mem_bytes,
        arch=arch,
        verbose=verbose,
    )


def compile_kernel(kernel_func):