loss_kernels.c File Reference

Loss function kernels (cross-entropy, etc.) More...

#include "ckernel_engine.h"
#include <math.h>
#include <stdlib.h>

Go to the source code of this file.

Functions
static int	ce_legacy_mode_enabled (void)
static int	ce_targets_all_valid_no_ignore (const int32_t *targets, int tokens, int vocab_size)
static int	parse_env_bool_on (const char *name)
void	softmax_cross_entropy_loss (const float *logits, const int32_t *targets, int tokens, int vocab_size, float *d_logits, float *loss_out)
static void	softmax_cross_entropy_loss_index_mean_impl (const float *logits, const int32_t *targets, int tokens, int vocab_size, float *d_logits, float *loss_out, int force_strict_math)
static void	softmax_cross_entropy_loss_legacy_mean_impl (const float *logits, const int32_t *targets, int tokens, int vocab_size, float *d_logits, float *loss_out)
void	softmax_cross_entropy_loss_ptref (const float *logits, const int32_t *targets, int tokens, int vocab_size, float *d_logits, float *loss_out)
static void	zero_row_f32 (float *row, int cols)

Detailed Description

Loss function kernels (cross-entropy, etc.)

CK-ENGINE KERNEL RULES:

1. NO malloc/free - memory via bump allocator, pointers passed in
2. NO OpenMP - parallelization at orchestrator/codegen layer
3. API must define: inputs, outputs, workspace, and memory layouts
4. Pure computation - deterministic, no side effects

After changes: make test && make llamacpp-parity-full

Cross-entropy: L = -log(softmax(logits)[target])

Definition in file loss_kernels.c.

Function Documentation

ce_legacy_mode_enabled()

static int ce_legacy_mode_enabled ( void  )

static

Definition at line 46 of file loss_kernels.c.

{
    static int cached = -1;
    if (cached < 0) {
        cached = parse_env_bool_on("CK_CE_LEGACY_MODE") ? 1 : 0;
    }
    return cached;
}

References parse_env_bool_on().

Referenced by softmax_cross_entropy_loss().

ce_targets_all_valid_no_ignore()

static int ce_targets_all_valid_no_ignore ( const int32_t *  targets, int  tokens, int  vocab_size  )

static

Definition at line 248 of file loss_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const int target = targets[t];
        if (target < 0 || target >= vocab_size) {
            return 0;
        }
    }
    return 1;
}

References vocab_size.

Referenced by softmax_cross_entropy_loss().

parse_env_bool_on()

static int parse_env_bool_on ( const char *  name )

static

Definition at line 29 of file loss_kernels.c.

{
    const char *v = getenv(name);
    if (!v || !v[0]) {
        return 0;
    }
    if (v[0] == '1' || v[0] == 'y' || v[0] == 'Y' || v[0] == 't' || v[0] == 'T') {
        return 1;
    }
    return 0;
}

Referenced by ce_legacy_mode_enabled().

softmax_cross_entropy_loss()

void softmax_cross_entropy_loss	(	const float *	logits,
		const int32_t *	targets,
		int	tokens,
		int	vocab_size,
		float *	d_logits,
		float *	loss_out
	)

Definition at line 259 of file loss_kernels.c.

{
    if (!logits || !targets || !d_logits || tokens <= 0 || vocab_size <= 0) {
        if (loss_out) {
            *loss_out = 0.0f;
        }
        return;
    }
 
    if (ce_legacy_mode_enabled() && ce_targets_all_valid_no_ignore(targets, tokens, vocab_size)) {
        softmax_cross_entropy_loss_legacy_mean_impl(
            logits, targets, tokens, vocab_size, d_logits, loss_out);
        return;
    }
 
    softmax_cross_entropy_loss_index_mean_impl(
        logits, targets, tokens, vocab_size, d_logits, loss_out, -1);
}

References ce_legacy_mode_enabled(), ce_targets_all_valid_no_ignore(), softmax_cross_entropy_loss_index_mean_impl(), softmax_cross_entropy_loss_legacy_mean_impl(), and vocab_size.

Referenced by softmax_cross_entropy_loss_bf16().

softmax_cross_entropy_loss_index_mean_impl()

static void softmax_cross_entropy_loss_index_mean_impl ( const float *  logits, const int32_t *  targets, int  tokens, int  vocab_size, float *  d_logits, float *  loss_out, int  force_strict_math  )

static

Definition at line 64 of file loss_kernels.c.

