| 157 | } |
| 158 | |
| 159 | void matmul_forward(Context& ctx, float* out, |
| 160 | const float* inp, const float* weight, const float* bias, |
| 161 | int B, int T, int C, int OC){ |
| 162 | struct MatmulParams { |
| 163 | uint32_t B; |
| 164 | uint32_t T; |
| 165 | uint32_t C; |
| 166 | uint32_t OC; |
| 167 | }; |
| 168 | unsigned long b = static_cast<unsigned long>(B); |
| 169 | unsigned long t = static_cast<unsigned long>(T); |
| 170 | unsigned long c = static_cast<unsigned long>(C); |
| 171 | unsigned long oc = static_cast<unsigned long>(OC); |
| 172 | setLogLevel(kError); |
| 173 | |
| 174 | Tensor inp_i = createTensor(ctx, Shape{b * t * c}, kf32, inp); |
| 175 | Tensor weight_i = createTensor(ctx, Shape{oc * c}, kf32, weight); |
| 176 | Tensor bias_i = bias == NULL ? createTensor(ctx, Shape{1}, kf32) : createTensor(ctx, Shape{oc}, kf32, bias); |
| 177 | Tensor out_o = createTensor(ctx, Shape{b * t * oc}, kf32); |
| 178 | std::promise<void> promise; |
| 179 | std::future<void> future = promise.get_future(); |
| 180 | assert ( (b*t) % 256 == 0 ); |
| 181 | Kernel op = createKernel(ctx, {kShaderMatmul, 256, kf32}, |
| 182 | Bindings{inp_i, weight_i, bias_i, out_o}, |
| 183 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 184 | /* params */ |
| 185 | MatmulParams{ |
| 186 | static_cast<uint32_t>(b), |
| 187 | static_cast<uint32_t>(t), |
| 188 | static_cast<uint32_t>(c), |
| 189 | static_cast<uint32_t>(oc) |
| 190 | }); |
| 191 | dispatchKernel(ctx, op, promise); |
| 192 | wait(ctx, future); |
| 193 | toCPU(ctx, out_o, out, b * t * oc * sizeof(float)); |
| 194 | } |
| 195 | |
| 196 | void matmul_backward(Context& ctx, float* dinp, float* dweight, float* dbias, |
| 197 | const float* dout, const float* inp, const float* weight, |
no test coverage detected