Module: Ignis::JIT::Kernels::Loss

Defined in:
lib/nvruby/jit/kernels/loss.rb

Overview

Loss function CUDA kernels for training. All are fused implementations for numerical stability and performance.

Class Method Summary collapse

Class Method Details

.bce_backward ⇒ Ignis::JIT::Kernel

BCE backward: σ(x) - y

Returns:



165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 165

def bce_backward
  source = "    extern \"C\" __global__\n    void bce_backward(const float* __restrict__ logits,\n                      const float* __restrict__ targets,\n                      const float* __restrict__ grad_output,\n                      float* __restrict__ grad_input,\n                      const int n) {\n      int idx = blockIdx.x * blockDim.x + threadIdx.x;\n      if (idx < n) {\n        float sig = 1.0f / (1.0f + expf(-logits[idx]));\n        grad_input[idx] = grad_output[idx] * (sig - targets[idx]);\n      }\n    }\n  CUDA\n  compile_cached(source, \"bce_backward\")\nend\n"

.bce_forward ⇒ Ignis::JIT::Kernel

Binary cross-entropy with logits: -[y*log(σ(x)) + (1-y)*log(1-σ(x))]

Returns:



143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 143

def bce_forward
  source = "    extern \"C\" __global__\n    void bce_forward(const float* __restrict__ logits,\n                     const float* __restrict__ targets,\n                     float* __restrict__ losses,\n                     const int n) {\n      int idx = blockIdx.x * blockDim.x + threadIdx.x;\n      if (idx < n) {\n        float x = logits[idx];\n        float y = targets[idx];\n        // Numerically stable: max(x,0) - x*y + log(1+exp(-|x|))\n        float max_val = fmaxf(x, 0.0f);\n        losses[idx] = max_val - x * y + logf(1.0f + expf(-fabsf(x)));\n      }\n    }\n  CUDA\n  compile_cached(source, \"bce_forward\")\nend\n"

.cross_entropy_backward ⇒ Ignis::JIT::Kernel

Cross-entropy backward: softmax(logits) - one_hot(target) Combined softmax + gradient in one kernel

Returns:



68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 68

def cross_entropy_backward
  source = "    extern \"C\" __global__\n    void cross_entropy_backward(const float* __restrict__ log_softmax,\n                                 const int* __restrict__ targets,\n                                 const float* __restrict__ grad_output,\n                                 float* __restrict__ grad_logits,\n                                 const int batch_size,\n                                 const int vocab_size,\n                                 const float label_smoothing) {\n      int idx = blockIdx.x * blockDim.x + threadIdx.x;\n      int total = batch_size * vocab_size;\n      if (idx < total) {\n        int row = idx / vocab_size;\n        int col = idx % vocab_size;\n        int target = targets[row];\n\n        float softmax_val = expf(log_softmax[idx]);\n        float grad_scale = grad_output[row];\n\n        if (label_smoothing > 0.0f) {\n          float smooth_target = label_smoothing / (float)vocab_size;\n          float hard_target = (col == target) ? (1.0f - label_smoothing + smooth_target) : smooth_target;\n          grad_logits[idx] = grad_scale * (softmax_val - hard_target);\n        } else {\n          float indicator = (col == target) ? 1.0f : 0.0f;\n          grad_logits[idx] = grad_scale * (softmax_val - indicator);\n        }\n      }\n    }\n  CUDA\n  compile_cached(source, \"cross_entropy_backward\")\nend\n"

.cross_entropy_forward ⇒ Ignis::JIT::Kernel

Fused cross-entropy forward: log_softmax + NLL in a single kernel Avoids materializing full log-softmax output

Returns:



13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 13

