| 83 | } |
| 84 | |
| 85 | void layernorm_forward(Context& ctx, float* out, float* mean, float* rstd, |
| 86 | float* inp, float* weight, float* bias, |
| 87 | int B, int T, int C){ |
| 88 | unsigned long b = static_cast<unsigned long>(B); |
| 89 | unsigned long t = static_cast<unsigned long>(T); |
| 90 | unsigned long c = static_cast<unsigned long>(C); |
| 91 | struct LayerNormParams { |
| 92 | uint32_t B; |
| 93 | uint32_t T; |
| 94 | uint32_t C; |
| 95 | }; |
| 96 | setLogLevel(kError); |
| 97 | Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32, inp); |
| 98 | Tensor weight_t = createTensor(ctx, Shape{c}, kf32, weight); |
| 99 | Tensor bias_t = createTensor(ctx, Shape{c}, kf32, bias); |
| 100 | Tensor out_t = createTensor(ctx, Shape{b * t * c}, kf32); |
| 101 | Tensor mean_t = createTensor(ctx, Shape{b * t}, kf32); |
| 102 | Tensor rstd_t = createTensor(ctx, Shape{b * t}, kf32); |
| 103 | std::promise<void> promise; |
| 104 | std::future<void> future = promise.get_future(); |
| 105 | Kernel op = createKernel(ctx, {kShaderLayerNorm, 256, kf32}, |
| 106 | Bindings{inp_t, weight_t, bias_t, out_t, mean_t, rstd_t}, |
| 107 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 108 | /* params */ |
| 109 | LayerNormParams{ |
| 110 | static_cast<uint32_t>(b), |
| 111 | static_cast<uint32_t>(t), |
| 112 | static_cast<uint32_t>(c) |
| 113 | }); |
| 114 | dispatchKernel(ctx, op, promise); |
| 115 | wait(ctx, future); |
| 116 | toCPU(ctx, out_t, out, b * t * c * sizeof(float)); |
| 117 | toCPU(ctx, mean_t, mean, b * t * sizeof(float)); |
| 118 | toCPU(ctx, rstd_t, rstd, b * t * sizeof(float)); |
| 119 | } |
| 120 | |
| 121 | void layernorm_backward(Context& ctx, float* dinp, float* dweight, float* dbias, |
| 122 | float* dout, float* inp, float* weight, float* mean, float* rstd, |
no test coverage detected