| 373 | } |
| 374 | |
| 375 | void softmax_forward(Context& ctx, float* probs, float* logits, int B, int T, int V, int Vp) { |
| 376 | struct SoftmaxParam { |
| 377 | uint32_t N; |
| 378 | uint32_t C; |
| 379 | uint32_t Cp; |
| 380 | }; |
| 381 | uint32_t b = static_cast<uint32_t>(B); |
| 382 | uint32_t t = static_cast<uint32_t>(T); |
| 383 | uint32_t c = static_cast<uint32_t>(V); |
| 384 | uint32_t cp = static_cast<uint32_t>(Vp); |
| 385 | Tensor input = createTensor(ctx, {b * t, cp}, kf32, logits); |
| 386 | Tensor output = createTensor(ctx, {b * t, cp}, kf32); |
| 387 | std::promise<void> promise; |
| 388 | std::future<void> future = promise.get_future(); |
| 389 | assert( (B*T) % 256 == 0); |
| 390 | Kernel op = createKernel( |
| 391 | ctx, {kShaderSoftmax1, 256, kf32}, Bindings{input, output}, |
| 392 | Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{b * t, c, cp}); |
| 393 | dispatchKernel(ctx, op, promise); |
| 394 | wait(ctx, future); |
| 395 | toCPU(ctx, output, probs, sizeof(float)*b*t*cp); |
| 396 | } |
| 397 | |
| 398 | void crossentropy_forward(Context& ctx, float* losses, |
| 399 | float* probs, int* targets, |
no test coverage detected