| 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, |
| 123 | int B, int T, int C){ |
| 124 | unsigned long b = static_cast<unsigned long>(B); |
| 125 | unsigned long t = static_cast<unsigned long>(T); |
| 126 | unsigned long c = static_cast<unsigned long>(C); |
| 127 | struct LayerNormParams { |
| 128 | uint32_t B; |
| 129 | uint32_t T; |
| 130 | uint32_t C; |
| 131 | }; |
| 132 | setLogLevel(kError); |
| 133 | Tensor dinp_t = createTensor(ctx, Shape{b * t * c}, kf32, dinp); |
| 134 | Tensor dweight_t = createTensor(ctx, Shape{c}, kf32, dweight); |
| 135 | Tensor dbias_t = createTensor(ctx, Shape{c}, kf32, dbias); |
| 136 | Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32, dout); |
| 137 | Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32, inp); |
| 138 | Tensor weight_t = createTensor(ctx, Shape{c}, kf32, weight); |
| 139 | Tensor mean_t = createTensor(ctx, Shape{b * t}, kf32, mean); |
| 140 | Tensor rstd_t = createTensor(ctx, Shape{b * t}, kf32, rstd); |
| 141 | std::promise<void> promise; |
| 142 | std::future<void> future = promise.get_future(); |
| 143 | Kernel op = createKernel(ctx, {kShaderLayerNormBackward, 256, kf32}, |
| 144 | Bindings{dinp_t, dweight_t, dbias_t, dout_t, inp_t, weight_t, mean_t, rstd_t}, |
| 145 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 146 | /* params */ |
| 147 | LayerNormParams{ |
| 148 | static_cast<uint32_t>(b), |
| 149 | static_cast<uint32_t>(t), |
| 150 | static_cast<uint32_t>(c) |
| 151 | }); |
| 152 | dispatchKernel(ctx, op, promise); |
| 153 | wait(ctx, future); |
| 154 | toCPU(ctx, dinp_t, dinp, b * t * c * sizeof(float)); |
| 155 | toCPU(ctx, dweight_t, dweight, c * sizeof(float)); |
| 156 | toCPU(ctx, dbias_t, dbias, c * sizeof(float)); |
| 157 | } |
| 158 | |
| 159 | void matmul_forward(Context& ctx, float* out, |
| 160 | const float* inp, const float* weight, const float* bias, |
no test coverage detected