wafer-core 0.1.0

Core utilities and environments for Wafer GPU kernel optimization

wafer-core

Core utilities and environments for Wafer GPU kernel optimization.

Features

• CUDA Compilation Tools: Compile CUDA files to PTX/SASS, detect GPU architecture
• Tool Finding: Locate CUDA toolkit executables (nvcc, nvdisasm)
• Telemetry: Decorator-based telemetry system with hooks
• Environments: GPU kernel optimization environments
• Remote Execution: SSH and Modal-based remote GPU execution
• Sessions: Agent session management
• Logging: Structured logging with color and JSON formatters

Installation

pip install wafer-core

Note: For full functionality including environments and sessions, you may also need to install rollouts:

pip install git+https://github.com/arb8020/research.git#subdirectory=rollouts

The core tools (CUDA compilation, tool finding, telemetry) work without rollouts.

Quick Start

from wafer_core import compile_cuda, detect_arch, find_nvcc

# Detect GPU architecture
arch_info = detect_arch()
print(f"Detected architecture: {arch_info['arch']}")

# Compile CUDA file
result = compile_cuda("kernel.cu", arch="sm_90")
print(f"PTX: {result.ptx}")
print(f"SASS: {result.sass}")

# Find nvcc
nvcc_path = find_nvcc()
if nvcc_path:
    print(f"Found nvcc at: {nvcc_path}")

Documentation

See the Wafer documentation for more information.

License

MIT License