diff --git a/.github/workflows/c-cpp.yml b/.github/workflows/c-cpp.yml index c27995b..16d5416 100644 --- a/.github/workflows/c-cpp.yml +++ b/.github/workflows/c-cpp.yml @@ -163,3 +163,19 @@ jobs: run: make clean && make asan=1 -j2 debug=1 - name: asan run: export DEVICE=cpu && make test + macos_cpu: + name: macOS Apple Silicon (CPU) + runs-on: macos-14 + steps: + - uses: actions/checkout@v2 + - name: build + run: make -j3 debug=1 + - name: test + run: make test + macos_metal: + name: macOS Apple Silicon (Metal) + runs-on: macos-14 + steps: + - uses: actions/checkout@v2 + - name: build metal=1 + run: make -j3 metal=1 debug=1 diff --git a/Makefile b/Makefile index 7902f9d..af28366 100755 --- a/Makefile +++ b/Makefile @@ -15,7 +15,6 @@ STATICLIB = $(BUILD_DIR)/libopenfish.a OBJ = $(BUILD_DIR)/misc.o \ $(BUILD_DIR)/error.o \ $(BUILD_DIR)/decode_cpu.o \ - $(BUILD_DIR)/nn_cpu.o \ $(BUILD_DIR)/openfish.o \ $(BUILD_DIR)/beam_search.o \ @@ -34,7 +33,7 @@ endif ifdef cuda CUDA_ROOT ?= /usr/local/cuda CUDA_LIB ?= $(CUDA_ROOT)/lib64 - CUDA_OBJ += $(BUILD_DIR)/decode_cuda.o $(BUILD_DIR)/nn_cuda.o + CUDA_OBJ += $(BUILD_DIR)/decode_cuda.o NVCC ?= $(CUDA_ROOT)/bin/nvcc CUDA_CFLAGS += -g -O2 -lineinfo $(CUDA_ARCH) -Xcompiler -Wall CUDA_LDFLAGS = -L$(CUDA_LIB) -lcudart_static -lrt -ldl @@ -50,12 +49,23 @@ else ifdef rocm # ifneq (,$(findstring gfx1150,$(ROCM_ARCH))) # ROCM_CFLAGS += -D__AMDGCN_WAVEFRONT_SIZE=32 # endif - ROCM_OBJ += $(BUILD_DIR)/decode_hip.o $(BUILD_DIR)/nn_hip.o + ROCM_OBJ += $(BUILD_DIR)/decode_hip.o GPU_LIB = $(BUILD_DIR)/hip_code.a ROCM_LDFLAGS = -L$(ROCM_LIB) -lamdhip64 -lrt -ldl CPPFLAGS += -DHAVE_ROCM=1 MAIN_CC = $(HIPCC) MAIN_CFLAGS = -x hip $(ROCM_CFLAGS) -fPIC +else ifdef metal + # Apple Silicon GPU backend. Objective-C++ glue is built with Apple clang (xcrun), + # and the .metal shaders are compiled at runtime via newLibraryWithSource:, so no + # offline metal toolchain (full Xcode) is required — only the Metal runtime framework. + METAL_CXX ?= xcrun clang++ + METAL_OBJ += $(BUILD_DIR)/decode_metal.o + GPU_LIB = $(BUILD_DIR)/metal_code.a + METAL_LDFLAGS = -framework Metal -framework Foundation -framework CoreFoundation -lc++ -lobjc + CPPFLAGS += -DHAVE_METAL=1 + MAIN_CC = $(CC) + MAIN_CFLAGS = $(CFLAGS) else GPU_LIB = $(BUILD_DIR)/cpu_decoy.a MAIN_CC = $(CC) @@ -68,13 +78,12 @@ endif ifdef debug CPPFLAGS += -DDEBUG=1 - CFLAGS += -fopenmp endif .PHONY: clean distclean test $(BINARY): $(BUILD_DIR)/main.o $(STATICLIB) - $(CC) $(CFLAGS) $(BUILD_DIR)/main.o $(STATICLIB) $(LDFLAGS) $(CUDA_LDFLAGS) $(ROCM_LDFLAGS) -o $@ + $(CC) $(CFLAGS) $(BUILD_DIR)/main.o $(STATICLIB) $(LDFLAGS) $(CUDA_LDFLAGS) $(ROCM_LDFLAGS) $(METAL_LDFLAGS) -o $@ $(STATICLIB): $(OBJ) $(GPU_LIB) cp $(GPU_LIB) $@ @@ -89,25 +98,21 @@ $(BUILD_DIR)/misc.o: src/misc.c src/misc.h $(BUILD_DIR)/error.o: src/error.c include/openfish/openfish_error.h $(CC) $(CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ -$(BUILD_DIR)/signal_prep.o: src/signal_prep.c - $(CC) $(CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ - $(BUILD_DIR)/decode_cpu.o: src/decode_cpu.c $(CC) $(CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ -$(BUILD_DIR)/nn_cpu.o: src/nn_cpu.c - $(CC) $(CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ - $(BUILD_DIR)/beam_search.o: src/beam_search.c $(CC) $(CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ $(BUILD_DIR)/openfish.o: src/openfish.c include/openfish/openfish.h $(CC) $(CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ -# cpu decoy +# cpu decoy: an empty static archive that $(STATICLIB) copies as its base before +# adding $(OBJ). BSD/macOS ar cannot create an archive with no members, so write +# the archive magic directly — both GNU and BSD ar accept this as a valid empty archive. $(BUILD_DIR)/cpu_decoy.a: rm -f $@ - $(AR) -r $@ + printf '!\n' > $@ # cuda $(BUILD_DIR)/cuda.a: $(CUDA_OBJ) @@ -116,9 +121,6 @@ $(BUILD_DIR)/cuda.a: $(CUDA_OBJ) $(BUILD_DIR)/decode_cuda.o: src/decode_cuda.c $(NVCC) -x cu $(CUDA_CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ -$(BUILD_DIR)/nn_cuda.o: src/nn_cuda.c - $(NVCC) -x cu $(CUDA_CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -c $< -o $@ - # hip $(BUILD_DIR)/hip_code.a: $(ROCM_OBJ) $(HIPCC) $(ROCM_CFLAGS) --emit-static-lib -fPIC --hip-link $^ -o $@ @@ -126,8 +128,22 @@ $(BUILD_DIR)/hip_code.a: $(ROCM_OBJ) $(BUILD_DIR)/decode_hip.o: src/decode_hip.c $(HIPCC) -x hip $(ROCM_CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -fPIC -c $< -o $@ -$(BUILD_DIR)/nn_hip.o: src/nn_hip.c - $(HIPCC) -x hip $(ROCM_CFLAGS) $(CPPFLAGS) $(DEPFLAGS) -fPIC -c $< -o $@ +# metal +$(BUILD_DIR)/metal_code.a: $(METAL_OBJ) + $(AR) rcs $@ $^ + +# embed the shader source as a C string (compiled at runtime with newLibraryWithSource:). +# openfish_defs.h is prepended so the shader and the host code share one copy of the constants, +# structs and arg blocks (newLibraryWithSource: can't resolve local #includes, so we concatenate +# at build time). each line is wrapped as a C string literal with a trailing \n; adjacent literals concatenate. +$(BUILD_DIR)/openfish_metal_src.h: src/openfish_defs.h src/openfish_metal.metal + printf 'static const char OPENFISH_METAL_SRC[] =\n' > $@ + sed -e 's/\\/\\\\/g' -e 's/"/\\"/g' -e 's/^/"/' -e 's/$$/\\n"/' \ + src/openfish_defs.h src/openfish_metal.metal >> $@ + printf ';\n' >> $@ + +$(BUILD_DIR)/decode_metal.o: src/decode_metal.mm src/openfish_defs.h $(BUILD_DIR)/openfish_metal_src.h + $(METAL_CXX) -x objective-c++ -fobjc-arc -std=c++17 $(CFLAGS) $(CPPFLAGS) -I$(BUILD_DIR) $(DEPFLAGS) -c $< -o $@ # pull in auto-generated header dependencies (.d files emitted by -MMD) -include $(BUILD_DIR)/*.d diff --git a/README.md b/README.md index 2777482..57e68e8 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # openfish -*openfish* is a library for CRF-CTC beam-search decoding used in nanopore basecalling. It supports CPU, NVIDIA GPU (CUDA) and AMD GPU (ROCm/HIP). +*openfish* is a library for CRF-CTC beam-search decoding used in nanopore basecalling. It supports CPU, NVIDIA GPU (CUDA), AMD GPU (ROCm/HIP) and Apple Silicon GPU (Metal). The CPU implementation was adopted from the C++ beam-search implementation in [ONT Dorado](https://github.com/nanoporetech/dorado) (licensed under the [Oxford Nanopore Technologies PLC. Public License Version 1.0](https://github.com/nanoporetech/dorado/blob/release-v0.8/LICENCE.txt)) and re-written in C. GPU backends were then built on top of that C implementation. @@ -12,6 +12,7 @@ The CPU implementation was adopted from the C++ beam-search implementation in [O - [CPU-only](#cpu-only) - [NVIDIA GPU (CUDA)](#nvidia-gpu-cuda) - [AMD GPU (ROCm)](#amd-gpu-rocm) + - [Apple Silicon GPU (Metal)](#apple-silicon-gpu-metal) - [Optional build flags](#optional-build-flags) - [Usage](#usage) - [Integrating as a library](#integrating-as-a-library) @@ -63,6 +64,22 @@ To target a specific GPU architecture, pass the `hipcc` architecture flag: make rocm=1 ROCM_ARCH="--offload-arch=gfx90a" ``` +### Apple Silicon GPU (Metal) + +Requires macOS on Apple Silicon (M-series). Only the **Xcode Command Line Tools** are needed +(`xcode-select --install`) — the full Xcode / offline Metal toolchain is *not* required, because +the shaders in `src/openfish_metal.metal` are compiled at runtime via `newLibraryWithSource:`. + +```sh +make metal=1 +``` + +The Objective-C++ glue (`src/decode_metal.mm`) is compiled with Apple clang (`xcrun clang++`) and +linked against the `Metal` and `Foundation` frameworks. All buffers use unified (shared) memory, so +there are no explicit host/device copies. Override the compiler with `METAL_CXX=... make metal=1`. + +The CPU-only build (`make`) also runs natively on Apple Silicon. + ### Optional build flags | Flag | Description | @@ -94,6 +111,12 @@ gcc [OPTIONS] -I path/to/openfish/include your_program.c \ path/to/openfish/lib/libopenfish.a -lz -lm -lpthread \ -L/opt/rocm/lib -lamdhip64 -lrt -ldl -o your_program +# static linking (Metal build, Apple Silicon) +clang [OPTIONS] -I path/to/openfish/include your_program.c \ + path/to/openfish/lib/libopenfish.a -lz -lm -lpthread \ + -framework Metal -framework Foundation -framework CoreFoundation \ + -lc++ -lobjc -o your_program + ``` *path/to/openfish/* is the absolute or relative path to the cloned repository. @@ -133,6 +156,18 @@ scripts/gpu_quick_run.sh hac scripts/gpu_quick_run.sh sup ``` +**GPU (Metal, Apple Silicon):** + +Like the CUDA/ROCm paths, the Metal backend consumes float16 scores, so `scripts/gpu_quick_run.sh` +works unchanged: + +```sh +make metal=1 debug=1 +scripts/gpu_quick_run.sh fast +scripts/gpu_quick_run.sh hac +scripts/gpu_quick_run.sh sup +``` + We have provided some expected values on several GPUs [here](test/ref_test_vals/) ## Acknowledgements diff --git a/docs/api.md b/docs/api.md index 0b612dc..9568e61 100644 --- a/docs/api.md +++ b/docs/api.md @@ -88,7 +88,7 @@ free(qstring); ## GPU decoding -Requires a `cuda=1` or `rocm=1` build. Scores must be in device memory and are **float16** for the GPU path (CPU path is float32). +Requires a `cuda=1`, `rocm=1` or `metal=1` build. Scores must be in device memory and are **float16** for the GPU path (CPU path is float32). For the Metal backend a device buffer is an `MTLBuffer` (the `scores_TNC` handle returned by the harness upload helper); on Apple Silicon's unified memory the `gpubuf` result pointers alias the shared buffers directly. ```c // Allocate persistent GPU working buffers (once per n_timesteps/batch_size/state_len combination) @@ -111,28 +111,3 @@ openfish_gpubuf_free(gpubuf); ``` `openfish_gpubuf_size(n_timesteps, batch_size, state_len)` returns the total number of device bytes a `gpubuf` of those dimensions occupies (useful for VRAM accounting before `openfish_gpubuf_init`). - -## Rotary embeddings - -Used by transformer-based models; applied in place to `x`. The GPU variant requires a `cuda=1` / `rocm=1` build and device pointers. - -```c -// CPU -void openfish_rotary_emb_cpu( - void *x, // [batch, seq, heads, head_dim] (modified in place) - const void *sin_buf, - const void *cos_buf, - int batch_size, int seq_len, int n_heads, int head_dim, int rotary_half, - int stride_batch, int stride_seq, int stride_head, - int n_threads -); - -// GPU (all pointers are device pointers) -void openfish_rotary_emb_gpu( - void *x, // modified in place - const void *sin_gpu, - const void *cos_gpu, - int batch_size, int seq_len, int n_heads, int head_dim, int rotary_half, - int stride_batch, int stride_seq, int stride_head -); -``` diff --git a/include/openfish/openfish.h b/include/openfish/openfish.h index f303004..49597d2 100644 --- a/include/openfish/openfish.h +++ b/include/openfish/openfish.h @@ -46,40 +46,13 @@ void openfish_decode_cpu( char **qstring ); -// Rotary position embedding (rotate-half convention). -// -// `sincos_width` is the width of the sin/cos tables: sin_buf and cos_buf MUST be -// laid out as [seq_len, sincos_width]. The kernel rotates 2*sincos_width elements -// of each head, pairing element i with element i+sincos_width: -// out_i = x_i*cos_j - x_{i+sincos_width}*sin_j -// out_{i+sincos_width} = x_i*sin_j + x_{i+sincos_width}*cos_j (j = i, table col) -// so it rotates 2*sincos_width dims (require 2*sincos_width <= head_dim). -// -// NOTE: fused GEMM+rotary kernels elsewhere may instead expect FULL-width tables -// [seq_len, 2*sincos_width] (each column duplicated). Do not confuse the two — the -// table's second dimension must always equal the width the consumer documents. -void openfish_rotary_emb_cpu( - void *x, - const void *sin_buf, - const void *cos_buf, - int batch_size, - int seq_len, - int n_heads, - int head_dim, - int sincos_width, // width of sin/cos tables: [seq_len, sincos_width]; rotates 2*sincos_width dims - int stride_batch, - int stride_seq, - int stride_head, - int n_threads -); - size_t openfish_gpubuf_size( int n_timesteps, int batch_size, int state_len ); -#if defined(HAVE_CUDA) || defined(HAVE_ROCM) +#if defined(HAVE_CUDA) || defined(HAVE_ROCM) || defined(HAVE_METAL) void openfish_decode_gpu( int n_timesteps, @@ -104,89 +77,7 @@ void openfish_gpubuf_free( openfish_gpubuf_t *gpubuf ); -// See openfish_rotary_emb_cpu: sin_gpu/cos_gpu are [seq_len, sincos_width] (rotate-half), -// rotating 2*sincos_width dims per head (require 2*sincos_width <= head_dim). -void openfish_rotary_emb_gpu( - void *x, - const void *sin_gpu, - const void *cos_gpu, - int batch_size, - int seq_len, - int n_heads, - int head_dim, - int sincos_width, // width of sin/cos tables: [seq_len, sincos_width]; rotates 2*sincos_width dims - int stride_batch, - int stride_seq, - int stride_head -); - -void openfish_flstm_step_gpu( - const void* scratch, - const void* ih_t, - void* cell, - void* hh_next, - int batch_size, - int hidden_dim -); - -void openfish_silu_mul_gpu( - const void *in, - void *out, - int n_tokens, - int hidden_dim -); - -void openfish_rmsnorm_gpu( - const void* in, - const void* residual, - const void* weight, - void* out, - int n_tokens, - int hidden_dim, - float alpha, - float eps -); - -void openfish_rmsnorm_quant_int8_gpu( - const void* in, - const void* weight, - void* residual, - void* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -); - -void openfish_rmsnorm_quant_fp8_gpu( - const void* in, - const void* weight, - void* residual, - void* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -); - -void openfish_quant_fp8_gpu( - const void* in, /* f16 [n_tokens, hidden_dim] input */ - void* out, /* uint8[n_tokens, hidden_dim] fp8 E4M3FN output */ - void* scale, /* f32 [n_tokens] per-token scale output */ - int n_tokens, - int hidden_dim -); - -void openfish_dequant_fp8_transpose_gpu( - const void* in, /* fp8 [n_timesteps, batch_size, n_channels] in */ - void* out, /* f16 [batch_size, n_timesteps, n_channels] out (dequant × scale, transposed) */ - int n_timesteps, - int batch_size, - int n_channels, - float scale -); - -#endif // defined(HAVE_CUDA) || defined(HAVE_ROCM) +#endif // defined(HAVE_CUDA) || defined(HAVE_ROCM) || defined(HAVE_METAL) #ifdef __cplusplus } diff --git a/scripts/gpu_quick_run.sh b/scripts/gpu_quick_run.sh index ee4fd80..c31bc42 100755 --- a/scripts/gpu_quick_run.sh +++ b/scripts/gpu_quick_run.sh @@ -20,12 +20,25 @@ DOWNLOAD_TEST_DATA() { mkdir -p test tar_path=test/data.tgz - wget -O $tar_path ${DATA_URL} || rm -rf $tar_path ${testdir} + if command -v wget >/dev/null 2>&1; then + wget -O $tar_path ${DATA_URL} || rm -rf $tar_path ${testdir} + else + curl -L -o $tar_path ${DATA_URL} || rm -rf $tar_path ${testdir} + fi echo "Extracting. Please wait." tar -xf $tar_path -C test || rm -rf $tar_path ${testdir} rm -f $tar_path } +# portable timing wrapper: GNU time uses --verbose, BSD/macOS time uses -l; fall back to none. +if /usr/bin/time --verbose true >/dev/null 2>&1; then + TIME="/usr/bin/time --verbose" +elif /usr/bin/time -l true >/dev/null 2>&1; then + TIME="/usr/bin/time -l" +else + TIME="" +fi + if [ "$#" -ne 1 ]; then die "usage: ./quick_run.sh " @@ -65,7 +78,7 @@ SCORES=${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_scores_TNC_half.blob DOWNLOAD_TEST_DATA -OMP_NUM_THREADS=1 /usr/bin/time --verbose ./openfish ${SCORES} ${BATCH_SIZE} ${STATE_LEN} || die "tool failed" +OMP_NUM_THREADS=1 $TIME ./openfish ${SCORES} ${BATCH_SIZE} ${STATE_LEN} || die "tool failed" ./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_bwd_NTC.blob bwd_NTC.blob ./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_post_NTC.blob post_NTC.blob diff --git a/scripts/pawsey_gpu_quick_run.sh b/scripts/pawsey_gpu_quick_run.sh deleted file mode 100755 index 69ab552..0000000 --- a/scripts/pawsey_gpu_quick_run.sh +++ /dev/null @@ -1,55 +0,0 @@ -#!/bin/bash - -# make sure to build with debug=1 - -die() { - echo "$1" >&2 - echo - exit 1 -} - -if [ "$#" -ne 1 ]; then - die "usage: ./quick_run.sh " -fi - -if [ ! -f "compare_blob" ]; then - g++ -o compare_blob test/compare_blob.cpp -fi - -MODEL=$1 - -STATE_LEN=3 -BATCH_SIZE=1000 -TIMESTEPS=1666 -TENS_LEN=0 -INTENS_LEN=0 -if [ "$MODEL" = "fast" ]; then - BATCH_SIZE=1000 - STATE_LEN=3 - TENS_LEN=$(( BATCH_SIZE*(TIMESTEPS+1)*64 )) - INTENS_LEN=$(( BATCH_SIZE*(TIMESTEPS) )) -fi -if [ "$MODEL" = "hac" ]; then - BATCH_SIZE=400 - STATE_LEN=4 - TENS_LEN=$(( BATCH_SIZE*(TIMESTEPS+1)*256 )) - INTENS_LEN=$(( BATCH_SIZE*(TIMESTEPS) )) -fi -if [ "$MODEL" = "sup" ]; then - BATCH_SIZE=200 - STATE_LEN=5 - TENS_LEN=$(( BATCH_SIZE*(TIMESTEPS+1)*1024 )) - INTENS_LEN=$(( BATCH_SIZE*(TIMESTEPS) )) -fi - -DATA_DIR=/software/projects/pawsey1099/bonwon/slorado_test_data/blobs -SCORES=${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_scores_TNC_half.blob - -/usr/bin/time --verbose ./openfish ${SCORES} models/dna_r10.4.1_e8.2_400bps_${MODEL}@v4.2.0 ${BATCH_SIZE} ${STATE_LEN} || die "tool failed" - -./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_bwd_NTC.blob bwd_NTC.blob $TENS_LEN -# ./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_fwd_NTC.blob fwd_NTC.blob $TENS_LEN -./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_post_NTC.blob post_NTC.blob $TENS_LEN -./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_qual_data.blob qual_data.blob $(( 4*(INTENS_LEN) )) -./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_total_probs.blob total_probs.blob $INTENS_LEN -# ./compare_blob ${DATA_DIR}/${MODEL}_${BATCH_SIZE}c_base_probs.blob base_probs.blob $INTENS_LEN diff --git a/src/beam_search.c b/src/beam_search.c index 0472357..3f65c88 100644 --- a/src/beam_search.c +++ b/src/beam_search.c @@ -1,5 +1,5 @@ #include "beam_search.h" -#include "decode.h" +#include "openfish_defs.h" #include #include diff --git a/src/beam_search.h b/src/beam_search.h index d5edf15..915e543 100644 --- a/src/beam_search.h +++ b/src/beam_search.h @@ -1,7 +1,7 @@ #ifndef BEAMSEARCH_CPU_H #define BEAMSEARCH_CPU_H -#include "decode.h" +#include "openfish_defs.h" #include #include diff --git a/src/beam_search_cuda.h b/src/beam_search_cuda.h index b299086..ee84c7a 100644 --- a/src/beam_search_cuda.h +++ b/src/beam_search_cuda.h @@ -5,7 +5,7 @@ #include #include -#include "decode.h" +#include "openfish_defs.h" __device__ static __forceinline__ void swapf(float *a, float *b) { float temp = *a; @@ -51,7 +51,7 @@ __device__ static __forceinline__ float log_sum_exp(float x, float y) { } static __global__ void generate_sequence( - const beam_args_t args, + const beam_params_t args, const uint8_t *_moves, const state_t *_states, const float *_qual_data, @@ -138,7 +138,9 @@ __device__ static __forceinline__ uint32_t crc32c(uint32_t crc, uint32_t new_bit } static __global__ void beam_search( - const beam_args_t beam_args, + const beam_params_t beam_args, + const void *_scores_TNC, + const float *_bwd_NTC, state_t *_states, uint8_t *_moves, beam_element_t *_beam_vector, @@ -168,8 +170,8 @@ static __global__ void beam_search( const size_t scores_block_stride = batch_size * n_channels; const float log_beam_cut = (beam_cut > 0.0f) ? __logf(beam_cut) : FLT_MAX; - const half *scores_TNC = (const half *)beam_args.scores_TNC + chunk * (num_states * NUM_BASES); - const float *bwd_NTC = beam_args.bwd_NTC + chunk * num_states * (n_timesteps + 1); + const half *scores_TNC = (const half *)_scores_TNC + chunk * (num_states * NUM_BASES); + const float *bwd_NTC = _bwd_NTC + chunk * num_states * (n_timesteps + 1); state_t *states = _states + chunk * n_timesteps; uint8_t *moves = _moves + chunk * n_timesteps; @@ -639,7 +641,8 @@ static __global__ void beam_search( } static __global__ void compute_qual_data( - const beam_args_t beam_args, + const beam_params_t beam_args, + const float *_post_NTC, state_t *_states, float *_qual_data, const float posts_scale @@ -654,7 +657,7 @@ static __global__ void compute_qual_data( const size_t num_states = 1ull << beam_args.num_state_bits; const size_t num_state_bits = beam_args.num_state_bits; - const float *post_NTC = beam_args.post_NTC + chunk * num_states * (n_timesteps + 1); + const float *post_NTC = _post_NTC + chunk * num_states * (n_timesteps + 1); state_t *states = _states + chunk * n_timesteps; float *qual_data = _qual_data + chunk * (n_timesteps * NUM_BASES); diff --git a/src/beam_search_hip.h b/src/beam_search_hip.h index 3627f54..c98cf76 100644 --- a/src/beam_search_hip.h +++ b/src/beam_search_hip.h @@ -10,7 +10,7 @@ #include #include -#include "decode.h" +#include "openfish_defs.h" __device__ static __forceinline__ void swapf(float *a, float *b) { float temp = *a; @@ -56,7 +56,7 @@ __device__ static __forceinline__ float log_sum_exp(float x, float y) { } static __global__ void generate_sequence( - const beam_args_t args, + const beam_params_t args, const uint8_t *_moves, const state_t *_states, const float *_qual_data, @@ -143,7 +143,9 @@ __device__ static __forceinline__ uint32_t crc32c(uint32_t crc, uint32_t new_bit } static __global__ void beam_search( - const beam_args_t beam_args, + const beam_params_t beam_args, + const void *_scores_TNC, + const float *_bwd_NTC, state_t *_states, uint8_t *_moves, beam_element_t *_beam_vector, @@ -173,8 +175,8 @@ static __global__ void beam_search( const size_t scores_block_stride = batch_size * n_channels; const float log_beam_cut = (beam_cut > 0.0f) ? __logf(beam_cut) : FLT_MAX; - const half *scores_TNC = (const half *)beam_args.scores_TNC + chunk * (num_states * NUM_BASES); - const float *bwd_NTC = beam_args.bwd_NTC + chunk * num_states * (n_timesteps + 1); + const half *scores_TNC = (const half *)_scores_TNC + chunk * (num_states * NUM_BASES); + const float *bwd_NTC = _bwd_NTC + chunk * num_states * (n_timesteps + 1); state_t *states = _states + chunk * n_timesteps; uint8_t *moves = _moves + chunk * n_timesteps; @@ -645,7 +647,8 @@ static __global__ void beam_search( } static __global__ void compute_qual_data( - const beam_args_t args, + const beam_params_t args, + const float *_post_NTC, state_t *_states, float *_qual_data, const float posts_scale @@ -660,7 +663,7 @@ static __global__ void compute_qual_data( const size_t num_states = 1ull << args.num_state_bits; const size_t num_state_bits = args.num_state_bits; - const float *post_NTC = args.post_NTC + chunk * num_states * (n_timesteps + 1); + const float *post_NTC = _post_NTC + chunk * num_states * (n_timesteps + 1); state_t *states = _states + chunk * n_timesteps; float *qual_data = _qual_data + chunk * (n_timesteps * NUM_BASES); diff --git a/src/decode.h b/src/decode.h deleted file mode 100644 index f74e854..0000000 --- a/src/decode.h +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef DECODE_H -#define DECODE_H - -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -typedef int32_t state_t; - -typedef struct beam_element { - state_t state; - uint8_t prev_element_index; - bool stay; -} beam_element_t; - -typedef struct beam_front_element { - uint32_t hash; - state_t state; - uint8_t prev_element_index; - bool stay; -} beam_front_element_t; - -typedef struct beam_args { - const void *scores_TNC; - float *bwd_NTC; - float *post_NTC; - size_t n_timesteps; - size_t batch_size; - size_t n_channels; - int num_state_bits; -} beam_args_t; - -typedef struct scan_args { - const void *scores_in; - uint64_t num_states; - uint64_t n_timesteps; - uint64_t batch_size; - uint64_t n_channels; - float fixed_stay_score; -} scan_args_t; - -#define NUM_BASE_BITS (2) -#define NUM_BASES (4) -#define NUM_TRANSITIONS (NUM_BASES + 1) -#define MAX_BEAM_WIDTH (32) -#define HASH_PRESENT_BITS (1024) -#define HASH_PRESENT_MASK (HASH_PRESENT_BITS - 1) -#define MAX_STATES (1024) -#define MAX_BEAM_CANDIDATES (NUM_TRANSITIONS * MAX_BEAM_WIDTH) -#define CRC_SEED (0x12345678u) - -#ifdef __cplusplus -} -#endif - -#endif // DECODE_H \ No newline at end of file diff --git a/src/decode_cuda.c b/src/decode_cuda.c index 4e46919..a790b0e 100644 --- a/src/decode_cuda.c +++ b/src/decode_cuda.c @@ -1,6 +1,6 @@ #include #include -#include "decode.h" +#include "openfish_defs.h" #include "scan_cuda.h" #include "beam_search_cuda.h" #include "error.h" @@ -108,8 +108,7 @@ void openfish_decode_gpu( OPENFISH_LOG_TRACE("scores tensor dim: %d, %d, %d", n_timesteps, batch_size, n_channels); - scan_args_t scan_args = {0}; - scan_args.scores_in = scores_TNC; + scan_params_t scan_args = {0}; scan_args.n_timesteps = n_timesteps; scan_args.batch_size = batch_size; scan_args.n_channels = n_channels; @@ -137,10 +136,7 @@ void openfish_decode_gpu( const float q_shift = options->q_shift; const float beam_cut = options->beam_cut; - beam_args_t beam_args = {0}; - beam_args.scores_TNC = (const half *)scores_TNC; - beam_args.bwd_NTC = gpubuf->bwd_NTC; - beam_args.post_NTC = gpubuf->post_NTC; + beam_params_t beam_args = {0}; beam_args.n_timesteps = n_timesteps; beam_args.batch_size = batch_size; beam_args.n_channels = n_channels; @@ -151,7 +147,7 @@ void openfish_decode_gpu( // beam search OPENFISH_LOG_TRACE("%s", "bwd scan..."); - bwd_scan<<>>(scan_args, gpubuf->bwd_NTC); + bwd_scan<<>>(scan_args, scores_TNC, gpubuf->bwd_NTC); checkCudaError(); cudaDeviceSynchronize(); checkCudaError(); @@ -163,6 +159,8 @@ void openfish_decode_gpu( // dynamic shared memory holds the back-guide sort scratch (num_states floats) beam_search<<>>( beam_args, + scores_TNC, + gpubuf->bwd_NTC, (state_t *)gpubuf->states, gpubuf->moves, (beam_element_t *)gpubuf->beam_vector, @@ -175,7 +173,7 @@ void openfish_decode_gpu( checkCudaError(); OPENFISH_LOG_TRACE("%s", "fwd + post scan..."); - fwd_post_scan<<>>(scan_args, gpubuf->bwd_NTC, gpubuf->post_NTC); + fwd_post_scan<<>>(scan_args, scores_TNC, gpubuf->bwd_NTC, gpubuf->post_NTC); checkCudaError(); cudaDeviceSynchronize(); checkCudaError(); @@ -183,6 +181,7 @@ void openfish_decode_gpu( OPENFISH_LOG_TRACE("%s", "compute qual data..."); compute_qual_data<<>>( beam_args, + gpubuf->post_NTC, (state_t *)gpubuf->states, gpubuf->qual_data, 1.0f diff --git a/src/decode_hip.c b/src/decode_hip.c index 25aad6e..0b2bb85 100644 --- a/src/decode_hip.c +++ b/src/decode_hip.c @@ -1,5 +1,5 @@ #include -#include "decode.h" +#include "openfish_defs.h" #include "scan_hip.h" #include "beam_search_hip.h" #include "error.h" @@ -112,8 +112,7 @@ void openfish_decode_gpu( OPENFISH_LOG_TRACE("scores tensor dim: %d, %d, %d", n_timesteps, batch_size, n_channels); - scan_args_t scan_args = {0}; - scan_args.scores_in = scores_TNC; + scan_params_t scan_args = {0}; scan_args.n_timesteps = n_timesteps; scan_args.batch_size = batch_size; scan_args.n_channels = n_channels; @@ -141,10 +140,7 @@ void openfish_decode_gpu( const float q_shift = options->q_shift; const float beam_cut = options->beam_cut; - beam_args_t beam_args = {0}; - beam_args.scores_TNC = (const half *)scores_TNC; - beam_args.bwd_NTC = gpubuf->bwd_NTC; - beam_args.post_NTC = gpubuf->post_NTC; + beam_params_t beam_args = {0}; beam_args.n_timesteps = n_timesteps; beam_args.batch_size = batch_size; beam_args.n_channels = n_channels; @@ -155,7 +151,7 @@ void openfish_decode_gpu( // beam search OPENFISH_LOG_TRACE("%s", "bwd scan..."); - bwd_scan<<>>(scan_args, gpubuf->bwd_NTC); + bwd_scan<<>>(scan_args, scores_TNC, gpubuf->bwd_NTC); checkHipError(); ret = hipDeviceSynchronize(); checkHipError(); HIP_CHECK(ret); @@ -167,6 +163,8 @@ void openfish_decode_gpu( // dynamic shared memory holds the back-guide sort scratch (num_states floats) beam_search<<>>( beam_args, + scores_TNC, + gpubuf->bwd_NTC, (state_t *)gpubuf->states, gpubuf->moves, (beam_element_t *)gpubuf->beam_vector, @@ -179,7 +177,7 @@ void openfish_decode_gpu( checkHipError(); HIP_CHECK(ret); OPENFISH_LOG_TRACE("%s", "fwd + post scan..."); - fwd_post_scan<<>>(scan_args, gpubuf->bwd_NTC, gpubuf->post_NTC); + fwd_post_scan<<>>(scan_args, scores_TNC, gpubuf->bwd_NTC, gpubuf->post_NTC); checkHipError(); ret = hipDeviceSynchronize(); checkHipError(); HIP_CHECK(ret); @@ -187,6 +185,7 @@ void openfish_decode_gpu( OPENFISH_LOG_TRACE("%s", "compute qual data..."); compute_qual_data<<>>( beam_args, + gpubuf->post_NTC, (state_t *)gpubuf->states, gpubuf->qual_data, 1.0f diff --git a/src/decode_metal.mm b/src/decode_metal.mm new file mode 100644 index 0000000..1c6a351 --- /dev/null +++ b/src/decode_metal.mm @@ -0,0 +1,363 @@ +// Metal (Apple Silicon) GPU backend for CRF-CTC decoding. +// +// Mirrors decode_cuda.c / decode_hip.c: one threadgroup per chunk, five compute kernels +// dispatched in sequence (bwd_scan -> beam_search -> fwd_post_scan -> compute_qual_data -> +// generate_sequence). The MSL kernels live in openfish_metal.metal and are compiled at +// runtime with newLibraryWithSource: (no offline metal toolchain required). +// +// On Apple Silicon all buffers use MTLResourceStorageModeShared (unified memory), so there +// are no explicit host<->device copies: gpubuf's public raw pointers alias the shared buffer +// contents directly, and results are read straight out of those buffers. + +#import +#import + +extern "C" { +#include +#include +#include "error.h" +} + +#include +#include +#include + +// constants, state_t, beam structs and the scan_params_t / beam_params_t argument blocks, +// shared verbatim with the shader (see openfish_defs.h). +#include "openfish_defs.h" + +// MSL source string generated from openfish_metal.metal at build time (see Makefile). +#include "openfish_metal_src.h" + +// internal handle: the public struct MUST be the first member so a openfish_gpubuf_t* +// can be cast back to metal_gpubuf*. +struct metal_gpubuf { + openfish_gpubuf_t pub; + id bwd_NTC; + id post_NTC; + id moves; + id sequence; + id qstring; + id beam_vector; + id states; + id qual_data; + id base_probs; + id total_probs; + int n_timesteps; + int batch_size; + int state_len; +}; + +// -------------------------------------------------------------- global metal state + +static id g_device = nil; +static id g_queue = nil; +static id g_pso_bwd = nil; +static id g_pso_fwd = nil; +static id g_pso_beam = nil; +static id g_pso_qual = nil; +static id g_pso_gen = nil; +static bool g_init = false; + +static id make_pso(id lib, const char *name) { + id fn = [lib newFunctionWithName:[NSString stringWithUTF8String:name]]; + if (!fn) { + OPENFISH_ERROR("metal: kernel function '%s' not found", name); + exit(EXIT_FAILURE); + } + NSError *err = nil; + id pso = [g_device newComputePipelineStateWithFunction:fn error:&err]; + if (!pso) { + OPENFISH_ERROR("metal: pipeline for '%s' failed: %s", name, [[err localizedDescription] UTF8String]); + exit(EXIT_FAILURE); + } + return pso; +} + +static void ensure_metal_init(void) { + if (g_init) return; + + // sanity: the persistent buffers store these structs, so device layout must match host. + static_assert(sizeof(beam_element_t) == 8, "beam_element_t layout must match MSL"); + // setBytes: payloads must match the shader's scan_params_t / beam_params_t. + static_assert(sizeof(scan_params_t) == 40, "scan_params_t layout must match MSL"); + static_assert(sizeof(beam_params_t) == 32, "beam_params_t layout must match MSL"); + + g_device = MTLCreateSystemDefaultDevice(); + if (!g_device) { + OPENFISH_ERROR("%s", "metal: no system default GPU device"); + exit(EXIT_FAILURE); + } + g_queue = [g_device newCommandQueue]; + + NSError *err = nil; + NSString *src = [NSString stringWithUTF8String:OPENFISH_METAL_SRC]; + MTLCompileOptions *opts = [[MTLCompileOptions alloc] init]; + id lib = [g_device newLibraryWithSource:src options:opts error:&err]; + if (!lib) { + OPENFISH_ERROR("metal: shader compile failed: %s", [[err localizedDescription] UTF8String]); + exit(EXIT_FAILURE); + } + + g_pso_bwd = make_pso(lib, "bwd_scan"); + g_pso_fwd = make_pso(lib, "fwd_post_scan"); + g_pso_beam = make_pso(lib, "beam_search"); + g_pso_qual = make_pso(lib, "compute_qual_data"); + g_pso_gen = make_pso(lib, "generate_sequence"); + + OPENFISH_LOG_TRACE("metal: initialised on device %s", [[g_device name] UTF8String]); + g_init = true; +} + +static id new_shared_buffer(size_t bytes) { + id b = [g_device newBufferWithLength:bytes options:MTLResourceStorageModeShared]; + if (!b) { + OPENFISH_ERROR("metal: failed to allocate %zu byte buffer", bytes); + exit(EXIT_FAILURE); + } + return b; +} + +// -------------------------------------------------------------- gpubuf lifecycle + +extern "C" openfish_gpubuf_t *openfish_gpubuf_init( + int n_timesteps, + int batch_size, + int state_len +) { + ensure_metal_init(); + + metal_gpubuf *mg = new metal_gpubuf(); + mg->n_timesteps = n_timesteps; + mg->batch_size = batch_size; + mg->state_len = state_len; + + const int num_states = (int)pow(NUM_BASES, state_len); + + mg->bwd_NTC = new_shared_buffer(sizeof(float) * (size_t)batch_size * (n_timesteps + 1) * num_states); + mg->post_NTC = new_shared_buffer(sizeof(float) * (size_t)batch_size * (n_timesteps + 1) * num_states); + mg->moves = new_shared_buffer(sizeof(uint8_t) * (size_t)batch_size * n_timesteps); + mg->sequence = new_shared_buffer(sizeof(char) * (size_t)batch_size * n_timesteps); + mg->qstring = new_shared_buffer(sizeof(char) * (size_t)batch_size * n_timesteps); + mg->beam_vector = new_shared_buffer(sizeof(beam_element_t) * (size_t)batch_size * MAX_BEAM_WIDTH * (n_timesteps + 1)); + mg->states = new_shared_buffer(sizeof(state_t) * (size_t)batch_size * n_timesteps); + mg->qual_data = new_shared_buffer(sizeof(float) * (size_t)batch_size * n_timesteps * NUM_BASES); + mg->base_probs = new_shared_buffer(sizeof(float) * (size_t)batch_size * n_timesteps); + mg->total_probs = new_shared_buffer(sizeof(float) * (size_t)batch_size * n_timesteps); + + // public raw pointers alias the shared buffer contents (unified memory). + mg->pub.bwd_NTC = (float *)[mg->bwd_NTC contents]; + mg->pub.post_NTC = (float *)[mg->post_NTC contents]; + mg->pub.moves = (uint8_t *)[mg->moves contents]; + mg->pub.sequence = (char *)[mg->sequence contents]; + mg->pub.qstring = (char *)[mg->qstring contents]; + mg->pub.beam_vector = [mg->beam_vector contents]; + mg->pub.states = [mg->states contents]; + mg->pub.qual_data = (float *)[mg->qual_data contents]; + mg->pub.base_probs = (float *)[mg->base_probs contents]; + mg->pub.total_probs = (float *)[mg->total_probs contents]; + + return &mg->pub; +} + +extern "C" void openfish_gpubuf_free(openfish_gpubuf_t *gpubuf) { + if (!gpubuf) return; + metal_gpubuf *mg = (metal_gpubuf *)gpubuf; + delete mg; // ARC releases the id members +} + +// -------------------------------------------------------------- decode + +extern "C" void openfish_decode_gpu( + int n_timesteps, + int batch_size, + int n_channels, + const void *scores_TNC, + int state_len, + const openfish_opt_t *options, + const openfish_gpubuf_t *gpubuf, + uint8_t **moves, + char **sequence, + char **qstring +) { + ensure_metal_init(); + + metal_gpubuf *mg = (metal_gpubuf *)gpubuf; + id scores = (__bridge id)scores_TNC; + + const int num_states = (int)pow(NUM_BASES, state_len); + const int num_state_bits = (int)log2((double)num_states); + + const float fixed_stay_score = options->blank_score; + const float q_scale = options->q_scale; + const float q_shift = options->q_shift; + const float beam_cut = options->beam_cut; + const float score_scale = 1.0f; + const float posts_scale = 1.0f; + + OPENFISH_LOG_TRACE("scores tensor dim: %d, %d, %d", n_timesteps, batch_size, n_channels); + + // host result buffers (freed by the caller, matching the CUDA/CPU API contract) + *moves = (uint8_t *)malloc((size_t)batch_size * n_timesteps * sizeof(uint8_t)); + MALLOC_CHK(*moves); + *sequence = (char *)malloc((size_t)batch_size * n_timesteps * sizeof(char)); + MALLOC_CHK(*sequence); + *qstring = (char *)malloc((size_t)batch_size * n_timesteps * sizeof(char)); + MALLOC_CHK(*qstring); + + // zero the output buffers (trailing positions past seq_len must stay 0) + memset([mg->moves contents], 0, (size_t)batch_size * n_timesteps * sizeof(uint8_t)); + memset([mg->sequence contents], 0, (size_t)batch_size * n_timesteps * sizeof(char)); + memset([mg->qstring contents], 0, (size_t)batch_size * n_timesteps * sizeof(char)); + + scan_params_t scan_args = {0}; + scan_args.num_states = num_states; + scan_args.n_timesteps = n_timesteps; + scan_args.batch_size = batch_size; + scan_args.n_channels = n_channels; + scan_args.fixed_stay_score = fixed_stay_score; + + beam_params_t beam_args = {0}; + beam_args.n_timesteps = n_timesteps; + beam_args.batch_size = batch_size; + beam_args.n_channels = n_channels; + beam_args.num_state_bits = num_state_bits; + + // the compact_offsets prefix-sum view overlays cand_scratch in the beam_search kernel, + // so the int offsets must fit within the bool bloom-filter storage. + ASSERT(MAX_BEAM_CANDIDATES * sizeof(int) <= HASH_PRESENT_BITS * sizeof(bool)); + + @autoreleasepool { + id cb = [g_queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + + const MTLSize grid = MTLSizeMake(batch_size, 1, 1); + + // ---- bwd scan (1 thread per state) ---- + [enc setComputePipelineState:g_pso_bwd]; + [enc setBytes:&scan_args length:sizeof(scan_args) atIndex:0]; + [enc setBuffer:scores offset:0 atIndex:1]; + [enc setBuffer:mg->bwd_NTC offset:0 atIndex:2]; + [enc dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(num_states, 1, 1)]; + [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + + // ---- beam search (MAX_BEAM_WIDTH * NUM_BASES threads) ---- + [enc setComputePipelineState:g_pso_beam]; + [enc setBytes:&beam_args length:sizeof(beam_args) atIndex:0]; + [enc setBuffer:scores offset:0 atIndex:1]; + [enc setBuffer:mg->bwd_NTC offset:0 atIndex:2]; + [enc setBuffer:mg->states offset:0 atIndex:3]; + [enc setBuffer:mg->moves offset:0 atIndex:4]; + [enc setBuffer:mg->beam_vector offset:0 atIndex:5]; + [enc setBytes:&beam_cut length:sizeof(float) atIndex:6]; + [enc setBytes:&fixed_stay_score length:sizeof(float) atIndex:7]; + [enc setBytes:&score_scale length:sizeof(float) atIndex:8]; + [enc dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(MAX_BEAM_WIDTH * NUM_BASES, 1, 1)]; + [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + + // ---- fwd + post scan (1 thread per state) ---- + [enc setComputePipelineState:g_pso_fwd]; + [enc setBytes:&scan_args length:sizeof(scan_args) atIndex:0]; + [enc setBuffer:scores offset:0 atIndex:1]; + [enc setBuffer:mg->bwd_NTC offset:0 atIndex:2]; + [enc setBuffer:mg->post_NTC offset:0 atIndex:3]; + [enc dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(num_states, 1, 1)]; + [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + + // ---- compute qual data (1 thread per chunk) ---- + [enc setComputePipelineState:g_pso_qual]; + [enc setBytes:&beam_args length:sizeof(beam_args) atIndex:0]; + [enc setBuffer:mg->post_NTC offset:0 atIndex:1]; + [enc setBuffer:mg->states offset:0 atIndex:2]; + [enc setBuffer:mg->qual_data offset:0 atIndex:3]; + [enc setBytes:&posts_scale length:sizeof(float) atIndex:4]; + [enc dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + + // ---- generate sequence (1 thread per chunk) ---- + [enc setComputePipelineState:g_pso_gen]; + [enc setBytes:&beam_args length:sizeof(beam_args) atIndex:0]; + [enc setBuffer:mg->moves offset:0 atIndex:1]; + [enc setBuffer:mg->states offset:0 atIndex:2]; + [enc setBuffer:mg->qual_data offset:0 atIndex:3]; + [enc setBuffer:mg->base_probs offset:0 atIndex:4]; + [enc setBuffer:mg->total_probs offset:0 atIndex:5]; + [enc setBuffer:mg->sequence offset:0 atIndex:6]; + [enc setBuffer:mg->qstring offset:0 atIndex:7]; + [enc setBytes:&q_shift length:sizeof(float) atIndex:8]; + [enc setBytes:&q_scale length:sizeof(float) atIndex:9]; + [enc dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status == MTLCommandBufferStatusError) { + OPENFISH_ERROR("metal: command buffer error: %s", [[cb.error localizedDescription] UTF8String]); + exit(EXIT_FAILURE); + } + } + + // copy results out of the shared buffers into the caller-owned arrays + memcpy(*moves, [mg->moves contents], (size_t)batch_size * n_timesteps * sizeof(uint8_t)); + memcpy(*sequence, [mg->sequence contents], (size_t)batch_size * n_timesteps * sizeof(char)); + memcpy(*qstring, [mg->qstring contents], (size_t)batch_size * n_timesteps * sizeof(char)); +} + +// -------------------------------------------------------------- test-harness helpers (test_utils_metal.h) + +extern "C" void set_device_metal(int device) { + (void)device; // Metal uses the system default device; no per-index selection + ensure_metal_init(); +} + +// scores_TNC is host float16 [T,N,C] (matching the CUDA/ROCm GPU path); upload into a shared +// buffer. returns a +1 retained MTLBuffer bridged to void* (release with free_scores_metal). +extern "C" void *upload_scores_to_metal( + int n_timesteps, + int batch_size, + int n_channels, + const void *scores_TNC +) { + ensure_metal_init(); + const size_t bytes = (size_t)n_timesteps * batch_size * n_channels * sizeof(uint16_t); + id buf = new_shared_buffer(bytes); + memcpy([buf contents], scores_TNC, bytes); + return (void *)CFBridgingRetain(buf); +} + +extern "C" void free_scores_metal(void *scores_gpu) { + if (scores_gpu) { + CFBridgingRelease(scores_gpu); + } +} + +#ifdef DEBUG +extern "C" void write_gpubuf_metal( + uint64_t n_timesteps, + uint64_t batch_size, + int state_len, + const openfish_gpubuf_t *gpubuf +) { + const int num_states = (int)pow(NUM_BASES, state_len); + const size_t tens = (size_t)batch_size * (n_timesteps + 1) * num_states; + const size_t intens = (size_t)batch_size * n_timesteps; + + struct { const char *name; const void *data; size_t count; size_t elem; } blobs[] = { + {"bwd_NTC.blob", gpubuf->bwd_NTC, tens, sizeof(float)}, + {"post_NTC.blob", gpubuf->post_NTC, tens, sizeof(float)}, + {"qual_data.blob", gpubuf->qual_data, intens * NUM_BASES, sizeof(float)}, + {"base_probs.blob", gpubuf->base_probs, intens, sizeof(float)}, + {"total_probs.blob",gpubuf->total_probs,intens, sizeof(float)}, + }; + for (size_t b = 0; b < sizeof(blobs) / sizeof(blobs[0]); ++b) { + FILE *fp = fopen(blobs[b].name, "w"); + F_CHK(fp, blobs[b].name); + if (fwrite(blobs[b].data, blobs[b].elem, blobs[b].count, fp) != blobs[b].count) { + fprintf(stderr, "error writing %s: %s\n", blobs[b].name, strerror(errno)); + exit(EXIT_FAILURE); + } + fclose(fp); + } +} +#endif diff --git a/src/main.c b/src/main.c index 50d8c38..686bd87 100644 --- a/src/main.c +++ b/src/main.c @@ -7,7 +7,6 @@ #include #include #include -#include #if defined HAVE_CUDA #include "test_utils_cuda.h" @@ -17,6 +16,14 @@ #include "test_utils_hip.h" #endif +#if defined HAVE_METAL +#include "test_utils_metal.h" +#endif + +#if defined(HAVE_CUDA) || defined(HAVE_ROCM) || defined(HAVE_METAL) +#define HAVE_GPU 1 +#endif + int main(int argc, char* argv[]) { #if defined DEBUG if (argc < 4) { @@ -26,11 +33,13 @@ int main(int argc, char* argv[]) { } set_openfish_log_level(OPENFISH_LOG_DBUG); - const int device = omp_get_thread_num(); + const int device = 0; #if defined HAVE_CUDA set_device_cuda(device); #elif defined HAVE_ROCM set_device_hip(device); +#elif defined HAVE_METAL + set_device_metal(device); #endif OPENFISH_LOG_DEBUG("simulating batches on device %d", device); @@ -44,7 +53,7 @@ int main(int argc, char* argv[]) { // read scores from file size_t scores_len = n_timesteps * batch_size * n_channels; -#if defined HAVE_CUDA || defined HAVE_ROCM +#if defined HAVE_CUDA || defined HAVE_ROCM || defined HAVE_METAL const int elem_size = sizeof(uint16_t); #else const int elem_size = sizeof(float); @@ -67,9 +76,11 @@ int main(int argc, char* argv[]) { void *scores_gpu = upload_scores_to_cuda(n_timesteps, batch_size, n_channels, scores); #elif defined HAVE_ROCM void *scores_gpu = upload_scores_to_hip(n_timesteps, batch_size, n_channels, scores); +#elif defined HAVE_METAL + void *scores_gpu = upload_scores_to_metal(n_timesteps, batch_size, n_channels, scores); #endif -#if defined HAVE_CUDA || defined HAVE_ROCM +#if defined HAVE_GPU openfish_gpubuf_t *gpubuf = openfish_gpubuf_init(n_timesteps, batch_size, state_len); #endif openfish_opt_t options = openfish_decoder_default_opts(); @@ -103,7 +114,7 @@ int main(int argc, char* argv[]) { #endif // decode scores -#if defined HAVE_CUDA || defined HAVE_ROCM +#if defined HAVE_GPU openfish_decode_gpu(n_timesteps, batch_size, n_channels, scores_gpu, state_len, &options, gpubuf, &moves, &sequence, &qstring); #else int n_threads = 8; @@ -163,7 +174,11 @@ int main(int argc, char* argv[]) { write_gpubuf_hip(n_timesteps, batch_size, state_len, gpubuf); #endif -#if defined HAVE_CUDA || defined HAVE_ROCM +#if defined DEBUG && defined HAVE_METAL + write_gpubuf_metal(n_timesteps, batch_size, state_len, gpubuf); +#endif + +#if defined HAVE_GPU openfish_gpubuf_free(gpubuf); #endif @@ -171,6 +186,8 @@ int main(int argc, char* argv[]) { free_scores_cuda(scores_gpu); #elif defined HAVE_ROCM free_scores_hip(scores_gpu); +#elif defined HAVE_METAL + free_scores_metal(scores_gpu); #endif #endif return 0; diff --git a/src/nn_cpu.c b/src/nn_cpu.c deleted file mode 100644 index 92dba45..0000000 --- a/src/nn_cpu.c +++ /dev/null @@ -1,132 +0,0 @@ -#include "error.h" -#include "misc.h" - -#include -#include - -#include -#include - -static void rotary_emb( - float *x, - const float *_cos, - const float *_sin, - uint64_t seq_len, - uint64_t stride_batch, - uint64_t stride_seq, - uint64_t stride_head, - uint64_t sincos_width, - uint64_t batch, - uint64_t head, - uint64_t rot -) { - float *_o0 = x + (batch * stride_batch) + (head * stride_head) + rot; - float *_o1 = x + (batch * stride_batch) + (head * stride_head) + sincos_width + rot; - - for (int seq = 0; seq < seq_len; ++seq) { - float cos_val = *(_cos + (seq * sincos_width) + rot); - float sin_val = *(_sin + (seq * sincos_width) + rot); - - float *o0 = _o0 + (seq * stride_seq); - float *o1 = _o1 + (seq * stride_seq); - - float x0 = *o0; - float x1 = *o1; - - *o0 = x0 * cos_val - x1 * sin_val; - *o1 = x0 * sin_val + x1 * cos_val; - } -} - -typedef struct { - float *x; - const float *sin_buf; - const float *cos_buf; - uint64_t start; - uint64_t end; - uint64_t seq_len; - uint64_t n_heads; - uint64_t head_dim; - uint64_t sincos_width; - uint64_t stride_batch; - uint64_t stride_seq; - uint64_t stride_head; -} rotary_emb_thread_arg_t; - -static void* pthread_single_rotary_emb(void* voidargs) { - rotary_emb_thread_arg_t* args = (rotary_emb_thread_arg_t*)voidargs; - - for (uint64_t batch = args->start; batch < args->end; ++batch) { - for (uint64_t head = 0; head < args->n_heads; ++head) { - for (uint64_t rot = 0; rot < args->sincos_width; ++rot) { - rotary_emb( - args->x, - args->cos_buf, - args->sin_buf, - args->seq_len, - args->stride_batch, - args->stride_seq, - args->stride_head, - args->sincos_width, - batch, head, rot - ); - } - } - } - - pthread_exit(0); -} - -void openfish_rotary_emb_cpu( - void *x, - const void *sin_buf, - const void *cos_buf, - int batch_size, - int seq_len, - int n_heads, - int head_dim, - int sincos_width, - int stride_batch, - int stride_seq, - int stride_head, - int n_threads -) { - // create threads - n_threads = batch_size < n_threads ? batch_size : n_threads; - const int chunks_per_thread = batch_size / n_threads; - const int num_threads_with_one_more_chunk = batch_size % n_threads; - - OPENFISH_LOG_TRACE("dispatching %d threads for cpu decoding", n_threads); - - pthread_t tids[n_threads]; - rotary_emb_thread_arg_t pt_args[n_threads]; - int32_t t, ret; - - // set the data structures - for (t = 0; t < n_threads; t++) { - int extra = t < num_threads_with_one_more_chunk ? t : num_threads_with_one_more_chunk; - pt_args[t].start = t * chunks_per_thread + extra; - pt_args[t].end = pt_args[t].start + chunks_per_thread + (int)(t < num_threads_with_one_more_chunk); - pt_args[t].x = (float *)x; - pt_args[t].sin_buf = (const float *)sin_buf; - pt_args[t].cos_buf = (const float *)cos_buf; - pt_args[t].seq_len = seq_len; - pt_args[t].n_heads = n_heads; - pt_args[t].head_dim = head_dim; - pt_args[t].sincos_width = sincos_width; - pt_args[t].stride_batch = stride_batch; - pt_args[t].stride_seq = stride_seq; - pt_args[t].stride_head = stride_head; - } - - // score tensors - for (t = 0; t < n_threads; t++) { - ret = pthread_create(&tids[t], NULL, pthread_single_rotary_emb, (void *)(&pt_args[t])); - NEG_CHK(ret); - } - - for (t = 0; t < n_threads; t++) { - ret = pthread_join(tids[t], NULL); - NEG_CHK(ret); - } -} diff --git a/src/nn_cuda.c b/src/nn_cuda.c deleted file mode 100644 index b36735e..0000000 --- a/src/nn_cuda.c +++ /dev/null @@ -1,141 +0,0 @@ -#include -#include "error.h" -#include "cuda_utils.h" -#include "nn_kernel_cuda.h" - -#include - -#include -#include - -void openfish_rmsnorm_quant_int8_gpu( - const void* in, - const void* weight, - void* residual, - void* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - ASSERT(hidden_dim <= 1024); - ASSERT(hidden_dim % 2 == 0); // kernel is half2-vectorized: one thread per adjacent pair - - int threads = hidden_dim / 2; - int blocks = n_tokens; - - rmsnorm_quant_int8<<>>( - (half *)in, (half *)weight, (int8_t *)residual, (float *)residual_scale, n_tokens, hidden_dim, alpha, eps - ); - checkCudaError(); - cudaDeviceSynchronize(); - checkCudaError(); -} - -void openfish_rmsnorm_quant_fp8_gpu( - const void* in, - const void* weight, - void* residual, - void* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - (void)in; (void)weight; (void)residual; (void)residual_scale; - (void)n_tokens; (void)hidden_dim; (void)alpha; (void)eps; - OPENFISH_ERROR("%s", "fp8 fused rmsnorm not implemented for CUDA"); - exit(EXIT_FAILURE); -} - -void openfish_rmsnorm_gpu( - const void* in, - const void* residual, - const void* weight, - void* out, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - ASSERT(hidden_dim <= 1024); - ASSERT(hidden_dim % 2 == 0); // kernel is half2-vectorized: one thread per adjacent pair - - int threads = hidden_dim / 2; - int blocks = n_tokens; - - rmsnorm<<>>( - (half *)in, (half *)residual, (half *)weight, (half *)out, n_tokens, hidden_dim, alpha, eps - ); - checkCudaError(); - cudaDeviceSynchronize(); - checkCudaError(); -} - -void openfish_silu_mul_gpu( - const void *in, - void *out, - int n_tokens, - int hidden_dim -) { - int threads = 1024; - int blocks = n_tokens; - - silu_mul<<>>( - (const half *)in, - (half *)out, - hidden_dim, - n_tokens - ); - checkCudaError(); - cudaDeviceSynchronize(); - checkCudaError(); -} - -void openfish_flstm_step_gpu( - const void* scratch, - const void* ih_t, - void* cell, - void* hh_next, - int batch_size, int hidden_dim -) { - int threads = (hidden_dim < 1024) ? hidden_dim : 1024; - flstm_step<<>>( - (const half*)scratch, (const half*)ih_t, - (half*)cell, (half*)hh_next, - 4 * hidden_dim, hidden_dim - ); - checkCudaError(); -} - -void openfish_rotary_emb_gpu( - void *x, - const void *sin_gpu, - const void *cos_gpu, - int batch_size, - int seq_len, - int n_heads, - int head_dim, - int sincos_width, - int stride_batch, - int stride_seq, - int stride_head -) { - int thread_h = 32; - dim3 block_size(sincos_width, thread_h, 1); - dim3 grid_size(batch_size, n_heads, 1); - - rotary_emb<<>>( - (half *)x, - (const float *)cos_gpu, - (const float *)sin_gpu, - seq_len, - stride_batch, - stride_seq, - stride_head, - sincos_width - ); - checkCudaError(); - cudaDeviceSynchronize(); - checkCudaError(); -} diff --git a/src/nn_hip.c b/src/nn_hip.c deleted file mode 100644 index 77fd4dd..0000000 --- a/src/nn_hip.c +++ /dev/null @@ -1,193 +0,0 @@ -#include -#include "error.h" -#include "nn_kernel_hip.h" -#include "hip_utils.h" - -#include - -#include -#include - -void openfish_rmsnorm_quant_int8_gpu( - const void* in, - const void* weight, - void* residual, - void* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - hipError_t ret; - ASSERT(hidden_dim <= 1024); - ASSERT(hidden_dim % 2 == 0); // kernel is half2-vectorized: one thread per adjacent pair - - int threads = hidden_dim / 2; - int blocks = n_tokens; - - rmsnorm_quant_int8<<>>( - (half *)in, (half *)weight, (int8_t *)residual, (float *)residual_scale, n_tokens, hidden_dim, alpha, eps - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} - -void openfish_rmsnorm_quant_fp8_gpu( - const void* in, - const void* weight, - void* residual, - void* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - hipError_t ret; - ASSERT(hidden_dim <= 1024); - ASSERT(hidden_dim % 2 == 0); // kernel is half2-vectorized: one thread per adjacent pair - - int threads = hidden_dim / 2; - int blocks = n_tokens; - - rmsnorm_quant_fp8<<>>( - (half *)in, (half *)weight, (uint8_t *)residual, (float *)residual_scale, n_tokens, hidden_dim, alpha, eps - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} - -void openfish_quant_fp8_gpu( - const void* in, - void* out, - void* scale, - int n_tokens, - int hidden_dim -) { - hipError_t ret; - ASSERT(hidden_dim <= 1024); - - quant_fp8<<>>( - (const half *)in, (uint8_t *)out, (float *)scale, n_tokens, hidden_dim - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} - -void openfish_dequant_fp8_transpose_gpu( - const void* in, /* fp8 [n_timesteps, batch_size, n_channels] */ - void* out, /* f16 [batch_size, n_timesteps, n_channels] */ - int n_timesteps, - int batch_size, - int n_channels, - float scale -) { - hipError_t ret; - ASSERT(n_channels <= 1024); - - dequant_fp8_transpose<<>>( - (const uint8_t *)in, (half *)out, n_timesteps, batch_size, n_channels, scale - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} - -void openfish_rmsnorm_gpu( - const void* in, - const void* residual, - const void* weight, - void* out, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - hipError_t ret; - ASSERT(hidden_dim <= 1024); - ASSERT(hidden_dim % 2 == 0); // kernel is half2-vectorized: one thread per adjacent pair - - int threads = hidden_dim / 2; - int blocks = n_tokens; - - rmsnorm<<>>( - (half *)in, (half *)residual, (half *)weight, (half *)out, n_tokens, hidden_dim, alpha, eps - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} - -void openfish_silu_mul_gpu( - const void *in, - void *out, - int n_tokens, - int hidden_dim -) { - hipError_t ret; - - int threads = 1024; - int blocks = n_tokens; - - silu_mul<<>>( - (const half *)in, - (half *)out, - hidden_dim, - n_tokens - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} - -void openfish_flstm_step_gpu( - const void* scratch, - const void* ih_t, - void* cell, - void* hh_next, - int batch_size, int hidden_dim -) { - int threads = (hidden_dim < 1024) ? hidden_dim : 1024; - flstm_step<<>>( - (const half*)scratch, (const half*)ih_t, - (half*)cell, (half*)hh_next, - 4 * hidden_dim, hidden_dim - ); - checkHipError(); -} - -void openfish_rotary_emb_gpu( - void *x, - const void *sin_gpu, - const void *cos_gpu, - int batch_size, - int seq_len, - int n_heads, - int head_dim, - int sincos_width, - int stride_batch, - int stride_seq, - int stride_head -) { - hipError_t ret; - - int thread_h = 32; - dim3 block_size(sincos_width, thread_h, 1); - dim3 grid_size(batch_size, n_heads, 1); - - rotary_emb<<>>( - (half *)x, - (const float *)cos_gpu, - (const float *)sin_gpu, - seq_len, - stride_batch, - stride_seq, - stride_head, - sincos_width - ); - checkHipError(); - ret = hipDeviceSynchronize(); - checkHipError(); HIP_CHECK(ret); -} diff --git a/src/nn_kernel_cuda.h b/src/nn_kernel_cuda.h deleted file mode 100644 index 84ee122..0000000 --- a/src/nn_kernel_cuda.h +++ /dev/null @@ -1,264 +0,0 @@ -// The MIT License (MIT) - -// Copyright (c) 2025 Bonson Wong - -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: - -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. - -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -// THE SOFTWARE. - -#ifndef NN_KERNEL_CUDA_H -#define NN_KERNEL_CUDA_H - -#include -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -static __global__ void rotary_emb( - half *x, - const float *_cos, - const float *_sin, - const uint64_t seq_len, - const uint64_t stride_batch, - const uint64_t stride_seq, - const uint64_t stride_head, - const uint64_t sincos_width -) { - const uint64_t batch = blockIdx.x; - const uint64_t head = blockIdx.y; - const uint64_t rot = threadIdx.x; - const uint64_t tid = threadIdx.y; - const uint64_t n_threads = blockDim.y; - - if (tid >= seq_len) return; - - half *_o0 = x + (batch * stride_batch) + (head * stride_head) + rot; - half *_o1 = x + (batch * stride_batch) + (head * stride_head) + sincos_width + rot; - - for (int seq = tid; seq < seq_len; seq += n_threads) { - float cos = *(_cos + (seq * sincos_width) + rot); - float sin = *(_sin + (seq * sincos_width) + rot); - - half *o0 = _o0 + (seq * stride_seq); - half *o1 = _o1 + (seq * stride_seq); - - float x0 = __half2float(*o0); - float x1 = __half2float(*o1); - - *o0 = __float2half(x0 * cos - x1 * sin); - *o1 = __float2half(x0 * sin + x1 * cos); - } -} - -static __global__ void silu_mul( - const half *in, - half *out, - const uint64_t hidden_dim, - const uint64_t n_tokens -) { - uint64_t j = blockIdx.x; - - for (uint64_t k = threadIdx.x; k < hidden_dim; k += blockDim.x) { - uint64_t i = k + j * (hidden_dim * 2); - - half y = in[i]; - half gate = in[i + hidden_dim]; - - float g = __half2float(gate); - float silu = g / (1.0f + __expf(-g)); - - out[k + j * hidden_dim] = __float2half(silu * __half2float(y)); - } -} - -// ── Block-wide reductions ───────────────────────────────────────────────────── -// Warp-level primitives, then a block-level reduction that broadcasts the result -// to all threads via a caller-supplied shared scratch buffer (>= blockDim.x/warpSize -// floats). The trailing double-sync leaves the buffer safe to reuse for a second -// reduction in the same kernel. blockDim.x <= 1024 => num_warps <= 32 == warpSize, -// so the final reduction fits in a single warp. -static __device__ __forceinline__ float warp_reduce_sum(float v) { - for (int offset = warpSize / 2; offset > 0; offset >>= 1) - v += __shfl_down_sync(0xffffffff, v, offset); - return v; -} -static __device__ __forceinline__ float warp_reduce_max(float v) { - for (int offset = warpSize / 2; offset > 0; offset >>= 1) - v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset)); - return v; -} -static __device__ __forceinline__ float block_reduce_sum(float v, float* warp_reduce_buf) { - int warp_id = threadIdx.x / warpSize; - int lane_id = threadIdx.x % warpSize; - v = warp_reduce_sum(v); - if (lane_id == 0) warp_reduce_buf[warp_id] = v; - __syncthreads(); - int num_warps = (blockDim.x + warpSize - 1) / warpSize; - float r = (threadIdx.x < num_warps) ? warp_reduce_buf[threadIdx.x] : 0.0f; - if (warp_id == 0) r = warp_reduce_sum(r); - if (threadIdx.x == 0) warp_reduce_buf[0] = r; - __syncthreads(); - float result = warp_reduce_buf[0]; - __syncthreads(); - return result; -} -static __device__ __forceinline__ float block_reduce_max(float v, float* warp_reduce_buf) { - int warp_id = threadIdx.x / warpSize; - int lane_id = threadIdx.x % warpSize; - v = warp_reduce_max(v); - if (lane_id == 0) warp_reduce_buf[warp_id] = v; - __syncthreads(); - int num_warps = (blockDim.x + warpSize - 1) / warpSize; - float r = (threadIdx.x < num_warps) ? warp_reduce_buf[threadIdx.x] : 0.0f; - if (warp_id == 0) r = warp_reduce_max(r); - if (threadIdx.x == 0) warp_reduce_buf[0] = r; - __syncthreads(); - float result = warp_reduce_buf[0]; - __syncthreads(); - return result; -} - -static __global__ void rmsnorm( - const half* in, - const half* residual, - const half* weight, - half* out, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - int row = blockIdx.x; // Which sequence/batch element - - if (row >= n_tokens) return; - - // Vectorized half2: each thread owns one adjacent pair. blockDim.x == hidden_dim/2. - const half2* x = reinterpret_cast(in + row * hidden_dim); - const half2* res = reinterpret_cast(residual + row * hidden_dim); - const half2* w2 = reinterpret_cast(weight); - half2* y = reinterpret_cast(out + row * hidden_dim); - int hd2 = hidden_dim >> 1; - - __shared__ float warp_reduce_buf[32]; - - float thread_sum = 0.0f; - float2 v_new; // valid because blockDim.x == hidden_dim/2, so this loop runs exactly once per thread - for (int i = threadIdx.x; i < hd2; i += blockDim.x) { - float2 xf = __half22float2(x[i]); - float2 rf = __half22float2(res[i]); - float2 val = make_float2(xf.x + rf.x * alpha, xf.y + rf.y * alpha); - v_new = val; - thread_sum += val.x * val.x + val.y * val.y; - } - - float sum_sq = block_reduce_sum(thread_sum, warp_reduce_buf); - float rms_inv = rsqrtf(sum_sq / hidden_dim + eps); - - for (int i = threadIdx.x; i < hd2; i += blockDim.x) { - float2 wf = __half22float2(w2[i]); - y[i] = __float22half2_rn(make_float2(v_new.x * rms_inv * wf.x, - v_new.y * rms_inv * wf.y)); - } -} - -static __global__ void rmsnorm_quant_int8( - const half* in, - const half* weight, - int8_t* residual, - float* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - int row = blockIdx.x; // Which sequence/batch element - int idx = threadIdx.x; // owns adjacent pair (2*idx, 2*idx+1); blockDim.x == hidden_dim/2 - - if (row >= n_tokens) return; - - // Vectorized half2 in / weight, char2 int8 residual. - const half2* inp = reinterpret_cast(in + row * hidden_dim); - const half2* w2 = reinterpret_cast(weight); - char2* res = reinterpret_cast(residual + row * hidden_dim); - float* res_scale = residual_scale + row; - - __shared__ float warp_reduce_buf[32]; - - float2 wf = __half22float2(w2[idx]); - - // Step 1: RMS over (in + alpha * dequant(residual)) - float2 xf = __half22float2(inp[idx]); - char2 rq = res[idx]; - float rs = *res_scale; - float2 val = make_float2(xf.x + ((float)rq.x * rs) * alpha, - xf.y + ((float)rq.y * rs) * alpha); - float sum_sq = block_reduce_sum(val.x * val.x + val.y * val.y, warp_reduce_buf); - float rms_inv = rsqrtf(sum_sq / hidden_dim + eps); - - // Step 2: amax of the normalized values for the output quant scale - float2 normalized = make_float2(val.x * rms_inv * wf.x, val.y * rms_inv * wf.y); - float abs_max = block_reduce_max(fmaxf(fabsf(normalized.x), fabsf(normalized.y)), warp_reduce_buf); - - float quant_scale = (abs_max > 0.0f) ? (127.0f / abs_max) : 1.0f; - if (idx == 0) *res_scale = 1.0f / quant_scale; // all threads already read old *res_scale into val - - // clamp and write quantized norm - char2 q; - q.x = (int8_t)max(-127, min(127, __float2int_rn(normalized.x * quant_scale))); - q.y = (int8_t)max(-127, min(127, __float2int_rn(normalized.y * quant_scale))); - res[idx] = q; -} - -// FLSTM epilogue: fuse bias-add, gate activations, cell update, and hidden state out. -// scratch: (N, 4*C) — result of addmm(up_bias_hh_, hh[t], W_hh_fused) (bias already included) -// ih_t: (N, 4*C) — precomputed in-hidden contribution for this timestep -// c: (N, C) — cell state, updated in-place -// hh_next: (N, C) — out hidden state h[t+1] -// Launch: flstm_step<<>>(...) -static __global__ void flstm_step( - const half* scratch, - const half* ih_t, - half* cell, - half* hh_next, - int gate_dim, int hidden_dim -) { - int n = blockIdx.x; - for (int ch = threadIdx.x; ch < hidden_dim; ch += blockDim.x) { - int base = n * gate_dim + ch; - float gi = __half2float(scratch[base + 0*hidden_dim]) + __half2float(ih_t[base + 0*hidden_dim]); - float gf = __half2float(scratch[base + 1*hidden_dim]) + __half2float(ih_t[base + 1*hidden_dim]); - float gg = __half2float(scratch[base + 2*hidden_dim]) + __half2float(ih_t[base + 2*hidden_dim]); - float go = __half2float(scratch[base + 3*hidden_dim]) + __half2float(ih_t[base + 3*hidden_dim]); - float i_g = fmaxf(0.f, fminf(1.f, gi * 0.2f + 0.5f)); - float f_g = fmaxf(0.f, fminf(1.f, gf * 0.2f + 0.5f)); - float g_g = fmaxf(-1.f, fminf(1.f, gg)); - float o_g = fmaxf(0.f, fminf(1.f, go * 0.2f + 0.5f)); - float c_new = f_g * __half2float(cell[n * hidden_dim + ch]) + i_g * g_g; - cell[n * hidden_dim + ch] = __float2half(c_new); - hh_next[n * hidden_dim + ch] = __float2half(o_g * tanhf(c_new)); - } -} - -#ifdef __cplusplus -} -#endif - -#endif // NN_KERNEL_CUDA_H \ No newline at end of file diff --git a/src/nn_kernel_hip.h b/src/nn_kernel_hip.h deleted file mode 100644 index 00c7463..0000000 --- a/src/nn_kernel_hip.h +++ /dev/null @@ -1,414 +0,0 @@ -// The MIT License (MIT) - -// Copyright (c) 2025 Bonson Wong - -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: - -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. - -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -// THE SOFTWARE. - -#ifndef NN_KERNEL_HUI -#define NN_KERNEL_HUI - -#include -#include -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -static __global__ void rotary_emb( - half *x, - const float *_cos, - const float *_sin, - const uint64_t seq_len, - const uint64_t stride_batch, - const uint64_t stride_seq, - const uint64_t stride_head, - const uint64_t sincos_width -) { - const uint64_t batch = blockIdx.x; - const uint64_t head = blockIdx.y; - const uint64_t rot = threadIdx.x; - const uint64_t tid = threadIdx.y; - const uint64_t n_threads = blockDim.y; - - if (tid >= seq_len) return; - - half *_o0 = x + (batch * stride_batch) + (head * stride_head) + rot; - half *_o1 = x + (batch * stride_batch) + (head * stride_head) + sincos_width + rot; - - for (int seq = tid; seq < seq_len; seq += n_threads) { - float cos = *(_cos + (seq * sincos_width) + rot); - float sin = *(_sin + (seq * sincos_width) + rot); - - half *o0 = _o0 + (seq * stride_seq); - half *o1 = _o1 + (seq * stride_seq); - - float x0 = __half2float(*o0); - float x1 = __half2float(*o1); - - *o0 = __float2half(x0 * cos - x1 * sin); - *o1 = __float2half(x0 * sin + x1 * cos); - } -} - -static __global__ void silu_mul( - const half *in, - half *out, - const uint64_t hidden_dim, - const uint64_t n_tokens -) { - uint64_t j = blockIdx.x; - - for (uint64_t k = threadIdx.x; k < hidden_dim; k += blockDim.x) { - uint64_t i = k + j * (hidden_dim * 2); - - half y = in[i]; - half gate = in[i + hidden_dim]; - - float g = __half2float(gate); - float silu = g / (1.0f + __expf(-g)); - - out[k + j * hidden_dim] = __float2half(silu * __half2float(y)); - } -} - -// ── Block-wide reductions ───────────────────────────────────────────────────── -// Warp-level primitives, then a block-level reduction that broadcasts the result -// to all threads via a caller-supplied shared scratch buffer (>= blockDim.x/warpSize -// floats). The trailing double-sync leaves the buffer safe to reuse for a second -// reduction in the same kernel. blockDim.x <= 1024 => num_warps <= 16 (warpSize 64), -// so the final reduction fits in a single warp. -static __device__ __forceinline__ float warp_reduce_sum(float v) { - for (int offset = warpSize / 2; offset > 0; offset >>= 1) - v += __shfl_down(v, offset); - return v; -} -static __device__ __forceinline__ float warp_reduce_max(float v) { - for (int offset = warpSize / 2; offset > 0; offset >>= 1) - v = fmaxf(v, __shfl_down(v, offset)); - return v; -} -static __device__ __forceinline__ float block_reduce_sum(float v, float* warp_reduce_buf) { - int warp_id = threadIdx.x / warpSize; - int lane_id = threadIdx.x % warpSize; - v = warp_reduce_sum(v); - if (lane_id == 0) warp_reduce_buf[warp_id] = v; - __syncthreads(); - int num_warps = (blockDim.x + warpSize - 1) / warpSize; - float r = (threadIdx.x < num_warps) ? warp_reduce_buf[threadIdx.x] : 0.0f; - if (warp_id == 0) r = warp_reduce_sum(r); - if (threadIdx.x == 0) warp_reduce_buf[0] = r; - __syncthreads(); - float result = warp_reduce_buf[0]; - __syncthreads(); - return result; -} -static __device__ __forceinline__ float block_reduce_max(float v, float* warp_reduce_buf) { - int warp_id = threadIdx.x / warpSize; - int lane_id = threadIdx.x % warpSize; - v = warp_reduce_max(v); - if (lane_id == 0) warp_reduce_buf[warp_id] = v; - __syncthreads(); - int num_warps = (blockDim.x + warpSize - 1) / warpSize; - float r = (threadIdx.x < num_warps) ? warp_reduce_buf[threadIdx.x] : 0.0f; - if (warp_id == 0) r = warp_reduce_max(r); - if (threadIdx.x == 0) warp_reduce_buf[0] = r; - __syncthreads(); - float result = warp_reduce_buf[0]; - __syncthreads(); - return result; -} - -static __global__ void rmsnorm( - const half* in, - const half* residual, - const half* weight, - half* out, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - int row = blockIdx.x; // Which sequence/batch element - - if (row >= n_tokens) return; - - // Vectorized half2: each thread owns one adjacent pair. blockDim.x == hidden_dim/2. - const half2* x = reinterpret_cast(in + row * hidden_dim); - const half2* res = reinterpret_cast(residual + row * hidden_dim); - const half2* w2 = reinterpret_cast(weight); - half2* y = reinterpret_cast(out + row * hidden_dim); - int hd2 = hidden_dim >> 1; - - __shared__ float warp_reduce_buf[64]; - - float thread_sum = 0.0f; - float2 v_new; // valid because blockDim.x == hidden_dim/2, so this loop runs exactly once per thread - for (int i = threadIdx.x; i < hd2; i += blockDim.x) { - float2 xf = __half22float2(x[i]); - float2 rf = __half22float2(res[i]); - float2 val = make_float2(xf.x + rf.x * alpha, xf.y + rf.y * alpha); - v_new = val; - thread_sum += val.x * val.x + val.y * val.y; - } - - float sum_sq = block_reduce_sum(thread_sum, warp_reduce_buf); - float rms_inv = rsqrtf(sum_sq / hidden_dim + eps); - - for (int i = threadIdx.x; i < hd2; i += blockDim.x) { - float2 wf = __half22float2(w2[i]); - y[i] = __float22half2_rn(make_float2(v_new.x * rms_inv * wf.x, - v_new.y * rms_inv * wf.y)); - } -} - -// Convert E4M3FN fp8 byte (IEEE, PyTorch kFloat8_e4m3fn) to float32. -// NaN encodings 0x7F/0xFF -> 0.0f. -static __device__ __forceinline__ float e4m3fn_to_float(uint8_t b) { - if ((b & 0x7F) == 0x7F) return 0.0f; - if ((b & 0x7F) == 0) return 0.0f; - uint8_t sign = b >> 7; - uint8_t exp = (b >> 3) & 0xF; - uint8_t mant = b & 0x7; - uint32_t f32_bits; - if (exp == 0) { - // Denormal: (-1)^s * mant * 2^(-9) - float val = (float)mant * 1.953125e-3f; - f32_bits = __float_as_uint(val) | ((uint32_t)sign << 31); - } else { - // Normal: bias 7->127 means +120 to f32 exponent; top-align 3-bit mantissa - f32_bits = ((uint32_t)sign << 31) - | ((uint32_t)(exp + 120) << 23) - | ((uint32_t)mant << 20); - } - return __uint_as_float(f32_bits); -} - -// Convert float32 to E4M3FN fp8 byte (IEEE, PyTorch kFloat8_e4m3fn). -// Values outside [-448, 448] are saturated; NaN is never produced. -static __device__ __forceinline__ uint8_t float_to_e4m3fn(float f) { - uint32_t bits = __float_as_uint(f); - uint32_t sign = bits >> 31; - uint32_t f32_exp = (bits >> 23) & 0xFF; - uint32_t f32_mant = bits & 0x7FFFFF; - - if (f32_exp == 0) return (uint8_t)(sign << 7); // zero / float denormal -> fp8 zero - - int e4m3_exp = (int)f32_exp - 120; // unbiased_f32 = f32_exp-127, E4M3FN biased = +7 - - if (e4m3_exp >= 16) return (uint8_t)((sign << 7) | 0x7E); // saturate to +-448 - - if (e4m3_exp > 0) { - // Normal E4M3FN: round 23-bit float mantissa down to 3 bits - uint32_t mant3 = (f32_mant + (1U << 19)) >> 20; - if (mant3 >= 8) { mant3 = 0; ++e4m3_exp; } - if (e4m3_exp >= 16) return (uint8_t)((sign << 7) | 0x7E); - if (e4m3_exp == 15 && mant3 == 7) mant3 = 6; // avoid NaN encoding - return (uint8_t)((sign << 7) | ((uint32_t)e4m3_exp << 3) | mant3); - } else { - // Denormal E4M3FN: value = mant * 2^(-9) - if (e4m3_exp <= -4) return (uint8_t)(sign << 7); // underflow to zero - int shift = 21 - e4m3_exp; // 21..24 - uint32_t full = (1U << 23) | f32_mant; - uint32_t mant3 = (full + (1U << (shift - 1))) >> shift; - if (mant3 >= 8) return (uint8_t)((sign << 7) | (1U << 3)); // round up to normal exp=1 - return (uint8_t)((sign << 7) | mant3); - } -} - -// Combined rmsnorm + fp8 E4M3FN quantization kernel. -// Fuses: dequant fp8 residual -> deepnorm add -> rmsnorm -> quant fp8 residual. -// half2-vectorized: each thread owns one adjacent pair; one block per row (n_tokens rows total). -static __global__ void rmsnorm_quant_fp8( - const half* in, - const half* weight, - uint8_t* residual, - float* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - int row = blockIdx.x; - int idx = threadIdx.x; // owns adjacent pair (2*idx, 2*idx+1); blockDim.x == hidden_dim/2 - - if (row >= n_tokens) return; - - const half2* inp = reinterpret_cast(in + (int64_t)row * hidden_dim); - const half2* w2 = reinterpret_cast(weight); - uchar2* res = reinterpret_cast(residual + (int64_t)row * hidden_dim); - float* res_scale = residual_scale + row; - - __shared__ float warp_reduce_buf[64]; - - float2 wf = __half22float2(w2[idx]); - - // Dequantize fp8 residual, add scaled to new in (DeepNorm: in + alpha*residual) - float rs = *res_scale; - uchar2 rq = res[idx]; - float2 xf = __half22float2(inp[idx]); - float2 val = make_float2(xf.x + e4m3fn_to_float(rq.x) * rs * alpha, - xf.y + e4m3fn_to_float(rq.y) * rs * alpha); - - // ── Reduce sum-of-squares for RMS ───────────────────────────────────────── - float sum_sq = block_reduce_sum(val.x * val.x + val.y * val.y, warp_reduce_buf); - float rms_inv = rsqrtf(sum_sq / hidden_dim + eps); - - float2 normalized = make_float2(val.x * rms_inv * wf.x, val.y * rms_inv * wf.y); - - // ── Reduce amax for fp8 quantization scale ──────────────────────────────── - float abs_max = block_reduce_max(fmaxf(fabsf(normalized.x), fabsf(normalized.y)), warp_reduce_buf); - - float fp8_scale = fmaxf(abs_max, 1e-12f) / 448.0f; - if (idx == 0) *res_scale = fp8_scale; // all threads already read old *res_scale into val - - // ── Write quantized fp8 residual ────────────────────────────────────────── - uchar2 out; - out.x = float_to_e4m3fn(fmaxf(-448.0f, fminf(448.0f, normalized.x / fp8_scale))); - out.y = float_to_e4m3fn(fmaxf(-448.0f, fminf(448.0f, normalized.y / fp8_scale))); - res[idx] = out; -} - -// Standalone f16 → fp8 E4M3FN per-token quantize. -// One block per row, hidden_dim threads per block. -// scale[row] = max(amax/448, 1e-12); out[row] = clamp(x/scale, -448, 448) as fp8. -static __global__ void quant_fp8( - const half* in, // [n_tokens, hidden_dim] f16 input (row-major) - uint8_t* out, // [n_tokens, hidden_dim] fp8 output - float* scale, // [n_tokens] per-token scale output - int n_tokens, - int hidden_dim -) { - int row = blockIdx.x; - int idx = threadIdx.x; - - if (row >= n_tokens || idx >= hidden_dim) return; - - float val = __half2float(in[(int64_t)row * hidden_dim + idx]); - - // ── Per-row amax reduce ─────────────────────────────────────────────────── - __shared__ float warp_reduce_buf[64]; - float abs_max = block_reduce_max(fabsf(val), warp_reduce_buf); - - float fp8_scale = fmaxf(abs_max / 448.0f, 1e-12f); - if (idx == 0) scale[row] = fp8_scale; - - float fp8_val = fmaxf(-448.0f, fminf(448.0f, val / fp8_scale)); - out[(int64_t)row * hidden_dim + idx] = float_to_e4m3fn(fp8_val); -} - -// Dequantize fp8 [n_timesteps,batch_size,n_channels] (× scalar scale) AND transpose to -// f16 [batch_size,n_timesteps,n_channels], in one pass. -// out[n,t,c] = e4m3fn_to_float(in[t,n,c]) * scale -// One block per (t,n) row, n_channels threads. Reads/writes are coalesced along c (the contiguous -// innermost of both layouts); the transpose lives only in the block→(n,t) index map. -// Replaces torch's .to(kHalf).transpose(0,1).contiguous() (3 ops + intermediates). -static __global__ void dequant_fp8_transpose( - const uint8_t* in, // [n_timesteps, batch_size, n_channels] fp8 (row-major, contiguous) - half* out, // [batch_size, n_timesteps, n_channels] f16 - int n_timesteps, int batch_size, int n_channels, float scale -) { - int tn = blockIdx.x; // 0 .. n_timesteps*batch_size-1 - int c = threadIdx.x; // 0 .. n_channels-1 - if (c >= n_channels) return; - int t = tn / batch_size; - int n = tn - t * batch_size; - float v = e4m3fn_to_float(in[(int64_t)tn * n_channels + c]) * scale; - out[(((int64_t)n * n_timesteps + t) * n_channels) + c] = __float2half(v); -} - -static __global__ void rmsnorm_quant_int8( // need to verify if works on rocm - const half* in, - const half* weight, - int8_t* residual, - float* residual_scale, - int n_tokens, - int hidden_dim, - float alpha, - float eps -) { - int row = blockIdx.x; // Which sequence/batch element - int idx = threadIdx.x; // owns adjacent pair (2*idx, 2*idx+1); blockDim.x == hidden_dim/2 - - if (row >= n_tokens) return; - - // Vectorized half2 in / weight, char2 int8 residual. - const half2* inp = reinterpret_cast(in + row * hidden_dim); - const half2* w2 = reinterpret_cast(weight); - char2* res = reinterpret_cast(residual + row * hidden_dim); - float* res_scale = residual_scale + row; - - __shared__ float warp_reduce_buf[64]; - - float2 wf = __half22float2(w2[idx]); - - // Step 1: RMS over (in + alpha * dequant(residual)) - float2 xf = __half22float2(inp[idx]); - char2 rq = res[idx]; - float rs = *res_scale; - float2 val = make_float2(xf.x + ((float)rq.x * rs) * alpha, - xf.y + ((float)rq.y * rs) * alpha); - float sum_sq = block_reduce_sum(val.x * val.x + val.y * val.y, warp_reduce_buf); - float rms_inv = rsqrtf(sum_sq / hidden_dim + eps); - - // Step 2: amax of the normalized values for the output quant scale - float2 normalized = make_float2(val.x * rms_inv * wf.x, val.y * rms_inv * wf.y); - float abs_max = block_reduce_max(fmaxf(fabsf(normalized.x), fabsf(normalized.y)), warp_reduce_buf); - - float quant_scale = (abs_max > 0.0f) ? (127.0f / abs_max) : 1.0f; - if (idx == 0) *res_scale = 1.0f / quant_scale; // all threads already read old *res_scale into val - - // clamp and write quantized norm - char2 q; - q.x = (int8_t)max(-127, min(127, __float2int_rn(normalized.x * quant_scale))); - q.y = (int8_t)max(-127, min(127, __float2int_rn(normalized.y * quant_scale))); - res[idx] = q; -} - -static __global__ void flstm_step( - const half* scratch, - const half* ih_t, - half* cell, - half* hh_next, - int gate_dim, int hidden_dim -) { - int n = blockIdx.x; - for (int ch = threadIdx.x; ch < hidden_dim; ch += blockDim.x) { - int base = n * gate_dim + ch; - float gi = __half2float(scratch[base + 0*hidden_dim]) + __half2float(ih_t[base + 0*hidden_dim]); - float gf = __half2float(scratch[base + 1*hidden_dim]) + __half2float(ih_t[base + 1*hidden_dim]); - float gg = __half2float(scratch[base + 2*hidden_dim]) + __half2float(ih_t[base + 2*hidden_dim]); - float go = __half2float(scratch[base + 3*hidden_dim]) + __half2float(ih_t[base + 3*hidden_dim]); - float i_g = fmaxf(0.f, fminf(1.f, gi * 0.2f + 0.5f)); - float f_g = fmaxf(0.f, fminf(1.f, gf * 0.2f + 0.5f)); - float g_g = fmaxf(-1.f, fminf(1.f, gg)); - float o_g = fmaxf(0.f, fminf(1.f, go * 0.2f + 0.5f)); - float c_new = f_g * __half2float(cell[n * hidden_dim + ch]) + i_g * g_g; - cell[n * hidden_dim + ch] = __float2half(c_new); - hh_next[n * hidden_dim + ch] = __float2half(o_g * tanhf(c_new)); - } -} - -#ifdef __cplusplus -} -#endif - -#endif // NN_KERNEL_HUI \ No newline at end of file diff --git a/src/openfish.c b/src/openfish.c index d6e659a..45ff46b 100644 --- a/src/openfish.c +++ b/src/openfish.c @@ -1,5 +1,5 @@ #include -#include "decode.h" +#include "openfish_defs.h" openfish_opt_t openfish_decoder_default_opts(void) { openfish_opt_t opt = {100.0f, 2.0f, 0.0f, 1.0f}; diff --git a/src/openfish_defs.h b/src/openfish_defs.h new file mode 100644 index 0000000..0b1aa26 --- /dev/null +++ b/src/openfish_defs.h @@ -0,0 +1,65 @@ +#ifndef OPENFISH_DEFS_H +#define OPENFISH_DEFS_H + +// Single source of truth for the decoding constants, index/state types, beam structs, and the +// kernel-argument blocks. Shared by every backend: +// - CPU / CUDA / HIP #include this header directly +// - Metal host glue (decode_metal.mm) includes this header directly +// - Metal shader (openfish_metal.metal) receives it by build-time concatenation into the +// runtime shader source (newLibraryWithSource: cannot resolve local #includes) +// +// It must therefore stay valid as C, C++, CUDA, HIP and MSL. The fixed-width integer names below are +// native to all of them (MSL provides uint64_t/int32_t/... too); only the standard-library includes +// are gated on #ifdef __METAL_VERSION__ (MSL has no ). + +#ifndef __METAL_VERSION__ +#include +#include +#include // size_t +#endif + +#define NUM_BASE_BITS (2) +#define NUM_BASES (4) +#define NUM_TRANSITIONS (NUM_BASES + 1) +#define MAX_BEAM_WIDTH (32) +#define HASH_PRESENT_BITS (1024) +#define HASH_PRESENT_MASK (HASH_PRESENT_BITS - 1) +#define MAX_STATES (1024) +#define MAX_BEAM_CANDIDATES (NUM_TRANSITIONS * MAX_BEAM_WIDTH) +#define CRC_SEED (0x12345678u) + +typedef int32_t state_t; + +typedef struct beam_element { + state_t state; + uint8_t prev_element_index; + bool stay; +} beam_element_t; + +typedef struct beam_front_element { + uint32_t hash; + state_t state; + uint8_t prev_element_index; + bool stay; +} beam_front_element_t; + +// Scalar-only kernel parameter blocks shared by every backend. Device pointers (the score/guide tensors) +// are passed to the kernels separately — as explicit kernel arguments on CUDA/HIP, and bound as separate +// MTLBuffers on Metal — rather than embedded here, so a single definition stays valid across all of them. +// On Metal these are filled host-side via setBytes:, and read device-side as `constant &`. +typedef struct scan_params { + uint64_t num_states; + uint64_t n_timesteps; + uint64_t batch_size; + uint64_t n_channels; + float fixed_stay_score; +} scan_params_t; + +typedef struct beam_params { + uint64_t n_timesteps; + uint64_t batch_size; + uint64_t n_channels; + int32_t num_state_bits; +} beam_params_t; + +#endif // OPENFISH_DEFS_H diff --git a/src/openfish_metal.metal b/src/openfish_metal.metal new file mode 100644 index 0000000..aebad7f --- /dev/null +++ b/src/openfish_metal.metal @@ -0,0 +1,782 @@ +// Metal Shading Language port of the CRF-CTC decoding kernels. +// +// This mirrors the CUDA kernels in scan_cuda.h and beam_search_cuda.h. Each threadgroup +// decodes one chunk (batch element), exactly like one CUDA block per chunk. CUDA warps map +// to Metal SIMD-groups (32-wide on Apple GPUs), __syncthreads() maps to a threadgroup barrier, +// __shared__ maps to threadgroup memory, and __shfl_down_sync maps to simd_shuffle_down. +// +// The kernels are dispatched with threadsPerThreadgroup == num_states (scan kernels), +// MAX_BEAM_WIDTH*NUM_BASES (beam search), or 1 (sequence/qual generation), so the +// chunk/thread bounds checks never actually fire and every thread reaches each barrier. + +#include +using namespace metal; + +// The constants, state_t, beam_element_t / beam_front_element_t and the scan_params_t / +// beam_params_t argument blocks come from openfish_defs.h, which the Makefile concatenates ahead +// of this source at build time (newLibraryWithSource: can't resolve local #includes). Only the +// two shader-private constants below are defined here. + +#define SIMD_WIDTH (32) +// -FLT_MAX sentinel (matches CUDA's -FLT_MAX) +#define OF_FLT_MAX (0x1.fffffep127f) + +// ------------------------------------------------------------------ scan kernels + +kernel void bwd_scan( + constant scan_params_t &args [[buffer(0)]], + device const half *scores_in [[buffer(1)]], + device float *out [[buffer(2)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]] +) { + const ulong chunk = tgid; + const ulong state = tid; + + const ulong num_states = args.num_states; + const ulong n_timesteps = args.n_timesteps; + const ulong batch_size = args.batch_size; + + if (chunk >= batch_size || tid >= num_states) { + return; + } + + const float fixed_stay_score = args.fixed_stay_score; + const ulong ts_states = num_states * NUM_BASES; + + device const half *const chunk_in = scores_in + chunk * ts_states; + device float *const chunk_out = out + chunk * (n_timesteps + 1) * num_states; + device float *const alpha_init = chunk_out + num_states * n_timesteps; + alpha_init[state] = 0.0f; + + for (ulong ts = 0; ts < n_timesteps; ++ts) { + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + device const half *const ts_in = chunk_in + batch_size * ts_states * (n_timesteps - ts - 1); + device float *const ts_alpha_in = alpha_init - num_states * ts; + device float *const ts_alpha_out = ts_alpha_in - num_states; + + const ulong stay_state_idx = state; + const ulong step_state_idx_a = (state * NUM_BASES) % num_states; + const ulong step_trans_idx_a = step_state_idx_a * NUM_BASES + ((state * NUM_BASES) / num_states); + + float vals[NUM_TRANSITIONS]; + vals[0] = ts_alpha_in[stay_state_idx] + fixed_stay_score; + float max_val = vals[0]; + for (ulong base = 0; base < NUM_BASES; ++base) { + vals[base + 1] = ts_alpha_in[step_state_idx_a + base] + float(ts_in[step_trans_idx_a + base * NUM_BASES]); + max_val = max_val > vals[base + 1] ? max_val : vals[base + 1]; + } + float sum = 0.0f; + for (ulong i = 0; i < NUM_TRANSITIONS; ++i) { + sum += exp(vals[i] - max_val); + } + ts_alpha_out[state] = max_val + log(sum); + } +} + +kernel void fwd_post_scan( + constant scan_params_t &args [[buffer(0)]], + device const half *scores_in [[buffer(1)]], + device const float *bwd [[buffer(2)]], + device float *out [[buffer(3)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint nthreads[[threads_per_threadgroup]], + uint lane_id [[thread_index_in_simdgroup]], + uint warp_id [[simdgroup_index_in_threadgroup]] +) { + const ulong chunk = tgid; + + const ulong num_states = args.num_states; + const ulong n_timesteps = args.n_timesteps; + const ulong n_ts = n_timesteps + 1; + const ulong batch_size = args.batch_size; + + if (chunk >= batch_size || tid >= num_states) { + return; + } + + const ulong state = tid; + const float fixed_stay_score = args.fixed_stay_score; + + const ulong msb = num_states / NUM_BASES; + const ulong ts_states = num_states * NUM_BASES; + + threadgroup float fwd_vals[MAX_STATES]; + threadgroup float fwd_maxs[32]; + threadgroup float exp_vals[MAX_STATES]; + threadgroup float exp_sums[32]; + threadgroup float ts_fwd[2][MAX_STATES]; + float warp_max; + + device const half *const chunk_scores = scores_in + chunk * ts_states; + + for (ulong s = tid; s < num_states; s += nthreads) { + ts_fwd[0][s] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + for (ulong ts = 0; ts < n_ts; ++ts) { + warp_max = -OF_FLT_MAX; + const ulong ts_idx = (chunk * n_ts + ts) * num_states; + + threadgroup const float *const ts_alpha_in = ts_fwd[ts & 1]; + threadgroup float *const ts_alpha_out = ts_fwd[(ts & 1) ^ 1]; + + if (ts < n_timesteps) { + device const half *const ts_scores = chunk_scores + batch_size * ts_states * ts; + + const ulong stay_state_idx = state; + const ulong step_state_idx_a = state / NUM_BASES; + const ulong step_trans_idx_a = state * NUM_BASES; + float vals[NUM_TRANSITIONS]; + float fwd_max_val = vals[0] = ts_alpha_in[stay_state_idx] + fixed_stay_score; + for (ulong base = 0; base < NUM_BASES; ++base) { + vals[base + 1] = ts_alpha_in[step_state_idx_a + base * msb] + + float(ts_scores[step_trans_idx_a + base]); + fwd_max_val = fwd_max_val > vals[base + 1] ? fwd_max_val : vals[base + 1]; + } + float fwd_sum = 0.0f; + for (ulong i = 0; i < NUM_TRANSITIONS; ++i) { + fwd_sum += exp(vals[i] - fwd_max_val); + } + ts_alpha_out[state] = fwd_max_val + log(fwd_sum); + } + + const float fwd_val = ts_alpha_in[state]; + const float val = fwd_val + bwd[ts_idx + state]; + + fwd_vals[state] = val; + warp_max = max(warp_max, val); + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // find max fwd val in simdgroup + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) { + warp_max = max(warp_max, simd_shuffle_down(warp_max, offset)); + } + if (lane_id == 0) fwd_maxs[warp_id] = warp_max; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // set max fwd vals across all simdgroups + if (warp_id == 0) { + warp_max = (tid < num_states / SIMD_WIDTH) ? fwd_maxs[lane_id] : 0; + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) { + warp_max = max(warp_max, simd_shuffle_down(warp_max, offset)); + } + if (tid == 0) fwd_maxs[0] = warp_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // enter exp vals + float warp_sum = 0.0f; + exp_vals[state] = exp(fwd_vals[state] - fwd_maxs[0]); + warp_sum += exp_vals[state]; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // sum exp vals in simdgroup + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) { + warp_sum += simd_shuffle_down(warp_sum, offset); + } + if (lane_id == 0) exp_sums[warp_id] = warp_sum; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // sum exp vals across all simdgroups + if (warp_id == 0) { + warp_sum = (tid < num_states / SIMD_WIDTH) ? exp_sums[lane_id] : 0; + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) { + warp_sum += simd_shuffle_down(warp_sum, offset); + } + if (tid == 0) exp_sums[0] = warp_sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + out[ts_idx + state] = exp_vals[state] / exp_sums[0]; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } +} + +// ------------------------------------------------------------------ beam search helpers + +static inline void swapf(thread float *a, thread float *b) { + float temp = *a; *a = *b; *b = temp; +} + +static inline int partitionf(threadgroup float *nums, int left, int right) { + int pivot = left; + { float t = nums[pivot]; nums[pivot] = nums[right]; nums[right] = t; } + + for (int i = left; i < right; ++i) { + if (nums[i] > nums[right]) { + float t = nums[pivot]; nums[pivot] = nums[i]; nums[i] = t; + ++pivot; + } + } + { float t = nums[pivot]; nums[pivot] = nums[right]; nums[right] = t; } + return pivot; +} + +static inline float kth_largestf(threadgroup float *nums, int k, int n) { + int left = 0; + int right = n - 1; + while (left <= right) { + int pivot = partitionf(nums, left, right); + if (pivot < k) { + left = pivot + 1; + } else if (pivot > k) { + right = pivot - 1; + } else { + return nums[pivot]; + } + } + return -OF_FLT_MAX; +} + +static inline float log_sum_exp(float x, float y) { + float abs_diff = fabs(x - y); + float m = x > y ? x : y; + return m + ((abs_diff < 17.0f) ? (log(1.0f + exp(-abs_diff))) : 0.0f); +} + +// Castagnoli CRC32 (CRC32C), reversed polynomial. +static inline uint crc32c(uint crc, uint new_bits, int num_new_bits) { + const uint POLYNOMIAL = 0x82f63b78u; + for (int i = 0; i < num_new_bits; ++i) { + uint b = (new_bits ^ crc) & 1; + crc >>= 1; + if (b) { + crc ^= POLYNOMIAL; + } + new_bits >>= 1; + } + return crc; +} + +// ------------------------------------------------------------------ beam search + +kernel void beam_search( + constant beam_params_t &beam_args [[buffer(0)]], + device const half *scores_TNC_in [[buffer(1)]], + device const float *bwd_NTC_in [[buffer(2)]], + device state_t *_states [[buffer(3)]], + device uchar *_moves [[buffer(4)]], + device beam_element_t *_beam_vector [[buffer(5)]], + constant float &beam_cut [[buffer(6)]], + constant float &fixed_stay_score [[buffer(7)]], + constant float &score_scale [[buffer(8)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint lane_id [[thread_index_in_simdgroup]], + uint warp_id [[simdgroup_index_in_threadgroup]] +) { + const ulong chunk = tgid; + const ulong n_threads = MAX_BEAM_WIDTH * NUM_BASES; + + if (chunk >= beam_args.batch_size || tid >= n_threads) { + return; + } + + const ulong n_timesteps = beam_args.n_timesteps; + const ulong batch_size = beam_args.batch_size; + const ulong n_channels = beam_args.n_channels; + + const int num_state_bits = beam_args.num_state_bits; + const ulong num_states = 1ull << num_state_bits; + const state_t states_mask = (state_t)(num_states - 1); + const ulong scores_block_stride = batch_size * n_channels; + const float log_beam_cut = (beam_cut > 0.0f) ? log(beam_cut) : OF_FLT_MAX; + + device const half *scores_TNC = scores_TNC_in + chunk * (num_states * NUM_BASES); + device const float *bwd_NTC = bwd_NTC_in + chunk * num_states * (n_timesteps + 1); + device state_t *states = _states + chunk * n_timesteps; + device uchar *moves = _moves + chunk * n_timesteps; + device beam_element_t *beam_vector = _beam_vector + chunk * MAX_BEAM_WIDTH * (n_timesteps + 1); + + threadgroup beam_front_element_t current_beam_front[MAX_BEAM_CANDIDATES]; + threadgroup beam_front_element_t prev_beam_front[MAX_BEAM_CANDIDATES]; + threadgroup float current_scores[MAX_BEAM_CANDIDATES]; + threadgroup float prev_scores[MAX_BEAM_CANDIDATES]; + + // candidate-phase scratch: k=1 Bloom filter during candidate/stay generation, then reused + // as the stream-compaction prefix-sum storage (compact_offsets) in the pruning phase. + threadgroup bool cand_scratch[HASH_PRESENT_BITS]; + threadgroup bool *const step_hash_present = cand_scratch; + threadgroup int *const compact_offsets = (threadgroup int *)cand_scratch; + + threadgroup ulong current_beam_width; + threadgroup float beam_init_threshold; + + // back-guide sort scratch (num_states floats; sized to the MAX_STATES upper bound here). + threadgroup float sorted_back_guides[MAX_STATES]; + + for (ulong beam_element = tid; beam_element < MAX_BEAM_CANDIDATES; beam_element += n_threads) { + current_beam_front[beam_element] = beam_front_element_t{0, 0, 0, false}; + prev_beam_front[beam_element] = beam_front_element_t{0, 0, 0, false}; + current_scores[beam_element] = 0.0f; + prev_scores[beam_element] = 0.0f; + } + + if (tid == 0) { + beam_init_threshold = -OF_FLT_MAX; + current_beam_width = MAX_BEAM_WIDTH < num_states ? MAX_BEAM_WIDTH : num_states; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + if (MAX_BEAM_WIDTH < num_states) { + for (ulong i = tid; i < num_states; i += n_threads) { + sorted_back_guides[i] = bwd_NTC[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (tid == 0) beam_init_threshold = kth_largestf(sorted_back_guides, MAX_BEAM_WIDTH - 1, (int)num_states); + } + + if (tid == 0) { + for (ulong state = 0, beam_element = 0; state < num_states && beam_element < MAX_BEAM_WIDTH; state++) { + if (bwd_NTC[state] >= beam_init_threshold) { + beam_front_element_t new_elem = {crc32c(CRC_SEED, (uint)state, 32), (state_t)state, 0, false}; + prev_beam_front[beam_element] = new_elem; + prev_scores[beam_element] = 0.0f; + ++beam_element; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + for (ulong element_idx = tid; element_idx < current_beam_width; element_idx += n_threads) { + beam_vector[element_idx].state = prev_beam_front[element_idx].state; + beam_vector[element_idx].prev_element_index = prev_beam_front[element_idx].prev_element_index; + beam_vector[element_idx].stay = prev_beam_front[element_idx].stay; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + threadgroup int elem_count; + threadgroup float warp_buf[64]; + threadgroup float *const max_buf = warp_buf; + threadgroup int *const count_buf = (threadgroup int *)warp_buf; + threadgroup float max_score; + threadgroup uint new_elem_count; + threadgroup float beam_cutoff_score; + threadgroup int entered_search; + threadgroup int bs_active; + + for (ulong block_idx = 0; block_idx < n_timesteps; ++block_idx) { + device const half *const block_scores = scores_TNC + (block_idx * scores_block_stride); + device const float *const block_back_scores = bwd_NTC + ((block_idx + 1) << num_state_bits); + + float warp_max = -OF_FLT_MAX; + if (tid == 0) { + new_elem_count = 0; + } + + for (uint i = tid; i < HASH_PRESENT_BITS; i += n_threads) { + step_hash_present[i] = false; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // generate candidate step elements for this timestep + for (ulong prev_elem_idx = (tid / NUM_BASES); prev_elem_idx < current_beam_width; prev_elem_idx += n_threads) { + threadgroup const beam_front_element_t *previous_element = &prev_beam_front[prev_elem_idx]; + const int new_base = tid % NUM_BASES; + + state_t new_state = ((state_t)((previous_element->state << NUM_BASE_BITS) & states_mask) | (state_t)(new_base)); + const state_t move_idx = (state_t)((new_state << NUM_BASE_BITS) + (((previous_element->state << NUM_BASE_BITS) >> num_state_bits))); + + float block_score = float(block_scores[move_idx]) * score_scale; + float new_score = prev_scores[prev_elem_idx] + block_score + (float)block_back_scores[new_state]; + + uint new_hash = crc32c(previous_element->hash, new_base, NUM_BASE_BITS); + step_hash_present[new_hash & HASH_PRESENT_MASK] = true; + + uint new_elem_idx = new_elem_count + (prev_elem_idx * NUM_BASES) + new_base; + + beam_front_element_t new_beam_elem = { + new_hash, + new_state, + (uchar)prev_elem_idx, + false + }; + current_beam_front[new_elem_idx] = new_beam_elem; + current_scores[new_elem_idx] = new_score; + warp_max = max(warp_max, new_score); + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (tid == 0) new_elem_count += current_beam_width * NUM_BASES; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // generate stays and fold duplicate step/stay paths + for (ulong prev_elem_idx = (tid / NUM_BASES); prev_elem_idx < current_beam_width; prev_elem_idx += n_threads) { + threadgroup const beam_front_element_t *previous_element = &prev_beam_front[prev_elem_idx]; + uint new_elem_idx = new_elem_count + prev_elem_idx; + + const float stay_score = prev_scores[prev_elem_idx] + + fixed_stay_score + + (float)block_back_scores[previous_element->state]; + + beam_front_element_t new_beam_elem = { + previous_element->hash, + previous_element->state, + (uchar)prev_elem_idx, + true + }; + current_beam_front[new_elem_idx] = new_beam_elem; + current_scores[new_elem_idx] = stay_score; + + warp_max = max(warp_max, stay_score); + + if (step_hash_present[previous_element->hash & HASH_PRESENT_MASK]) { + ulong stay_elem_idx = (current_beam_width << NUM_BASE_BITS) + prev_elem_idx; + int stay_latest_base = (int)(previous_element->state & 3); + + for (ulong prev_elem_comp_idx = (tid % NUM_BASES); prev_elem_comp_idx < current_beam_width; prev_elem_comp_idx += NUM_BASES) { + ulong step_elem_idx = (prev_elem_comp_idx << NUM_BASE_BITS) | stay_latest_base; + + if (current_beam_front[stay_elem_idx].hash == current_beam_front[step_elem_idx].hash) { + const float folded_score = log_sum_exp(current_scores[stay_elem_idx], current_scores[step_elem_idx]); + if (current_scores[stay_elem_idx] > current_scores[step_elem_idx]) { + current_scores[stay_elem_idx] = folded_score; + current_scores[step_elem_idx] = -OF_FLT_MAX; + } else { + current_scores[step_elem_idx] = folded_score; + current_scores[stay_elem_idx] = -OF_FLT_MAX; + } + warp_max = max(warp_max, folded_score); + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (tid == 0) { new_elem_count += current_beam_width; } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // block-wide max reduction + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) { + warp_max = max(warp_max, simd_shuffle_down(warp_max, offset)); + } + if (lane_id == 0) max_buf[warp_id] = warp_max; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + if (warp_id == 0) { + warp_max = (tid < n_threads / SIMD_WIDTH) ? max_buf[lane_id] : 0; + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) { + warp_max = max(warp_max, simd_shuffle_down(warp_max, offset)); + } + if (tid == 0) max_score = warp_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + if (tid == 0) { + beam_cutoff_score = max_score - log_beam_cut; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // count elements >= cutoff (block-wide reduction; result in count_buf[0]) + { + int local = 0; + for (int i = tid; i < (int)new_elem_count; i += n_threads) { + if (current_scores[i] >= beam_cutoff_score) ++local; + } + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) local += simd_shuffle_down(local, offset); + if (lane_id == 0) count_buf[warp_id] = local; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (warp_id == 0) { + int v = (tid < (int)(n_threads / SIMD_WIDTH)) ? count_buf[lane_id] : 0; + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) v += simd_shuffle_down(v, offset); + if (tid == 0) count_buf[0] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + if (tid == 0) { + elem_count = count_buf[0]; + entered_search = (elem_count > MAX_BEAM_WIDTH) ? 1 : 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // binary search for a cutoff score that keeps the beam within [80%, 100%] of MAX_BEAM_WIDTH. + if (entered_search) { + const ulong min_beam_width = (MAX_BEAM_WIDTH * 8) / 10; + float low_score = beam_cutoff_score; + float hi_score = max_score; + int num_guesses = 1; + const int MAX_GUESSES = 10; + + while (true) { + if (tid == 0) { + if ((elem_count > MAX_BEAM_WIDTH || elem_count < (int)min_beam_width) && num_guesses < MAX_GUESSES) { + if (elem_count > MAX_BEAM_WIDTH) { + low_score = beam_cutoff_score; + beam_cutoff_score = (beam_cutoff_score + hi_score) / 2.0f; + } else { + hi_score = beam_cutoff_score; + beam_cutoff_score = (beam_cutoff_score + low_score) / 2.0f; + } + ++num_guesses; + bs_active = 1; + } else { + bs_active = 0; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (!bs_active) break; + + { + int local = 0; + for (int i = tid; i < (int)new_elem_count; i += n_threads) { + if (current_scores[i] >= beam_cutoff_score) ++local; + } + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) local += simd_shuffle_down(local, offset); + if (lane_id == 0) count_buf[warp_id] = local; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (warp_id == 0) { + int v = (tid < (int)(n_threads / SIMD_WIDTH)) ? count_buf[lane_id] : 0; + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) v += simd_shuffle_down(v, offset); + if (tid == 0) count_buf[0] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + if (tid == 0) elem_count = count_buf[0]; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + + if (tid == 0) { + bs_active = (num_guesses == MAX_GUESSES) ? 1 : 0; + if (bs_active) beam_cutoff_score = hi_score; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (bs_active) { + int local = 0; + for (int i = tid; i < (int)new_elem_count; i += n_threads) { + if (current_scores[i] >= beam_cutoff_score) ++local; + } + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) local += simd_shuffle_down(local, offset); + if (lane_id == 0) count_buf[warp_id] = local; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (warp_id == 0) { + int v = (tid < (int)(n_threads / SIMD_WIDTH)) ? count_buf[lane_id] : 0; + for (int offset = SIMD_WIDTH / 2; offset > 0; offset >>= 1) v += simd_shuffle_down(v, offset); + if (tid == 0) count_buf[0] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (tid == 0) elem_count = count_buf[0]; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + + if (tid == 0) elem_count = elem_count < MAX_BEAM_WIDTH ? elem_count : MAX_BEAM_WIDTH; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + + // stream compaction: exclusive prefix sum of the >= cutoff predicate gives each passing + // element its serial write index; gate on dst < MAX_BEAM_WIDTH to keep the first 32 in order. + for (int i = tid; i < (int)new_elem_count; i += n_threads) { + compact_offsets[i] = (current_scores[i] >= beam_cutoff_score) ? 1 : 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + for (int d = 1; d < (int)new_elem_count; d <<= 1) { + int a0 = 0, a1 = 0; + int i0 = tid, i1 = tid + (int)n_threads; + if (i0 < (int)new_elem_count && i0 >= d) a0 = compact_offsets[i0 - d]; + if (i1 < (int)new_elem_count && i1 >= d) a1 = compact_offsets[i1 - d]; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (i0 < (int)new_elem_count && i0 >= d) compact_offsets[i0] += a0; + if (i1 < (int)new_elem_count && i1 >= d) compact_offsets[i1] += a1; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + for (int i = tid; i < (int)new_elem_count; i += n_threads) { + if (current_scores[i] >= beam_cutoff_score) { + int dst = compact_offsets[i] - 1; + if (dst < MAX_BEAM_WIDTH) { + prev_beam_front[dst] = current_beam_front[i]; + prev_scores[dst] = current_scores[i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (tid == 0) { + int total_pass = (new_elem_count > 0) ? compact_offsets[new_elem_count - 1] : 0; + elem_count = total_pass < MAX_BEAM_WIDTH ? total_pass : MAX_BEAM_WIDTH; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // at the last timestep make the best path element 0 + if (tid == 0 && block_idx == n_timesteps - 1) { + float best_score = -OF_FLT_MAX; + ulong best_score_index = 0; + for (ulong i = 0; i < (ulong)elem_count; i++) { + if (prev_scores[i] > best_score) { + best_score = prev_scores[i]; + best_score_index = i; + } + } + beam_front_element_t temp0 = prev_beam_front[0]; + prev_beam_front[0] = prev_beam_front[best_score_index]; + prev_beam_front[best_score_index] = temp0; + + float temp1 = prev_scores[0]; + prev_scores[0] = prev_scores[best_score_index]; + prev_scores[best_score_index] = temp1; + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + + // copy this new beam front into persistent beam state + ulong beam_offset = (block_idx + 1) * MAX_BEAM_WIDTH; + for (ulong i = tid; i < (ulong)elem_count; i += n_threads) { + prev_scores[i] -= (float)block_back_scores[prev_beam_front[i].state]; + + beam_vector[beam_offset + i].state = prev_beam_front[i].state; + beam_vector[beam_offset + i].prev_element_index = prev_beam_front[i].prev_element_index; + beam_vector[beam_offset + i].stay = prev_beam_front[i].stay; + } + if (tid == 0) current_beam_width = elem_count; + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + } + + // trace back the best path into states/moves + if (tid == 0) { + uchar element_index = 0; + for (ulong beam_idx = n_timesteps; beam_idx != 0; --beam_idx) { + ulong beam_addr = beam_idx * MAX_BEAM_WIDTH + element_index; + states[beam_idx - 1] = (state_t)beam_vector[beam_addr].state; + moves[beam_idx - 1] = beam_vector[beam_addr].stay ? 0 : 1; + element_index = beam_vector[beam_addr].prev_element_index; + } + moves[0] = 1; + } +} + +// ------------------------------------------------------------------ qual / sequence + +kernel void compute_qual_data( + constant beam_params_t &beam_args [[buffer(0)]], + device const float *post_NTC_in [[buffer(1)]], + device state_t *_states [[buffer(2)]], + device float *_qual_data [[buffer(3)]], + constant float &posts_scale [[buffer(4)]], + uint tgid [[threadgroup_position_in_grid]] +) { + const ulong chunk = tgid; + if (chunk >= beam_args.batch_size) { + return; + } + + const ulong n_timesteps = beam_args.n_timesteps; + const ulong num_states = 1ull << beam_args.num_state_bits; + const ulong num_state_bits = beam_args.num_state_bits; + + device const float *post_NTC = post_NTC_in + chunk * num_states * (n_timesteps + 1); + device state_t *states = _states + chunk * n_timesteps; + device float *qual_data = _qual_data + chunk * (n_timesteps * NUM_BASES); + + int shifted_states[2 * NUM_BASES]; + + for (ulong block_idx = 0; block_idx < n_timesteps; ++block_idx) { + int state = states[block_idx]; + states[block_idx] = states[block_idx] % NUM_BASES; + int base_to_emit = states[block_idx]; + + device const float *const timestep_posts = post_NTC + ((block_idx + 1) << num_state_bits); + + float block_prob = (float)(timestep_posts[state]) * posts_scale; + + int l_shift_idx = state >> NUM_BASE_BITS; + int r_shift_idx = (state << NUM_BASE_BITS) % num_states; + int msb = ((int)num_states) >> NUM_BASE_BITS; + int l_shift_state, r_shift_state; + for (int shift_base = 0; shift_base < NUM_BASES; ++shift_base) { + l_shift_state = l_shift_idx + msb * shift_base; + shifted_states[2 * shift_base] = l_shift_state; + + r_shift_state = r_shift_idx + shift_base; + shifted_states[2 * shift_base + 1] = r_shift_state; + } + + int candidate_state; + for (ulong state_idx = 0; state_idx < 2 * NUM_BASES; ++state_idx) { + candidate_state = shifted_states[state_idx]; + bool count_state = (candidate_state != state); + if (count_state) { + for (ulong inner_state = 0; inner_state < state_idx; ++inner_state) { + if (shifted_states[inner_state] == candidate_state) { + count_state = false; + break; + } + } + } + if (count_state) { + block_prob += (float)(timestep_posts[candidate_state]) * posts_scale; + } + } + block_prob = fmin(fmax(block_prob, 0.0f), 1.0f); + block_prob = pow(block_prob, 0.4f); + + const float wrong_base_prob = (1.0f - block_prob) / 3.0f; + + for (ulong base = 0; base < NUM_BASES; base++) { + qual_data[block_idx * NUM_BASES + base] = ((int)base == base_to_emit ? block_prob : wrong_base_prob); + } + } +} + +kernel void generate_sequence( + constant beam_params_t &args [[buffer(0)]], + device const uchar *_moves [[buffer(1)]], + device const state_t *_states [[buffer(2)]], + device const float *_qual_data [[buffer(3)]], + device float *_base_probs [[buffer(4)]], + device float *_total_probs [[buffer(5)]], + device char *_sequence [[buffer(6)]], + device char *_qstring [[buffer(7)]], + constant float &shift [[buffer(8)]], + constant float &scale [[buffer(9)]], + uint tgid [[threadgroup_position_in_grid]] +) { + const ulong chunk = tgid; + if (chunk >= args.batch_size) { + return; + } + + const ulong n_timesteps = args.n_timesteps; + device const uchar *moves = _moves + chunk * n_timesteps; + device const state_t *states = _states + chunk * n_timesteps; + device const float *qual_data = _qual_data + chunk * n_timesteps * NUM_BASES; + device float *base_probs = _base_probs + chunk * n_timesteps; + device float *total_probs = _total_probs + chunk * n_timesteps; + device char *sequence = _sequence + chunk * n_timesteps; + device char *qstring = _qstring + chunk * n_timesteps; + + ulong seq_len = 0; + for (ulong i = 0; i < n_timesteps; ++i) { + seq_len += moves[i]; + base_probs[i] = 0.0f; + total_probs[i] = 0.0f; + } + + ulong seq_pos = 0; + const char alphabet[4] = {'A', 'C', 'G', 'T'}; + + for (ulong blk = 0; blk < n_timesteps; ++blk) { + int state = states[blk]; + int move = (int)moves[blk]; + int base = state & 3; + int offset = (blk == 0) ? 0 : move - 1; + int probPos = (int)(seq_pos + offset); + + base_probs[probPos] += qual_data[blk * NUM_BASES + base]; + for (ulong k = 0; k < NUM_BASES; ++k) { + total_probs[probPos] += qual_data[blk * NUM_BASES + k]; + } + + if (blk == 0) { + sequence[seq_pos++] = (char)base; + } else { + for (int j = 0; j < move; ++j) { + sequence[seq_pos++] = (char)base; + } + } + } + + for (ulong i = 0; i < seq_len; ++i) { + sequence[i] = alphabet[(int)sequence[i]]; + base_probs[i] = 1.0f - (base_probs[i] / total_probs[i]); + base_probs[i] = -10.0f * log10(base_probs[i]); + float qscore = fmin(fmax(base_probs[i] * scale + shift, 1.0f), 50.0f); + qstring[i] = (char)((int)(33.5f + qscore)); + } +} diff --git a/src/scan_cuda.h b/src/scan_cuda.h index 7cb8aac..c9964b0 100644 --- a/src/scan_cuda.h +++ b/src/scan_cuda.h @@ -5,17 +5,18 @@ #include #include -#include "decode.h" +#include "openfish_defs.h" static __global__ void bwd_scan( - const scan_args_t args, + const scan_params_t args, + const void *_scores_in, float *out ) { const uint64_t chunk = blockIdx.x + (blockIdx.y * gridDim.x); const uint64_t tid = threadIdx.x + (threadIdx.y * blockDim.x); const uint64_t state = tid; - const half *scores_in = (const half *)args.scores_in; + const half *scores_in = (const half *)_scores_in; const uint64_t num_states = args.num_states; const uint64_t n_timesteps = args.n_timesteps; const uint64_t batch_size = args.batch_size; @@ -59,7 +60,8 @@ static __global__ void bwd_scan( } static __global__ void fwd_post_scan( - const scan_args_t args, + const scan_params_t args, + const void *_scores_in, const float *bwd, float *out ) { @@ -72,7 +74,7 @@ static __global__ void fwd_post_scan( (void)mask; const uint64_t state = tid; - const half *scores_in = (const half *)args.scores_in; + const half *scores_in = (const half *)_scores_in; const uint64_t num_states = args.num_states; const uint64_t n_timesteps = args.n_timesteps; const uint64_t n_ts = args.n_timesteps + 1; diff --git a/src/scan_hip.h b/src/scan_hip.h index c82eecd..e4b3e84 100644 --- a/src/scan_hip.h +++ b/src/scan_hip.h @@ -7,17 +7,18 @@ #include #include -#include "decode.h" +#include "openfish_defs.h" static __global__ void bwd_scan( - const scan_args_t args, + const scan_params_t args, + const void *_scores_in, float *out ) { const uint64_t chunk = blockIdx.x + (blockIdx.y * gridDim.x); const uint64_t tid = threadIdx.x + (threadIdx.y * blockDim.x); const uint64_t state = tid; - const half *scores_in = (const half *)args.scores_in; + const half *scores_in = (const half *)_scores_in; const uint64_t num_states = args.num_states; const uint64_t n_timesteps = args.n_timesteps; const uint64_t batch_size = args.batch_size; @@ -61,7 +62,8 @@ static __global__ void bwd_scan( } static __global__ void fwd_post_scan( - const scan_args_t args, + const scan_params_t args, + const void *_scores_in, const float *bwd, float *out ) { @@ -74,7 +76,7 @@ static __global__ void fwd_post_scan( (void)mask; const uint64_t state = tid; - const half *scores_in = (const half *)args.scores_in; + const half *scores_in = (const half *)_scores_in; const uint64_t num_states = args.num_states; const uint64_t n_timesteps = args.n_timesteps; const uint64_t n_ts = args.n_timesteps + 1; diff --git a/src/test_utils_cuda.h b/src/test_utils_cuda.h index 664e647..2066bd3 100644 --- a/src/test_utils_cuda.h +++ b/src/test_utils_cuda.h @@ -1,7 +1,7 @@ #pragma once #include -#include "decode.h" +#include "openfish_defs.h" #include "error.h" #include "cuda_utils.h" diff --git a/src/test_utils_hip.h b/src/test_utils_hip.h index 9c0d2b4..98e573d 100644 --- a/src/test_utils_hip.h +++ b/src/test_utils_hip.h @@ -1,7 +1,7 @@ #pragma once #include -#include "decode.h" +#include "openfish_defs.h" #include "error.h" #include "hip_utils.h" diff --git a/src/test_utils_metal.h b/src/test_utils_metal.h new file mode 100644 index 0000000..487df61 --- /dev/null +++ b/src/test_utils_metal.h @@ -0,0 +1,38 @@ +#pragma once + +// Test-harness helpers for the Metal backend, mirroring test_utils_cuda.h / test_utils_hip.h. +// These are implemented in decode_metal.mm and exposed with C linkage so main.c (compiled as C) +// can call them. + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +void set_device_metal(int device); + +// scores_TNC is host float16 [T,N,C] (same as the CUDA/ROCm GPU path). +// returns an opaque handle (an MTLBuffer) to pass as scores_TNC into openfish_decode_gpu. +void *upload_scores_to_metal( + int n_timesteps, + int batch_size, + int n_channels, + const void *scores_TNC +); + +void free_scores_metal(void *scores_gpu); + +#ifdef DEBUG +void write_gpubuf_metal( + uint64_t n_timesteps, + uint64_t batch_size, + int state_len, + const openfish_gpubuf_t *gpubuf +); +#endif + +#ifdef __cplusplus +} +#endif diff --git a/test/requirements.txt b/test/requirements.txt deleted file mode 100644 index 643a8f0..0000000 --- a/test/requirements.txt +++ /dev/null @@ -1,9 +0,0 @@ -# Reference/test deps for the rotary embedding kernel (test/*.py). -# -# test_rotary_cuda.py JIT-compiles a CUDA binding with ninja and runs it on the -# GPU, so it needs a CUDA-enabled torch (NOT the +cpu build) plus a CUDA toolkit -# (CUDA_HOME, default /usr/local/cuda). Matches the working ../quantrnn setup. -# -# pip install --index-url https://download.pytorch.org/whl/cu126 -r test/requirements.txt -torch==2.7.1 -ninja==1.11.* diff --git a/test/test_rmsnorm_cuda.py b/test/test_rmsnorm_cuda.py deleted file mode 100644 index 0da1933..0000000 --- a/test/test_rmsnorm_cuda.py +++ /dev/null @@ -1,234 +0,0 @@ -"""Build-and-run test for the CUDA RMSNorm kernels. - -Single script: it JIT-compiles the real CUDA kernels (src/nn_cuda.c) plus a tiny -torch binding with ninja (torch.utils.cpp_extension), runs openfish_rmsnorm_gpu -and openfish_rmsnorm_quant_int8_gpu on the GPU, and diffs the result against a -torch reference -> PASS/FAIL. Also benchmarks each kernel with CUDA events. - - ./pyvenv/bin/python test/test_rmsnorm_cuda.py - -Requires a CUDA-enabled torch (see test/requirements.txt) and a CUDA toolkit. -(The fp8 fused kernel is not implemented on CUDA, so it is not covered here.) -""" - -import os -import sys - -# CUDA_HOME must be set before importing torch.utils.cpp_extension (it is -# resolved at import time). The ninja and nvcc binaries also need to be on PATH -# for the subprocess build (the venv's bin holds the pip-installed ninja). -CUDA_HOME = os.environ.setdefault("CUDA_HOME", "/usr/local/cuda") -os.environ["PATH"] = os.pathsep.join([ - os.path.dirname(sys.executable), - os.path.join(CUDA_HOME, "bin"), - os.environ.get("PATH", ""), -]) -# A100 = 8.0; override via env for other GPUs. -os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.0") - -import torch -from torch.utils.cpp_extension import load_inline - -ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - -# C++ binding (compiled by the host compiler). It calls the real -# openfish_rmsnorm_gpu / openfish_rmsnorm_quant_int8_gpu launch wrappers. -CPP_SRC = r""" -#include -#include -#include - -// out = rmsnorm(in + alpha*residual) * weight (fp32 math, fp16 out) -void rmsnorm(torch::Tensor in, torch::Tensor residual, torch::Tensor weight, - torch::Tensor out, double alpha, double eps) { - TORCH_CHECK(in.is_cuda() && residual.is_cuda() && weight.is_cuda() && out.is_cuda(), - "tensors must be CUDA"); - TORCH_CHECK(in.scalar_type() == torch::kHalf, "in must be fp16"); - const int n_tokens = in.size(0); - const int hidden_dim = in.size(1); - openfish_rmsnorm_gpu(in.data_ptr(), residual.data_ptr(), weight.data_ptr(), - out.data_ptr(), n_tokens, hidden_dim, (float)alpha, (float)eps); -} - -// In-place: residual (int8) and residual_scale (f32, per-token) are read as the -// previous quantized residual and overwritten with the new quantized rmsnorm. -void rmsnorm_quant_int8(torch::Tensor in, torch::Tensor weight, - torch::Tensor residual, torch::Tensor residual_scale, - double alpha, double eps) { - TORCH_CHECK(in.is_cuda() && weight.is_cuda() && residual.is_cuda() && residual_scale.is_cuda(), - "tensors must be CUDA"); - TORCH_CHECK(in.scalar_type() == torch::kHalf, "in must be fp16"); - TORCH_CHECK(residual.scalar_type() == torch::kChar, "residual must be int8"); - const int n_tokens = in.size(0); - const int hidden_dim = in.size(1); - openfish_rmsnorm_quant_int8_gpu(in.data_ptr(), weight.data_ptr(), residual.data_ptr(), - residual_scale.data_ptr(), n_tokens, hidden_dim, - (float)alpha, (float)eps); -} - -double rmsnorm_bench(torch::Tensor in, torch::Tensor residual, torch::Tensor weight, - torch::Tensor out, double alpha, double eps, - int64_t warmup, int64_t iters) { - const int n_tokens = in.size(0); - const int hidden_dim = in.size(1); - for (int i = 0; i < warmup; ++i) - openfish_rmsnorm_gpu(in.data_ptr(), residual.data_ptr(), weight.data_ptr(), - out.data_ptr(), n_tokens, hidden_dim, (float)alpha, (float)eps); - cudaDeviceSynchronize(); - cudaEvent_t start, stop; - cudaEventCreate(&start); cudaEventCreate(&stop); - cudaEventRecord(start); - for (int i = 0; i < iters; ++i) - openfish_rmsnorm_gpu(in.data_ptr(), residual.data_ptr(), weight.data_ptr(), - out.data_ptr(), n_tokens, hidden_dim, (float)alpha, (float)eps); - cudaEventRecord(stop); - cudaEventSynchronize(stop); - float ms = 0.0f; cudaEventElapsedTime(&ms, start, stop); - cudaEventDestroy(start); cudaEventDestroy(stop); - return (double)ms / (double)iters; -} - -double rmsnorm_quant_int8_bench(torch::Tensor in, torch::Tensor weight, - torch::Tensor residual, torch::Tensor residual_scale, - double alpha, double eps, int64_t warmup, int64_t iters) { - const int n_tokens = in.size(0); - const int hidden_dim = in.size(1); - for (int i = 0; i < warmup; ++i) - openfish_rmsnorm_quant_int8_gpu(in.data_ptr(), weight.data_ptr(), residual.data_ptr(), - residual_scale.data_ptr(), n_tokens, hidden_dim, - (float)alpha, (float)eps); - cudaDeviceSynchronize(); - cudaEvent_t start, stop; - cudaEventCreate(&start); cudaEventCreate(&stop); - cudaEventRecord(start); - for (int i = 0; i < iters; ++i) - openfish_rmsnorm_quant_int8_gpu(in.data_ptr(), weight.data_ptr(), residual.data_ptr(), - residual_scale.data_ptr(), n_tokens, hidden_dim, - (float)alpha, (float)eps); - cudaEventRecord(stop); - cudaEventSynchronize(stop); - float ms = 0.0f; cudaEventElapsedTime(&ms, start, stop); - cudaEventDestroy(start); cudaEventDestroy(stop); - return (double)ms / (double)iters; -} -""" - -# Compile the real kernel translation unit straight into the extension (nvcc, -# -fPIC) instead of linking lib/libopenfish.a, whose objects are built without -# -fPIC and so cannot go into a shared object. This still exercises the actual -# launch wrappers from src/nn_cuda.c. -CUDA_SRC = r""" -#include "error.c" -#include "nn_cuda.c" -""" - - -def ref_rmsnorm(in_, residual, weight, alpha, eps): - val = in_.to(torch.float32) + residual.to(torch.float32) * alpha - rms = torch.rsqrt(val.pow(2).mean(dim=-1, keepdim=True) + eps) - return (val * rms * weight.to(torch.float32)).to(torch.float16) - - -def ref_rmsnorm_quant_int8(in_, weight, residual, residual_scale, alpha, eps): - val = in_.to(torch.float32) + (residual.to(torch.float32) * residual_scale[:, None]) * alpha - rms = torch.rsqrt(val.pow(2).mean(dim=-1, keepdim=True) + eps) - normalized = val * rms * weight.to(torch.float32) - amax = normalized.abs().amax(dim=-1, keepdim=True) - quant_scale = torch.where(amax > 0, 127.0 / amax, torch.ones_like(amax)) - q = torch.clamp(torch.round(normalized * quant_scale), -127, 127).to(torch.int8) - new_scale = (1.0 / quant_scale).squeeze(-1) - return q, new_scale, normalized - - -def main(): - if not torch.cuda.is_available(): - sys.exit("CUDA torch not available; install torch with CUDA (see test/requirements.txt)") - - print(">> JIT-compiling CUDA kernels + binding with ninja") - mod = load_inline( - name="openfish_rmsnorm_test", - cpp_sources=CPP_SRC, - cuda_sources=CUDA_SRC, - functions=["rmsnorm", "rmsnorm_quant_int8", "rmsnorm_bench", "rmsnorm_quant_int8_bench"], - extra_include_paths=[os.path.join(ROOT, "include"), os.path.join(ROOT, "src")], - extra_cflags=["-DHAVE_CUDA=1", "-O2"], - extra_cuda_cflags=["-DHAVE_CUDA=1", "-O2", "--expt-relaxed-constexpr"], - with_cuda=True, - verbose=False, - ) - - dev = "cuda" - torch.manual_seed(0) - alpha, eps = 2.0, 1e-5 - ok = True - - # ---- correctness: rmsnorm ---- - n_tokens, hidden_dim = 512, 512 - in_ = torch.randn(n_tokens, hidden_dim, device=dev, dtype=torch.float16) - residual = torch.randn(n_tokens, hidden_dim, device=dev, dtype=torch.float16) - weight = torch.randn(hidden_dim, device=dev, dtype=torch.float16) - out = torch.empty_like(in_) - mod.rmsnorm(in_, residual, weight, out, alpha, eps) - torch.cuda.synchronize() - expected = ref_rmsnorm(in_, residual, weight, alpha, eps) - max_diff = (out.to(torch.float32) - expected.to(torch.float32)).abs().max().item() - tol = 5e-3 - print(f"[rmsnorm] max abs diff: {max_diff:.3e} (tol {tol:.1e})") - ok &= max_diff < tol - - # ---- correctness: rmsnorm_quant_int8 ---- - in8 = torch.randn(n_tokens, hidden_dim, device=dev, dtype=torch.float16) - weight8 = torch.randn(hidden_dim, device=dev, dtype=torch.float16) - residual8 = torch.randint(-127, 128, (n_tokens, hidden_dim), device=dev, dtype=torch.int8) - res_scale = torch.rand(n_tokens, device=dev, dtype=torch.float32) * 0.02 + 0.001 - q_ref, scale_ref, norm_ref = ref_rmsnorm_quant_int8( - in8, weight8, residual8.clone(), res_scale.clone(), alpha, eps) - q_got = residual8.clone() - scale_got = res_scale.clone() - mod.rmsnorm_quant_int8(in8, weight8, q_got, scale_got, alpha, eps) - torch.cuda.synchronize() - # Compare in dequantized space (a single int8 level near boundaries can differ - # by 1 due to fp rounding order); scale should match closely. - scale_diff = (scale_got - scale_ref).abs().max().item() - dq_got = q_got.to(torch.float32) * scale_got[:, None] - dq_ref = norm_ref - dq_diff = (dq_got - dq_ref).abs().max().item() - lvl_mismatch = (q_got.to(torch.int32) - q_ref.to(torch.int32)).abs() - lvl_bad = (lvl_mismatch > 1).float().mean().item() - # Correctness: per-token scale must match, and no quantized level may differ by - # more than 1 (a single-level delta is just fp round-to-nearest tie-breaking at a - # bin boundary; dq_diff is therefore ~one quant step by construction, not an error). - scale_tol = 1e-4 - print(f"[rmsnorm_quant_int8] scale max diff: {scale_diff:.3e} (tol {scale_tol:.1e}) " - f"dequant max diff: {dq_diff:.3e} (~1 quant step, informational) levels>1 off: {lvl_bad*100:.3f}%") - ok &= (scale_diff < scale_tol) and (lvl_bad == 0.0) - - print("PASS" if ok else "FAIL") - - # ---- timing ---- - print("\n>> timing (mean over 100 iters, after 20 warmup)") - print(f"{'kernel':>20} {'tokens':>8} {'hidden':>7} {'ms/call':>10} {'GB/s':>8}") - for n_tokens, hidden_dim in [(512, 512), (4096, 768), (8192, 1024)]: - in_ = torch.randn(n_tokens, hidden_dim, device=dev, dtype=torch.float16) - residual = torch.randn(n_tokens, hidden_dim, device=dev, dtype=torch.float16) - weight = torch.randn(hidden_dim, device=dev, dtype=torch.float16) - out = torch.empty_like(in_) - ms = mod.rmsnorm_bench(in_, residual, weight, out, alpha, eps, 20, 100) - # in + residual read (2B each) + out write (2B) + weight read (2B) - bytes_moved = n_tokens * hidden_dim * (2 + 2 + 2) + hidden_dim * 2 - gbps = bytes_moved / (ms * 1e-3) / 1e9 - print(f"{'rmsnorm':>20} {n_tokens:>8} {hidden_dim:>7} {ms:>10.4f} {gbps:>8.1f}") - - residual8 = torch.randint(-127, 128, (n_tokens, hidden_dim), device=dev, dtype=torch.int8) - res_scale = torch.rand(n_tokens, device=dev, dtype=torch.float32) * 0.02 + 0.001 - ms = mod.rmsnorm_quant_int8_bench(in_, weight, residual8, res_scale, alpha, eps, 20, 100) - # in read (2B) + residual read+write (1B each) + weight read (2B) - bytes_moved = n_tokens * hidden_dim * (2 + 1 + 1) + hidden_dim * 2 - gbps = bytes_moved / (ms * 1e-3) / 1e9 - print(f"{'rmsnorm_quant_int8':>20} {n_tokens:>8} {hidden_dim:>7} {ms:>10.4f} {gbps:>8.1f}") - - sys.exit(0 if ok else 1) - - -if __name__ == "__main__": - main() diff --git a/test/test_rotary_cuda.py b/test/test_rotary_cuda.py deleted file mode 100644 index 5f6e388..0000000 --- a/test/test_rotary_cuda.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Build-and-run test for the CUDA rotary embedding kernel. - -Single script: it JIT-compiles the real CUDA kernel (src/nn_cuda.c) plus a tiny -torch binding with ninja (torch.utils.cpp_extension), runs openfish_rotary_emb_gpu -on the GPU, and diffs the result against a torch reference -> PASS/FAIL. - - ./pyvenv/bin/python test/test_rotary_cuda.py - -Requires a CUDA-enabled torch (see test/requirements.txt) and a CUDA toolkit. -""" - -import os -import sys - -# CUDA_HOME must be set before importing torch.utils.cpp_extension (it is -# resolved at import time). The ninja and nvcc binaries also need to be on PATH -# for the subprocess build (the venv's bin holds the pip-installed ninja). -CUDA_HOME = os.environ.setdefault("CUDA_HOME", "/usr/local/cuda") -os.environ["PATH"] = os.pathsep.join([ - os.path.dirname(sys.executable), - os.path.join(CUDA_HOME, "bin"), - os.environ.get("PATH", ""), -]) -# A100 = 8.0; override via env for other GPUs. -os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.0") - -import torch -from torch.utils.cpp_extension import load_inline - -ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - -# TxModel.cpp constants (MultiHeadAttentionImpl / RotaryEmbeddingImpl). -THETA = 10000.0 -MAX_SEQ_LEN = 2048 -ROTARY_DIM = 32 - -# C++ binding (compiled by the host compiler). It calls the real -# openfish_rotary_emb_gpu for q (chunk 0) and k (chunk 1), passing the *full -# qkv* strides exactly as TxModel.cpp does. v (chunk 2) is untouched. -CPP_SRC = r""" -#include -#include -#include - -// Run the kernel on q (chunk 0) and k (chunk 1) of qkv, passing the *full qkv* -// strides exactly as TxModel.cpp does. v (chunk 2) is untouched. -// qkv: [N, seq_len, 3, n_heads, head_dim], fp16, CUDA, contiguous. -// sin/cos: [max_seq_len, head_dim/2], fp32, CUDA. -static void run_rotary(torch::Tensor &qkv, torch::Tensor &sin, torch::Tensor &cos, int rotary_dim) { - const int N = qkv.size(0); - const int seq_len = qkv.size(1); - const int n_heads = qkv.size(3); - const int head_dim = qkv.size(4); - const int sb = qkv.stride(0); - const int ss = qkv.stride(1); - const int sh = qkv.stride(3); - const int chunk_stride = qkv.stride(2); // = n_heads * head_dim when contiguous - - char *base = static_cast(qkv.data_ptr()); - const size_t elem = qkv.element_size(); - - // API arg order is (x, sin, cos). - for (int chunk = 0; chunk < 2; ++chunk) { // q, k - openfish_rotary_emb_gpu( - base + (size_t)chunk * chunk_stride * elem, - sin.data_ptr(), cos.data_ptr(), - N, seq_len, n_heads, head_dim, rotary_dim, sb, ss, sh); - } -} - -void rotary_emb(torch::Tensor qkv, torch::Tensor sin, torch::Tensor cos, int64_t rotary_dim) { - TORCH_CHECK(qkv.is_cuda() && sin.is_cuda() && cos.is_cuda(), "tensors must be CUDA"); - TORCH_CHECK(qkv.scalar_type() == torch::kHalf, "qkv must be fp16"); - TORCH_CHECK(qkv.dim() == 5 && qkv.size(2) == 3, "qkv must be [N,seq,3,nhead,head_dim]"); - run_rotary(qkv, sin, cos, (int)rotary_dim); -} - -// Returns mean milliseconds per rotary_emb call (q+k), timed with CUDA events. -double rotary_emb_bench(torch::Tensor qkv, torch::Tensor sin, torch::Tensor cos, - int64_t rotary_dim, int64_t warmup, int64_t iters) { - for (int i = 0; i < warmup; ++i) run_rotary(qkv, sin, cos, (int)rotary_dim); - cudaDeviceSynchronize(); - - cudaEvent_t start, stop; - cudaEventCreate(&start); - cudaEventCreate(&stop); - cudaEventRecord(start); - for (int i = 0; i < iters; ++i) run_rotary(qkv, sin, cos, (int)rotary_dim); - cudaEventRecord(stop); - cudaEventSynchronize(stop); - - float ms = 0.0f; - cudaEventElapsedTime(&ms, start, stop); - cudaEventDestroy(start); - cudaEventDestroy(stop); - return (double)ms / (double)iters; -} -""" - -# Compile the real kernel translation unit straight into the extension (nvcc, -# -fPIC) instead of linking lib/libopenfish.a, whose objects are built without -# -fPIC and so cannot go into a shared object. This still exercises the actual -# openfish_rotary_emb_gpu launch wrapper from src/nn_cuda.c. -CUDA_SRC = r""" -#include "error.c" -#include "nn_cuda.c" -""" - - -def build_cos_sin(head_dim, device): - inv_freq = torch.pow(THETA, torch.arange(0, head_dim, 2, dtype=torch.float64) / head_dim).reciprocal() - freqs = torch.outer(torch.arange(MAX_SEQ_LEN, dtype=torch.float64), inv_freq) - cos = torch.cos(freqs).to(torch.float32).contiguous().to(device) - sin = torch.sin(freqs).to(torch.float32).contiguous().to(device) - return cos, sin - - -def reference(qkv, cos, sin, rotary_dim): - """Rotate-half RoPE on q and k chunks; fp32 math, fp16 output (matches kernel).""" - out = qkv.clone() - seq_len = qkv.size(1) - c = cos[:seq_len, :rotary_dim].view(1, seq_len, 1, rotary_dim) - s = sin[:seq_len, :rotary_dim].view(1, seq_len, 1, rotary_dim) - for chunk in (0, 1): - x = out[:, :, chunk, :, :] - x0 = x[..., :rotary_dim].to(torch.float32) - x1 = x[..., rotary_dim:2 * rotary_dim].to(torch.float32) - x[..., :rotary_dim] = (x0 * c - x1 * s).to(torch.float16) - x[..., rotary_dim:2 * rotary_dim] = (x0 * s + x1 * c).to(torch.float16) - return out - - -def main(): - if not torch.cuda.is_available(): - sys.exit("CUDA torch not available; install torch with CUDA (see test/requirements.txt)") - - print(">> JIT-compiling CUDA kernel + binding with ninja") - mod = load_inline( - name="openfish_rotary_test", - cpp_sources=CPP_SRC, - cuda_sources=CUDA_SRC, - functions=["rotary_emb", "rotary_emb_bench"], - extra_include_paths=[os.path.join(ROOT, "include"), os.path.join(ROOT, "src")], - extra_cflags=["-DHAVE_CUDA=1", "-O2"], - extra_cuda_cflags=["-DHAVE_CUDA=1", "-O2", "--expt-relaxed-constexpr"], - with_cuda=True, - verbose=False, - ) - - dev = "cuda" - head_dim = 2 * ROTARY_DIM # = 64, matching the model - cos, sin = build_cos_sin(head_dim, dev) - - # ---- correctness ---- - torch.manual_seed(0) - batch, seq_len, n_heads = 4, 128, 8 - qkv = torch.randn(batch, seq_len, 3, n_heads, head_dim, device=dev).to(torch.float16).contiguous() - expected = reference(qkv, cos, sin, ROTARY_DIM) - - got = qkv.clone() - mod.rotary_emb(got, sin, cos, ROTARY_DIM) - torch.cuda.synchronize() - - # v chunk must be untouched. - assert torch.equal(got[:, :, 2], qkv[:, :, 2]), "v chunk was modified" - - max_diff = (got.to(torch.float32) - expected.to(torch.float32)).abs().max().item() - # fp16 rounding: a single fp16 ulp near magnitude ~4 is ~4e-3, and GPU - # FMA/rounding order can differ from torch by ~1 ulp, so allow a small margin. - tol = 5e-3 - print(f"kernel vs reference max abs diff: {max_diff:.3e} (tol {tol:.1e})") - ok = max_diff < tol - print("PASS" if ok else "FAIL") - - # ---- timing ---- - # bytes moved per call: q+k each read+write 2*rotary_dim halves per (N,seq,head). - print("\n>> timing (mean over 100 iters, after 20 warmup)") - print(f"{'N':>5} {'seq':>6} {'heads':>6} {'ms/call':>10} {'GB/s':>8}") - for batch, seq_len, n_heads in [(4, 128, 8), (64, 512, 8), (256, 1666, 8), (512, 2048, 8)]: - qkv = torch.randn(batch, seq_len, 3, n_heads, head_dim, device=dev).to(torch.float16).contiguous() - ms = mod.rotary_emb_bench(qkv, sin, cos, ROTARY_DIM, 20, 100) - # 2 chunks * (read+write) * 2*rotary_dim halves(2B) + cos/sin reads (2*rotary_dim*4B) - elems = batch * seq_len * n_heads - bytes_moved = 2 * elems * (2 * (2 * ROTARY_DIM) * 2 + (2 * ROTARY_DIM) * 4) - gbps = bytes_moved / (ms * 1e-3) / 1e9 - print(f"{batch:>5} {seq_len:>6} {n_heads:>6} {ms:>10.4f} {gbps:>8.1f}") - - sys.exit(0 if ok else 1) - - -if __name__ == "__main__": - main()