def cross_entropy_forward
  source = "    extern \"C\" __global__\n    void cross_entropy_forward(const float* __restrict__ logits,\n                                const int* __restrict__ targets,\n                                float* __restrict__ losses,\n                                float* __restrict__ log_softmax_out,\n                                const int batch_size,\n                                const int vocab_size,\n                                const float label_smoothing) {\n      int row = blockIdx.x * blockDim.x + threadIdx.x;\n      if (row < batch_size) {\n        const float* row_logits = logits + row * vocab_size;\n        float* row_lsm = log_softmax_out + row * vocab_size;\n        int target = targets[row];\n\n        // Find max for numerical stability\n        float max_val = row_logits[0];\n        for (int j = 1; j < vocab_size; j++) {\n          max_val = fmaxf(max_val, row_logits[j]);\n        }\n\n        // log_softmax = x - max - log(sum(exp(x - max)))\n        float log_sum_exp = 0.0f;\n        for (int j = 0; j < vocab_size; j++) {\n          log_sum_exp += expf(row_logits[j] - max_val);\n        }\n        log_sum_exp = logf(log_sum_exp);\n\n        // Compute log_softmax and store\n        for (int j = 0; j < vocab_size; j++) {\n          row_lsm[j] = row_logits[j] - max_val - log_sum_exp;\n        }\n\n        // NLL loss with optional label smoothing\n        if (label_smoothing > 0.0f) {\n          float smooth_loss = 0.0f;\n          for (int j = 0; j < vocab_size; j++) {\n            smooth_loss -= row_lsm[j];\n          }\n          smooth_loss /= (float)vocab_size;\n          float nll = -row_lsm[target];\n          losses[row] = (1.0f - label_smoothing) * nll + label_smoothing * smooth_loss;\n        } else {\n          losses[row] = -row_lsm[target];\n        }\n      }\n    }\n  CUDA\n  compile_cached(source, \"cross_entropy_forward\")\nend\n"

.mean_reduce ⇒ Ignis::JIT::Kernel

Mean reduction: compute mean of array

Returns:



185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 185

def mean_reduce
  source = "    extern \"C\" __global__\n    void mean_reduce(const float* __restrict__ input,\n                     float* __restrict__ output,\n                     const int n) {\n      // Single-thread simple reduction (for loss scalar)\n      if (blockIdx.x == 0 && threadIdx.x == 0) {\n        float sum = 0.0f;\n        for (int i = 0; i < n; i++) {\n          sum += input[i];\n        }\n        output[0] = sum / (float)n;\n      }\n    }\n  CUDA\n  compile_cached(source, \"mean_reduce\")\nend\n"

.mse_backward ⇒ Ignis::JIT::Kernel

MSE backward: 2 * (pred - target) / n

Returns:



123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 123

def mse_backward
  source = "    extern \"C\" __global__\n    void mse_backward(const float* __restrict__ predictions,\n                      const float* __restrict__ targets,\n                      const float* __restrict__ grad_output,\n                      float* __restrict__ grad_input,\n                      const int n,\n                      const float scale) {\n      int idx = blockIdx.x * blockDim.x + threadIdx.x;\n      if (idx < n) {\n        grad_input[idx] = grad_output[idx] * 2.0f * (predictions[idx] - targets[idx]) * scale;\n      }\n    }\n  CUDA\n  compile_cached(source, \"mse_backward\")\nend\n"

.mse_forward ⇒ Ignis::JIT::Kernel

MSE forward: (pred - target)^2, per element

Returns:



104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
 
# File 'lib/nvruby/jit/kernels/loss.rb', line 104

def mse_forward
  source = "    extern \"C\" __global__\n    void mse_forward(const float* __restrict__ predictions,\n                     const float* __restrict__ targets,\n                     float* __restrict__ losses,\n                     const int n) {\n      int idx = blockIdx.x * blockDim.x + threadIdx.x;\n      if (idx < n) {\n        float diff = predictions[idx] - targets[idx];\n        losses[idx] = diff * diff;\n      }\n    }\n  CUDA\n  compile_cached(source, \"mse_forward\")\nend\n"