| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #pragma once |
| 7 | #include "vec.h" |
| 8 | |
| 9 | template <typename scalar_t> |
| 10 | inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) { |
| 11 | using Vec = at::vec::Vectorized<scalar_t>; |
| 12 | const Vec data_vec(val); |
| 13 | at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size); |
| 14 | } |
| 15 | |
| 16 | template <typename scalar_t> |
| 17 | inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { |
| 18 | using Vec = at::vec::Vectorized<scalar_t>; |
| 19 | constexpr int kVecSize = Vec::size(); |
| 20 | int64_t d; |
| 21 | #pragma GCC unroll 4 |
| 22 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 23 | Vec data = Vec::loadu(input + d); |
| 24 | data.store(out + d); |
| 25 | } |
| 26 | for (; d < size; ++d) { |
| 27 | out[d] = input[d]; |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | template <typename scalar_t> |
| 32 | inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) { |
| 33 | using bVec = at::vec::Vectorized<scalar_t>; |
| 34 | using fVec = at::vec::Vectorized<float>; |
| 35 | constexpr int kVecSize = bVec::size(); |
| 36 | int64_t d; |
| 37 | #pragma GCC unroll 4 |
| 38 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 39 | auto [x0, x1] = load_float_vec2(input + d); |
| 40 | bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1); |
| 41 | out_vec.store(out + d); |
| 42 | } |
| 43 | for (; d < size; ++d) { |
| 44 | out[d] = static_cast<scalar_t>(input[d]); |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | template <> |
| 49 | inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) { |
| 50 | // size might be 64x + 32 |
| 51 | std::memcpy(out, input, size * sizeof(uint8_t)); |
| 52 | } |
| 53 | |
| 54 | template <typename scalar_t, typename input_t> |
| 55 | inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict__ input, float weight, int64_t size) { |
| 56 | static_assert( |
| 57 | std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>, |
| 58 | "copy_mul_stub only supports input_t == float or input_t == scalar_t"); |
| 59 | using bVec = at::vec::Vectorized<scalar_t>; |
| 60 | using fVec = at::vec::Vectorized<float>; |
| 61 | constexpr int kVecSize = bVec::size(); |
| 62 | const fVec weight_vec = fVec(weight); |
| 63 | int64_t d; |
| 64 | #pragma GCC unroll 4 |
| 65 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 66 | auto [x0, x1] = load_float_vec2(input + d); |
| 67 | bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec); |
| 68 | out_vec.store(out + d); |
| 69 | } |
| 70 | for (; d < size; ++d) { |
| 71 | out[d] = static_cast<scalar_t>(input[d] * weight); |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | // acc from [topk, K] to [K] |
| 76 | template <typename scalar_t> |
| 77 | inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { |
| 78 | using bVec = at::vec::Vectorized<scalar_t>; |
| 79 | using fVec = at::vec::Vectorized<float>; |
| 80 | constexpr int kVecSize = bVec::size(); |
| 81 | if (topk == 1) { |
| 82 | // do copy for topk = 1 |
| 83 | copy_stub(out, input, K); |
| 84 | } else { |
| 85 | // do sum for topk != 1 |
| 86 | int64_t d; |
| 87 | #pragma GCC unroll 4 |
| 88 | for (d = 0; d <= K - kVecSize; d += kVecSize) { |
| 89 | fVec sum_fvec0 = fVec(0.f); |
| 90 | fVec sum_fvec1 = fVec(0.f); |
| 91 | for (int t = 0; t < topk; ++t) { |
| 92 | auto [x_fvec0, x_fvec1] = load_float_vec2(input + t * K + d); |
| 93 | sum_fvec0 += x_fvec0; |
| 94 | sum_fvec1 += x_fvec1; |
| 95 | } |
| 96 | bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1); |
| 97 | out_bvec.store(out + d); |
| 98 | } |
| 99 | for (; d < K; ++d) { |
| 100 | float sum_val = 0.f; |
| 101 | for (int t = 0; t < topk; ++t) { |
| 102 | sum_val += static_cast<float>(input[t * K + d]); |
| 103 | } |
| 104 | out[d] = static_cast<scalar_t>(sum_val); |
| 105 | } |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | // out = input + input2 * scale |
| 110 | template <typename scalar_t, typename input_t> |
| 111 | inline void add_mul_stub( |
| 112 | scalar_t* __restrict__ out, |
| 113 | const input_t* __restrict__ input, |
| 114 | const scalar_t* __restrict__ input2, |
| 115 | float scale, |
| 116 | int64_t size) { |
| 117 | static_assert( |
| 118 | std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>, |
| 119 | "add_mul_stub only supports input_t == float or input_t == scalar_t"); |
| 120 | |
| 121 | // out = input (without scale factor) |
| 122 | if (input2 == nullptr) { |
| 123 | copy_stub(out, input, size); |
| 124 | return; |
| 125 | } |
| 126 | |
| 127 | using bVec = at::vec::Vectorized<scalar_t>; |
| 128 | using fVec = at::vec::Vectorized<float>; |
| 129 | constexpr int kVecSize = bVec::size(); |
| 130 | const fVec s_vec = fVec(scale); |
| 131 | int64_t d; |
| 132 | #pragma GCC unroll 4 |
| 133 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 134 | auto [x0, x1] = load_float_vec2(input + d); |
| 135 | auto [y0, y1] = load_float_vec2(input2 + d); |
| 136 | x0 = x0 + y0 * s_vec; |
| 137 | x1 = x1 + y1 * s_vec; |
| 138 | bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1); |
| 139 | out_vec.store(out + d); |
| 140 | } |
| 141 | for (; d < size; ++d) { |
| 142 | out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale); |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | template <typename scalar_t, typename input_t> |
| 147 | inline void silu_and_mul_stub( |
| 148 | scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) { |
| 149 | static_assert( |
| 150 | std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>, |
| 151 | "silu_and_mul_stub only supports input_t == float or input_t == scalar_t"); |
| 152 | using bVec = at::vec::Vectorized<scalar_t>; |
| 153 | using fVec = at::vec::Vectorized<float>; |
| 154 | |
| 155 | // no remainder |
| 156 | #pragma GCC unroll 4 |
| 157 | for (int64_t d = 0; d < size; d += bVec::size()) { |
| 158 | auto [x0, x1] = load_float_vec2(input + d); |
| 159 | auto [y0, y1] = load_float_vec2(input2 + d); |
| 160 | x0 = fast_silu(x0) * y0; |
| 161 | x1 = fast_silu(x1) * y1; |
| 162 | bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1); |
| 163 | out_vec.store(out + d); |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | template <typename scalar_t, typename input_t> |
| 168 | inline void clamp_sigmoid_and_mul_stub( |
| 169 | scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) { |
| 170 | static_assert( |
| 171 | std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>, |
| 172 | "clamp_sigmoid_and_mul_stub only supports input_t == float or input_t == scalar_t"); |
| 173 | using bVec = at::vec::Vectorized<scalar_t>; |
| 174 | using fVec = at::vec::Vectorized<float>; |
| 175 | const fVec one = fVec(1.f); |
| 176 | const fVec limit_v = fVec(limit); |
| 177 | const fVec nlimit_v = fVec(-limit); |
| 178 | const fVec alpha_v = fVec(alpha); |
| 179 | |
| 180 | #pragma GCC unroll 4 |
| 181 | for (int64_t d = 0; d < 2 * size; d += bVec::size()) { |
| 182 | auto [x0_, y0_] = load_float_vec2(input + d); |
| 183 | auto [x0, y0] = at::vec::deinterleave2<float>(x0_, y0_); |
| 184 | |
| 185 | x0 = at::vec::minimum(x0, limit_v); |
| 186 | y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0)); |
| 187 | x0 = fast_sigmoid_glu(x0, alpha_v) * (y0 + one); |
| 188 | store_from_float_ext(out + d / 2, x0); |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | template <typename scalar_t> |
| 193 | inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) { |
| 194 | using bVec = at::vec::Vectorized<scalar_t>; |
| 195 | using fVec = at::vec::Vectorized<float>; |
| 196 | constexpr int kVecSize = bVec::size(); |
| 197 | const fVec weight_vec = fVec(weight); |
| 198 | int64_t d; |
| 199 | #pragma GCC unroll 4 |
| 200 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 201 | auto [x0, x1] = load_float_vec2(input + d); |
| 202 | bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec); |
| 203 | out_vec.store(out + d); |
| 204 | } |
| 205 | for (; d < size; ++d) { |
| 206 | out[d] = static_cast<scalar_t>(input[d] * weight); |
| 207 | } |
| 208 | } |
| 209 | |
| 210 | // input = input + input2 |
| 211 | inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) { |
| 212 | using fVec = at::vec::Vectorized<float>; |
| 213 | constexpr int kVecSize = fVec::size(); |
| 214 | int64_t d; |
| 215 | #pragma GCC unroll 4 |
| 216 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 217 | fVec x_fvec = fVec::loadu(input + d); |
| 218 | fVec y_fvec = fVec::loadu(input2 + d); |
| 219 | x_fvec = x_fvec + y_fvec; |
| 220 | x_fvec.store(input + d); |
| 221 | } |
| 222 | for (; d < size; ++d) { |
| 223 | input[d] = input[d] + input2[d]; |
| 224 | } |
| 225 | } |
| 226 | |
| 227 | template <typename scalar_t> |
| 228 | inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) { |
| 229 | using bVec = at::vec::Vectorized<scalar_t>; |
| 230 | using fVec = at::vec::Vectorized<float>; |
| 231 | constexpr int kVecSize = bVec::size(); |
| 232 | const fVec weight_vec = fVec(weight); |
| 233 | int64_t d; |
| 234 | #pragma GCC unroll 4 |
| 235 | for (d = 0; d <= size - kVecSize; d += kVecSize) { |
| 236 | auto [x0, x1] = load_float_vec2(input + d); |
| 237 | bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec); |
| 238 | out_vec.store(out + d); |
| 239 | } |
| 240 | for (; d < size; ++d) { |
| 241 | out[d] = static_cast<scalar_t>(input[d] * weight); |
| 242 | } |
| 243 | } |