{
    if (!logits || !targets || !d_logits || tokens <= 0 || vocab_size <= 0) {
        if (loss_out) {
            *loss_out = 0.0f;
        }
        return;
    }
 
    const int ignore_index = -100;
    const int strict = (force_strict_math >= 0) ? (force_strict_math ? 1 : 0)
                                                : ck_strict_parity_enabled();
    double total_loss = 0.0;
    int valid_tokens = 0;
    int invalid_target_seen = 0;
 
    for (int t = 0; t < tokens; ++t) {
        const float *row = logits + (size_t)t * (size_t)vocab_size;
        float *drow = d_logits + (size_t)t * (size_t)vocab_size;
        const int target = targets[t];
 
        if (target == ignore_index) {
            zero_row_f32(drow, vocab_size);
            continue;
        }
        if (target < 0 || target >= vocab_size) {
            zero_row_f32(drow, vocab_size);
            invalid_target_seen = 1;
            continue;
        }
 
        if (strict) {
            double max_logit = (double)row[0];
            for (int v = 1; v < vocab_size; ++v) {
                const double rv = (double)row[v];
                if (rv > max_logit) {
                    max_logit = rv;
                }
            }
 
            double sum_exp = 0.0;
            for (int v = 0; v < vocab_size; ++v) {
                sum_exp += exp((double)row[v] - max_logit);
            }
            const double inv_sum = 1.0 / sum_exp;
            for (int v = 0; v < vocab_size; ++v) {
                drow[v] = (float)(exp((double)row[v] - max_logit) * inv_sum);
            }
 
            total_loss += -(double)row[target] + max_logit + log(sum_exp);
        } else {
            float max_logit = row[0];
            for (int v = 1; v < vocab_size; ++v) {
                if (row[v] > max_logit) {
                    max_logit = row[v];
                }
            }
 
            double sum_exp = 0.0;
            for (int v = 0; v < vocab_size; ++v) {
                const float e = expf(row[v] - max_logit);
                drow[v] = e;
                sum_exp += (double)e;
            }
            const float inv_sum = 1.0f / (float)sum_exp;
            for (int v = 0; v < vocab_size; ++v) {
                drow[v] *= inv_sum;
            }
 
            total_loss += -(double)row[target] + (double)max_logit + log(sum_exp);
        }
 
        drow[target] -= 1.0f;
        ++valid_tokens;
    }
 
    if (valid_tokens > 0) {
        const float scale = 1.0f / (float)valid_tokens;
        for (int t = 0; t < tokens; ++t) {
            float *drow = d_logits + (size_t)t * (size_t)vocab_size;
            for (int v = 0; v < vocab_size; ++v) {
                drow[v] *= scale;
            }
        }
    }
 
    if (loss_out) {
        if (invalid_target_seen || valid_tokens == 0) {
            *loss_out = NAN;
        } else {
            *loss_out = (float)(total_loss / (double)valid_tokens);
        }
    }
}

References ck_strict_parity_enabled(), vocab_size, and zero_row_f32().

Referenced by softmax_cross_entropy_loss(), and softmax_cross_entropy_loss_ptref().

softmax_cross_entropy_loss_legacy_mean_impl()

static void softmax_cross_entropy_loss_legacy_mean_impl ( const float *  logits, const int32_t *  targets, int  tokens, int  vocab_size, float *  d_logits, float *  loss_out  )

static

Definition at line 169 of file loss_kernels.c.

{
    const int strict = ck_strict_parity_enabled();
    const double scale = 1.0 / (double)tokens;
    double total_loss = 0.0;
 
    for (int t = 0; t < tokens; ++t) {
        const float *row = logits + (size_t)t * (size_t)vocab_size;
        float *drow = d_logits + (size_t)t * (size_t)vocab_size;
        const int target = targets[t];
 
        if (strict) {
            double max_logit = (double)row[0];
            for (int v = 1; v < vocab_size; ++v) {
                const double rv = (double)row[v];
                if (rv > max_logit) {
                    max_logit = rv;
                }
            }
 
            double sum_exp = 0.0;
            for (int v = 0; v < vocab_size; ++v) {
                sum_exp += exp((double)row[v] - max_logit);
            }
            const double inv_sum = 1.0 / sum_exp;
            for (int v = 0; v < vocab_size; ++v) {
                const double p = exp((double)row[v] - max_logit) * inv_sum;
                drow[v] = (float)(p * scale);
            }
 
            if (target >= 0 && target < vocab_size) {
                const double log_sum_exp = log(sum_exp);
                const double target_logit = (double)row[target];
                total_loss += -(target_logit - max_logit - log_sum_exp);
                drow[target] -= (float)scale;
            }
        } else {
            float max_logit = row[0];
            for (int v = 1; v < vocab_size; ++v) {
                if (row[v] > max_logit) {
                    max_logit = row[v];
                }
            }
 
            double sum_exp = 0.0;
            for (int v = 0; v < vocab_size; ++v) {
                const float e = expf(row[v] - max_logit);
                drow[v] = e;
                sum_exp += (double)e;
            }
 
            const float inv_sum = 1.0f / (float)sum_exp;
            for (int v = 0; v < vocab_size; ++v) {
                drow[v] *= inv_sum;
            }
 
            if (target >= 0 && target < vocab_size) {
                const double log_sum_exp = log(sum_exp);
                const double target_logit = (double)row[target];
                total_loss += -(target_logit - (double)max_logit - log_sum_exp);
                drow[target] -= 1.0f;
            }
 
            for (int v = 0; v < vocab_size; ++v) {
                drow[v] *= (float)scale;
            }
        }
    }
 
    if (loss_out) {
        *loss_out = (float)(total_loss / (double)tokens);
    }
}

References ck_strict_parity_enabled(), and vocab_size.

Referenced by softmax_cross_entropy_loss().

softmax_cross_entropy_loss_ptref()

void softmax_cross_entropy_loss_ptref	(	const float *	logits,
		const int32_t *	targets,
		int	tokens,
		int	vocab_size,
		float *	d_logits,
		float *	loss_out
	)

Definition at line 292 of file loss_kernels.c.

{
    /*
     * Keep a strict reference variant for parity experiments:
     * - always uses strict math path,
     * - same reduction / ignore semantics as default kernel.
     */
    softmax_cross_entropy_loss_index_mean_impl(
        logits, targets, tokens, vocab_size, d_logits, loss_out, 1);
}

References softmax_cross_entropy_loss_index_mean_impl(), and vocab_size.

zero_row_f32()

static void zero_row_f32 ( float *  row, int  cols  )

static

Definition at line 22 of file loss_kernels.c.

{
    for (int i = 0; i < cols; ++i) {
        row[i] = 0.0f;
    }
}

Referenced by softmax_cross_entropy_loss_index_mean_impl().