| 396 | } |
| 397 | |
| 398 | void crossentropy_forward(Context& ctx, float* losses, |
| 399 | float* probs, int* targets, |
| 400 | int B, int T, int Vp){ |
| 401 | struct CrossEntropyParams { |
| 402 | uint32_t B; |
| 403 | uint32_t T; |
| 404 | uint32_t VP; |
| 405 | }; |
| 406 | unsigned long b = static_cast<unsigned long>(B); |
| 407 | unsigned long t = static_cast<unsigned long>(T); |
| 408 | unsigned long vp = static_cast<unsigned long>(Vp); |
| 409 | setLogLevel(kError); |
| 410 | Tensor losses_t = createTensor(ctx, Shape{b * t}, kf32, losses); |
| 411 | Tensor probs_t = createTensor(ctx, Shape{b * t * vp}, kf32, probs); |
| 412 | Tensor targets_t = createTensor(ctx, Shape{b * t}, ki32, targets); |
| 413 | std::promise<void> promise; |
| 414 | std::future<void> future = promise.get_future(); |
| 415 | Kernel op = createKernel(ctx, {kShaderCrossEntropyForward, 256, kf32}, |
| 416 | Bindings{losses_t, probs_t, targets_t}, |
| 417 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 418 | /* params */ |
| 419 | CrossEntropyParams{ |
| 420 | static_cast<uint32_t>(b), |
| 421 | static_cast<uint32_t>(t), |
| 422 | static_cast<uint32_t>(vp) |
| 423 | }); |
| 424 | dispatchKernel(ctx, op, promise); |
| 425 | wait(ctx, future); |
| 426 | toCPU(ctx, losses_t, losses, b * t * sizeof(float)); |
| 427 | } |
| 428 | |
| 429 | void crossentropy_softmax_backward(Context& ctx, float* dlogits, |
| 430 | float* dlosses, float* probs, int* targets, |
no test coverage detected