|
| 1 | +#include "opt-step.hpp" |
| 2 | + |
| 3 | +#define SYCL_OPT_STEP_BLOCK_SIZE 256 |
| 4 | + |
| 5 | +template <typename T> |
| 6 | +static void opt_step_adamw_f32_kernel( |
| 7 | + T * __restrict__ x, |
| 8 | + const T * __restrict__ g, |
| 9 | + T * __restrict__ g_m, |
| 10 | + T * __restrict__ g_v, |
| 11 | + const T * __restrict__ pars, |
| 12 | + const int64_t k, |
| 13 | + const sycl::nd_item<1> & item) { |
| 14 | + |
| 15 | + const int64_t i = (int64_t) item.get_global_id(0); |
| 16 | + if (i >= k) { |
| 17 | + return; |
| 18 | + } |
| 19 | + |
| 20 | + const float alpha = pars[0]; |
| 21 | + const float beta1 = pars[1]; |
| 22 | + const float beta2 = pars[2]; |
| 23 | + const float eps = pars[3]; |
| 24 | + const float wd = pars[4]; |
| 25 | + const float beta1h = pars[5]; |
| 26 | + const float beta2h = pars[6]; |
| 27 | + |
| 28 | + const float gi = g[i]; |
| 29 | + const float gmi = g_m[i] * beta1 + gi * (1.0f - beta1); |
| 30 | + const float gvi = g_v[i] * beta2 + gi * gi * (1.0f - beta2); |
| 31 | + |
| 32 | + g_m[i] = gmi; |
| 33 | + g_v[i] = gvi; |
| 34 | + |
| 35 | + const float mh = gmi * beta1h; |
| 36 | + const float vh = sycl::sqrt(gvi * beta2h) + eps; |
| 37 | + |
| 38 | + x[i] = x[i] * (1.0f - alpha * wd) - alpha * mh / vh; |
| 39 | +} |
| 40 | + |
| 41 | +template <typename T> |
| 42 | +static void opt_step_sgd_f32_kernel( |
| 43 | + T * __restrict__ x, |
| 44 | + const T * __restrict__ g, |
| 45 | + const T * __restrict__ pars, |
| 46 | + const int64_t k, |
| 47 | + const sycl::nd_item<1> & item) { |
| 48 | + |
| 49 | + const int64_t i = (int64_t) item.get_global_id(0); |
| 50 | + if (i >= k) { |
| 51 | + return; |
| 52 | + } |
| 53 | + |
| 54 | + x[i] = x[i] * (1.0f - pars[0] * pars[1]) - pars[0] * g[i]; |
| 55 | +} |
| 56 | + |
| 57 | +void ggml_sycl_opt_step_adamw(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { |
| 58 | + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/5); |
| 59 | + |
| 60 | + const ggml_tensor * src0 = dst->src[0]; |
| 61 | + const ggml_tensor * src0_grad = dst->src[1]; |
| 62 | + const ggml_tensor * src0_grad_m = dst->src[2]; |
| 63 | + const ggml_tensor * src0_grad_v = dst->src[3]; |
| 64 | + const ggml_tensor * adamw_params = dst->src[4]; |
| 65 | + |
| 66 | + GGML_ASSERT(src0->type == GGML_TYPE_F32); |
| 67 | + GGML_ASSERT(src0_grad->type == GGML_TYPE_F32); |
| 68 | + GGML_ASSERT(src0_grad_m->type == GGML_TYPE_F32); |
| 69 | + GGML_ASSERT(src0_grad_v->type == GGML_TYPE_F32); |
| 70 | + GGML_ASSERT(adamw_params->type == GGML_TYPE_F32); |
| 71 | + GGML_ASSERT(ggml_is_contiguous(src0)); |
| 72 | + GGML_ASSERT(ggml_is_contiguous(src0_grad)); |
| 73 | + GGML_ASSERT(ggml_is_contiguous(src0_grad_m)); |
| 74 | + GGML_ASSERT(ggml_is_contiguous(src0_grad_v)); |
| 75 | + GGML_ASSERT(ggml_is_contiguous(adamw_params)); |
| 76 | + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); |
| 77 | + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_m)); |
| 78 | + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_v)); |
| 79 | + GGML_ASSERT(ggml_nelements(adamw_params) == 7); |
| 80 | + |
| 81 | + dpct::queue_ptr stream = ctx.stream(); |
| 82 | + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); |
| 83 | + |
| 84 | + float * src0_d = (float *) src0->data; |
| 85 | + const float * src0_grad_d = (const float *) src0_grad->data; |
| 86 | + float * src0_grad_m_d = (float *) src0_grad_m->data; |
| 87 | + float * src0_grad_v_d = (float *) src0_grad_v->data; |
| 88 | + const float * adamw_params_d = (const float *) adamw_params->data; |
| 89 | + |
| 90 | + const int64_t ne = ggml_nelements(src0); |
| 91 | + const int64_t num_blocks = (ne + SYCL_OPT_STEP_BLOCK_SIZE - 1) / SYCL_OPT_STEP_BLOCK_SIZE; |
| 92 | + |
| 93 | + stream->parallel_for( |
| 94 | + sycl::nd_range<1>(num_blocks * SYCL_OPT_STEP_BLOCK_SIZE, SYCL_OPT_STEP_BLOCK_SIZE), |
| 95 | + [=](sycl::nd_item<1> item) { |
| 96 | + opt_step_adamw_f32_kernel(src0_d, src0_grad_d, src0_grad_m_d, src0_grad_v_d, adamw_params_d, ne, item); |
| 97 | + }); |
| 98 | +} |
| 99 | + |
| 100 | +void ggml_sycl_opt_step_sgd(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { |
| 101 | + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3); |
| 102 | + |
| 103 | + const ggml_tensor * src0 = dst->src[0]; |
| 104 | + const ggml_tensor * src0_grad = dst->src[1]; |
| 105 | + const ggml_tensor * sgd_params = dst->src[2]; |
| 106 | + |
| 107 | + GGML_ASSERT(src0->type == GGML_TYPE_F32); |
| 108 | + GGML_ASSERT(src0_grad->type == GGML_TYPE_F32); |
| 109 | + GGML_ASSERT(sgd_params->type == GGML_TYPE_F32); |
| 110 | + GGML_ASSERT(ggml_is_contiguous(src0)); |
| 111 | + GGML_ASSERT(ggml_is_contiguous(src0_grad)); |
| 112 | + GGML_ASSERT(ggml_is_contiguous(sgd_params)); |
| 113 | + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); |
| 114 | + GGML_ASSERT(ggml_nelements(sgd_params) == 2); |
| 115 | + |
| 116 | + dpct::queue_ptr stream = ctx.stream(); |
| 117 | + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); |
| 118 | + |
| 119 | + float * src0_d = (float *) src0->data; |
| 120 | + const float * src0_grad_d = (const float *) src0_grad->data; |
| 121 | + const float * sgd_params_d = (const float *) sgd_params->data; |
| 122 | + |
| 123 | + const int64_t ne = ggml_nelements(src0); |
| 124 | + const int64_t num_blocks = (ne + SYCL_OPT_STEP_BLOCK_SIZE - 1) / SYCL_OPT_STEP_BLOCK_SIZE; |
| 125 | + |
| 126 | + stream->parallel_for( |
| 127 | + sycl::nd_range<1>(num_blocks * SYCL_OPT_STEP_BLOCK_SIZE, SYCL_OPT_STEP_BLOCK_SIZE), |
| 128 | + [=](sycl::nd_item<1> item) { |
| 129 | + opt_step_sgd_f32_kernel(src0_d, src0_grad_d, sgd_params_d, ne, item); |
| 130 | + }); |
| 131 | +} |
0 commit comments