| 10 | using namespace gpu; |
| 11 | |
| 12 | void encoder_forward(Context& ctx, float* out, |
| 13 | int* inp, float* wte, float* wpe, |
| 14 | int B, int T, int C){ |
| 15 | unsigned long b = static_cast<unsigned long>(B); |
| 16 | unsigned long t = static_cast<unsigned long>(T); |
| 17 | unsigned long c = static_cast<unsigned long>(C); |
| 18 | unsigned long v = VOCAB_SIZE; |
| 19 | struct EncoderParams { |
| 20 | uint32_t B; |
| 21 | uint32_t T; |
| 22 | uint32_t C; |
| 23 | }; |
| 24 | setLogLevel(kError); |
| 25 | printf("Creating tensors\n"); |
| 26 | printf("Creating input tensor\%pn", inp); |
| 27 | Tensor input = createTensor(ctx, Shape{b * t}, ki32, inp); |
| 28 | printf("Created input tensor\n"); |
| 29 | Tensor wte_t = createTensor(ctx, Shape{v, c}, kf32, wte); |
| 30 | printf("Created wte tensor\n"); |
| 31 | Tensor wpe_t = createTensor(ctx, Shape{t, c}, kf32, wpe); |
| 32 | printf("Created wpe tensor\n"); |
| 33 | Tensor output = createTensor(ctx, Shape{b * t * c}, kf32); |
| 34 | printf("Created tensors\n"); |
| 35 | std::promise<void> promise; |
| 36 | std::future<void> future = promise.get_future(); |
| 37 | Kernel op = createKernel(ctx, {kShaderEncoder, 256, kf32}, |
| 38 | Bindings{input, wte_t, wpe_t, output}, |
| 39 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 40 | /* params */ |
| 41 | EncoderParams{ |
| 42 | static_cast<uint32_t>(b), |
| 43 | static_cast<uint32_t>(t), |
| 44 | static_cast<uint32_t>(c) |
| 45 | }); |
| 46 | dispatchKernel(ctx, op, promise); |
| 47 | wait(ctx, future); |
| 48 | toCPU(ctx, output, out, b * t * c * sizeof(float)); |
| 49 | } |
| 50 | |
| 51 | void encoder_backward(Context& ctx, float* dwte, float* dwpe, |
| 52 | float* dout, int* inp, |
no test coverage detected