/*Kernels for softmax forward pass. Compile example:nvcc -O3 --use_fast_math -lcublas -lcublasLt softmax_forward.cu -o softmax_forward version 1 is naive port from CPU code to kernel: parallelizes over B,T, loops over C./softmax_forward 1 version 2 is a fused kernel that parallelizes over all of B,T,C./softmax_forward 2 version 3 uses intra-warp reductions for maxval and sumval, must use block_size=32./softmax_forward 3 version 4 uses both intra-warp reductions and shared memory for inter-warp reductionsso it can tolerate any block_size % 32 == 0. this is hopefully the most efficient version./softmax_forward 4 version 5 is naive port from CPU code (softmax_online) to kernel: parallelizes over B,T, loops over C./softmax_forward 5 version 6 is softmax_online that parallelizes over all of B,T,C./softmax_forward 6 version 7 is softmax optimized for very large C../softmax_forward 7*/ #include <stdio.h>#include <stdlib.h>#include <assert.h>#include <cuda_runtime.h>#include <cooperative_groups.h>#include <cooperative_groups/reduce.h>#include "common.h" // ----------------------------------------------------------------------------// CPU code reference void softmax_forward_cpu(float* out, const float* inp, int N, int C) {    // inp is (N, C)    // out is (N, C), each row of inp will get softmaxed    for (int i = 0; i < N; i++) {        const float* inp_row = inp + i * C;        float* out_row = out + i * C;         float maxval = -INFINITY;        for (int j = 0; j < C; j++) {            if (inp_row[j] > maxval) {                maxval = inp_row[j];            }        }        // Note: since we want to ensure that the CUDA-kernels are accurate,        // we do this accumulation in higher precision, so we can be assured        // that our ground-truth is of high quality.        double sum = 0.0;        for (int j = 0; j < C; j++) {            out_row[j] = expf(inp_row[j] - maxval);            sum += out_row[j];        }        float norm = 1.f / (float)sum;        for (int j = 0; j < C; j++) {            out_row[j] *= norm;        }    }}  // online version of softmax on CPU from the paper "Online normalizer calculation for softmax"void softmax_forward_online_cpu(float* out, const float* inp, int N, int C) {    // inp is (N, C)    // out is (N, C), each row of inp will get softmaxed    for (int i = 0; i < N; i++) {        const float* inp_row = inp + i * C;        float* out_row = out + i * C;         float maxval = -INFINITY;        float sum = 0.0f;		for (int j = 0; j < C; j++) {			float maxval_prev = maxval;			if (inp_row[j] > maxval) {				maxval = inp_row[j];				sum = sum * expf(maxval_prev - maxval) + expf(inp_row[j] - maxval);			} else {				sum += expf(inp_row[j] - maxval);			}		}         for (int j = 0; j < C; j++) {            out_row[j] = expf(inp_row[j] - maxval) / sum;        }    }} // ----------------------------------------------------------------------------// GPU kernels __global__ void softmax_forward_kernel1(float* out, const float* inp, int N, int C) {    // inp is (N, C)    // out is (N, C), each row of inp will get softmaxed    int i = blockIdx.x * blockDim.x + threadIdx.x;    if (i < N) {        const float* inp_row = inp + i * C;        float* out_row = out + i * C;         float maxval = -INFINITY;        for (int j = 0; j < C; j++) {            if (inp_row[j] > maxval) {                maxval = inp_row[j];            }        }        double sum = 0.0;        for (int j = 0; j < C; j++) {            out_row[j] = expf(inp_row[j] - maxval);            sum += out_row[j];        }        for (int j = 0; j < C; j++) {            out_row[j] /= (float)sum;        }    }} __global__ void softmax_forward_kernel2(float* out, const float* inp, int N, int C) {    // inp is (N, C)    // in each row of C elements, first calculates maxval, then returns expf(val - maxval)    extern __shared__ float shared[];    int idx = blockIdx.x; // ranges [0, N)    int tid = threadIdx.x; // ranges [0, block_size)    int block_size = blockDim.x;    const float* x = inp + idx * C; // idx-th row of inp    // thread coarsening    float maxval = -INFINITY;    for (int i = tid; i < C; i += block_size) {        maxval = fmaxf(maxval, x[i]);    }    shared[tid] = maxval;    __syncthreads();    // reductions    for (int stride = block_size / 2; stride >= 1; stride /= 2) {        __syncthreads();        if (tid < stride) {            shared[tid] = fmaxf(shared[tid], shared[tid + stride]);        }    }    __syncthreads();    float offset = shared[0];    // compute expf and write the result to global memory    for (int i = tid; i < C; i += block_size) {        out[idx * C + i] = expf(x[i] - offset);    }    __syncthreads();    // thread coarsening again, for the sum    x = out + idx * C; // idx-th row of out    float sumval = 0.0f;    for (int i = tid; i < C; i += block_size) {        sumval += x[i];    }    shared[tid] = sumval;    __syncthreads();    // reductions    for (int stride = block_size / 2; stride >= 1; stride /= 2) {        __syncthreads();        if (tid < stride) {            shared[tid] += shared[tid + stride];        }    }    // broadcast the sum to all threads in the block    __syncthreads();    float sum = shared[0];    // divide the input values by the sum    for (int i = tid; i < C; i += block_size) {        out[idx * C + i] = x[i] / sum;    }} // warp-level reduction for finding the maximum value__device__ float warpReduceMax(float val) {    for (int offset = 16; offset > 0; offset /= 2) {        val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));    }    return val;} __global__ void softmax_forward_kernel3(float* out, const float* inp, int N, int C) {    // kernel must use block size of 32    extern __shared__ float shared[];    int idx = blockIdx.x;    int tid = threadIdx.x;    const float* x = inp + idx * C;     // Thread coarsening and within-warp reduction for maxval    float maxval = -INFINITY;    for (int i = tid; i < C; i += blockDim.x) {        maxval = fmaxf(maxval, x[i]);    }    maxval = warpReduceMax(maxval);     // Broadcast maxval within the warp    float offset = __shfl_sync(0xFFFFFFFF, maxval, 0);     // Compute expf and write the result to global memory    for (int i = tid; i < C; i += blockDim.x) {        out[idx * C + i] = expf(x[i] - offset);    }     // Thread coarsening and within-warp reduction for sumval    x = out + idx * C;    float sumval = 0.0f;    for (int i = tid; i < C; i += blockDim.x) {        sumval += x[i];    }    sumval = warpReduceSum(sumval);     // Broadcast sumval within the warp    float sum = __shfl_sync(0xFFFFFFFF, sumval, 0);     // Divide the input values by the sum    for (int i = tid; i < C; i += blockDim.x) {        out[idx * C + i] = x[i] / sum;    }} __global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {    // out is (N, C) just like inp. Each row of inp will get softmaxed.    // same as kernel3, but can handle any block size (multiple of 32)    // each row of C elements is handled by block_size threads    // furthermore, each block_size threads get executed in warps of 32 threads     // special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions    // shared memory is used for inter-warp reduction    extern __shared__ float shared[];    int idx = blockIdx.x;    int tid = threadIdx.x;    int warpId = threadIdx.x / 32; // warp index within a block    int laneId = threadIdx.x % 32; // thread index within a warp     // the number of warps per block. recall that blockDim.x is block_size    int warpsPerBlock = blockDim.x / 32;     // shared[] must be allocated to have 2 * warpsPerBlock elements    // first half for max values, the second half for sum values    float* maxvals = shared;    float* sumvals = &shared[warpsPerBlock];     // one row of inp, i.e. inp[idx, :] of shape (C,)    const float* x = inp + idx * C;     // first, thread coarsening by directly accessing global memory in series    float maxval = -INFINITY;    for (int i = tid; i < C; i += blockDim.x) {        maxval = fmaxf(maxval, x[i]);    }    // now within-warp reductions for maxval    maxval = warpReduceMax(maxval);     // the 0th thread of each warp writes the maxval of that warp to shared memory    if (laneId == 0) maxvals[warpId] = maxval;    __syncthreads();     // now the 0th thread reduces the maxvals in shared memory, i.e. across warps    if (tid == 0) {        float val = maxvals[tid];        for (int i = 1; i < warpsPerBlock; i++) {            val = fmaxf(val, maxvals[i]);        }        // store the final max in the first position        maxvals[0] = val;    }    __syncthreads();    // broadcast the max to all threads    float offset = maxvals[0];     // compute expf and write the result to global memory    for (int i = tid; i < C; i += blockDim.x) {        out[idx * C + i] = expf(x[i] - offset);    }     // okay now we calculated exp(x - max(x))    // step 2: sum all the values and divide by the sum     // thread coarsening for sum    x = out + idx * C;    float sumval = 0.0f;    for (int i = tid; i < C; i += blockDim.x) {        sumval += x[i];    }    // within-warp reduction for sumval    sumval = warpReduceSum(sumval);     // write sumval to shared memory    if (laneId == 0) sumvals[warpId] = sumval;    __syncthreads();     // inter-thread reduction of sum    if (tid == 0) {        float val = sumvals[tid];        for (int i = 1; i < warpsPerBlock; ++i) {            val += sumvals[i];        }        sumvals[0] = val;    }    __syncthreads();    // broadcast the sum to all threads    float sum = sumvals[0];     // divide the whole row by the sum    for (int i = tid; i < C; i += blockDim.x) {        out[idx * C + i] = x[i] / sum;    }} __global__ void softmax_forward_online_kernel1(float* out, const float* inp, int N, int C) {    // inp is (N, C)    // out is (N, C), each row of inp will get softmaxed    int i = blockIdx.x * blockDim.x + threadIdx.x;    if (i < N) {        const float* inp_row = inp + i * C;        float* out_row = out + i * C;         float maxval = -INFINITY;        double sum = 0.0;        for (int j = 0; j < C; j++) {            float maxval_prev = maxval;			if (inp_row[j] > maxval) {				maxval = inp_row[j];				sum = sum * expf(maxval_prev - maxval) + expf(inp_row[j] - maxval);			}			else {				sum += expf(inp_row[j] - maxval);			}		}         for (int j = 0; j < C; j++) {            out_row[j] = expf(inp_row[j] - maxval) / sum;        }    }} // struct for the reduction operation, guarantees 8-byte alignmentstruct __align__(8) SumMax{    float maxval;    float sum;}; // forceinline helps avoid function call overhead__device__ __forceinline__ SumMax reduce_sum_max_op(SumMax a, SumMax b) {    bool a_bigger = (a.maxval > b.maxval);    SumMax bigger_m = a_bigger ? a : b;    SumMax smaller_m = a_bigger ? b : a;    SumMax res;    res.maxval = bigger_m.maxval;    res.sum = bigger_m.sum + smaller_m.sum * expf(smaller_m.maxval - bigger_m.maxval);    return res;} __global__ void softmax_forward_online_kernel2(float* out, const float* inp, int N, int C) {	namespace cg = cooperative_groups;	cg::thread_block block = cg::this_thread_block();	cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);	int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();	if (idx >= N) {		return;	} 	// one row of inp, i.e. inp[idx, :] of shape (C,)	const float* x = inp + idx * C;     // base case for the reduction    SumMax sm_partial;	sm_partial.maxval = -INFINITY;	sm_partial.sum = 0.0f; 	// first, thread coarsening by directly accessing global memory in series	for (int i = warp.thread_rank(); i < C; i += warp.size()) {		sm_partial = reduce_sum_max_op(sm_partial, { x[i], 1.0f });	}     // second, the reduction	SumMax sm_total = cg::reduce(warp, sm_partial, reduce_sum_max_op); 	// divide the whole row by the sum	for (int i = warp.thread_rank(); i < C; i += warp.size()) {        // the below is equivalent to        // out[idx * C + i] = expf(x[i] - sm_total.maxval) / sm_total.sum;        // but uses special instruction that bypasses the cache        __stcs(out + idx * C + i, expf(x[i] - sm_total.maxval) / sm_total.sum);	}} __global__ void softmax_forward_kernel7(float* out, const float* inp, int N, int C) {    // out is (N, C) just like inp. Each row of inp will get softmaxed.    // same as kernel4, but optimised for very large Cs with advanced unrolling     // The trick is to read into a register array (all indices known at compile time)    // and always read UNROLL_FACTOR values to maximise memory level parallelism    // even if we would be out of bounds, we set the index to min(C-1, idx)    // so we just do some unnecessary reads (obviously bad for small C)    // the writes are in a separate loop with a conditional check for out of bounds    // making it separate is necessary to convince the compiler to do the right thing    const int UNROLL_FACTOR = 8;    const int warpsPerBlock = blockDim.x / 32;     extern __shared__ float shared[];    int idx = blockIdx.x;    int tid = threadIdx.x;    int warpId = threadIdx.x / 32; // warp index within a block    int laneId = threadIdx.x % 32; // thread index within a warp     // shared[] must be allocated to have 2 * warpsPerBlock elements    // first half for max values, the second half for sum values    float* maxvals = shared;    float* sumvals = &shared[warpsPerBlock];     if (tid >= C) {        maxvals[warpId] = -INFINITY;        sumvals[warpId] = 0.0f;        return;    }     const float* x = inp + idx * C; // input    float* y = out + idx * C; // output     // first, thread coarsening by directly accessing global memory in series    float maxval = -INFINITY;    for (int i = tid; i < C; i += blockDim.x * UNROLL_FACTOR) {        #pragma unroll        for (int u = 0; u < UNROLL_FACTOR; u++) {            maxval = fmaxf(maxval, x[min(C - 1, i + u*blockDim.x)]);        }    }     // now within-warp reductions for maxval    maxval = warpReduceMax(maxval);    // the 0th thread of each warp writes the maxval of that warp to shared memory    if (laneId == 0) maxvals[warpId] = maxval;    __syncthreads();    // now the 0th thread reduces the maxvals in shared memory, i.e. across warps    if (tid == 0) {        float val = maxvals[tid];        #pragma unroll        for (int i = 1; i < warpsPerBlock; i++) {            val = fmaxf(val, maxvals[i]);        }        // store the final max in the first position        maxvals[0] = val;    }    __syncthreads();    // broadcast the max to all threads    float offset = maxvals[0];     // compute expf and write the result to global memory    // + thread coarsening for sum    float sumval = 0.0f;    for (int i = tid; i < C; i += blockDim.x * UNROLL_FACTOR) {        float reg_array[UNROLL_FACTOR];        #pragma unroll        for (int u = 0; u < UNROLL_FACTOR; u++) {            reg_array[u] = __ldcs(&x[min(C - 1, i + u*blockDim.x)]);        }        #pragma unroll        for (int u = 0; u < UNROLL_FACTOR; u++) {            if (i + u*blockDim.x < C) {                float output = expf(reg_array[u] - offset);                y[min(C - 1, i + u*blockDim.x)] = output; // compiler likes redundant min()?!                sumval += output; // combined into the same loop unlike kernel3            }        }    }     // okay now we calculated exp(x - max(x))    // step 2: sum all the values and divide by the sum     // within-warp reduction for sumval    sumval = warpReduceSum(sumval);    // write sumval to shared memory    if (laneId == 0) sumvals[warpId] = sumval;    __syncthreads();    // inter-thread reduction of sum    if (tid == 0) {        float val = sumvals[tid];        #pragma unroll        for (int i = 1; i < warpsPerBlock; ++i) {            val += sumvals[i];        }        sumvals[0] = val;    }    __syncthreads();    // broadcast the sum to all threads    float sum = sumvals[0];     // divide the whole row by the sum    for (int i = tid; i < C; i += blockDim.x * UNROLL_FACTOR) {        float reg_array[UNROLL_FACTOR];        #pragma unroll        for (int u = 0; u < UNROLL_FACTOR; u++) {            reg_array[u] = y[min(C - 1, i + u*blockDim.x)];        }        #pragma unroll

Lines 1–500 of 735. Load more