@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e69c990c21
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 01:15 UTC

[Feature][CPU Backend]: Optimize ARM vectorization backend (#30329)

Signed-off-by: Radu Salavat <radu.salavat@arm.com>

Radu Salavat committed Feb 3, 2026 at 04:17 UTC e69c990c216c623b1de22f055926602a336f9352
5 files changed +552 -597
csrc/cpu/cpu_attn_impl.hpp
-11
@@ -816,14 +816,10 @@ struct VecTypeTrait<float> {
816 using vec_t = vec_op::FP32Vec16;
817 };
818
819 -// ARM only supports BF16 with ARMv8.6-A extension
820 -#if (defined(__aarch64__) && !defined(ARM_BF16_SUPPORT))
821 -#else
819 template <>
820 struct VecTypeTrait<c10::BFloat16> {
821 using vec_t = vec_op::BF16Vec16;
822 };
826 -#endif
823
824 #if !defined(__powerpc__) && !defined(__s390x__)
825 template <>
@@ -1585,17 +1581,10 @@ class AttentionMainLoop {
1581
1582 if (use_sink) {
1583 alignas(64) float s_aux_fp32[16];
1588 -#if defined(__aarch64__) && !defined(ARM_BF16_SUPPORT)
1589 - // ARM without native BF16 support: manual conversion
1590 - for (int i = 0; i < 16; ++i) {
1591 - s_aux_fp32[i] = static_cast<float>(curr_s_aux[i]);
1592 - }
1593 -#else
1584 // All other platforms have BF16Vec16 available
1585 vec_op::BF16Vec16 vec_bf16(curr_s_aux);
1586 vec_op::FP32Vec16 vec_fp32(vec_bf16);
1587 vec_fp32.save(s_aux_fp32);
1598 -#endif
1588
1589 float* __restrict__ curr_sum_buffer = sum_buffer;
1590 float* __restrict__ curr_max_buffer = max_buffer;
csrc/cpu/cpu_types_arm.hpp
+551 -579
@@ -1,23 +1,24 @@
1 +#include <cmath>
2 +#include <type_traits>
3 +
4 #include <arm_neon.h>
5 +
6 #include <torch/all.h>
3 -#include <cmath>
7 +#include <ATen/cpu/vec/functional.h>
8 +#include <ATen/cpu/vec/vec.h>
9
10 #if defined(__APPLE__)
11 #include "omp.h"
12 #endif
13
14 +using namespace at::vec;
15 +
16 namespace vec_op {
17
11 -#ifdef ARM_BF16_SUPPORT
12 - #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
13 - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
14 - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
15 - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
16 -#else
17 - #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
18 - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
19 - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
20 -#endif
18 +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
19 + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
20 + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
21 + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
22
23 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
24 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -45,667 +46,632 @@ constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
46
47 template <typename T, T count, typename F,
48 typename = std::enable_if_t<std::is_invocable_v<F, T>>>
48 -constexpr void unroll_loop(F&& f) {
49 +inline constexpr void unroll_loop(F&& f) {
50 unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
51 }
52
52 -template <typename T>
53 -struct Vec {
54 - constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
55 -};
53 +template <typename T, typename... Ts>
54 +struct is_one_of : std::bool_constant<(std::is_same_v<T, Ts> || ...)> {};
55
57 -struct FP32Vec8;
58 -struct FP32Vec16;
56 +template <typename T, typename... Ts>
57 +inline constexpr bool is_one_of_v = is_one_of<T, Ts...>::value;
58
60 -struct FP16Vec8 : public Vec<FP16Vec8> {
61 - constexpr static int VEC_ELEM_NUM = 8;
59 +struct uninit_t {
60 + explicit constexpr uninit_t() = default;
61 +};
62 +inline constexpr uninit_t uninit{};
63
63 - float16x8_t reg;
64 +template <typename NxVectorizedTVecReg, typename T, int VEC_ELEM_NUM>
65 +union AliasReg {
66 + NxVectorizedTVecReg reg;
67 + T values[VEC_ELEM_NUM];
68 +};
69
65 - explicit FP16Vec8(const void* ptr)
66 - : reg(vld1q_f16(static_cast<const __fp16*>(ptr))) {};
70 +// Template over at::vec::Vectorized<T> to support
71 +// multiple vectorised registers into 1 of length VEC_REG_NUM val
72 +template <int N, typename T>
73 +struct NxVectorizedTVecReg {
74 + using value_t = T;
75 + using VectorizedT = Vectorized<T>;
76
68 - explicit FP16Vec8(const FP32Vec8&);
77 + VectorizedT val[N];
78
70 - void save(void* ptr) const { vst1q_f16(static_cast<__fp16*>(ptr), reg); }
71 -};
79 + NxVectorizedTVecReg() = default;
80 + NxVectorizedTVecReg(const NxVectorizedTVecReg&) = default;
81 + NxVectorizedTVecReg(NxVectorizedTVecReg&&) = default;
82 + NxVectorizedTVecReg& operator=(const NxVectorizedTVecReg&) = default;
83 + NxVectorizedTVecReg& operator=(NxVectorizedTVecReg&&) = default;
84
73 -struct FP16Vec16 : public Vec<FP16Vec16> {
74 - constexpr static int VEC_ELEM_NUM = 16;
85 + explicit NxVectorizedTVecReg(uninit_t) noexcept {};
86
76 - float16x8x2_t reg;
87 + FORCE_INLINE explicit NxVectorizedTVecReg(const VectorizedT& vec_t) {
88 + unroll_loop<int, N>([&](int i) { val[i] = vec_t; });
89 + };
90
78 - explicit FP16Vec16(const void* ptr) {
79 - reg.val[0] = vld1q_f16(reinterpret_cast<const __fp16*>(ptr));
80 - reg.val[1] = vld1q_f16(reinterpret_cast<const __fp16*>(ptr) + 8);
91 + FORCE_INLINE explicit NxVectorizedTVecReg(T v) noexcept {
92 + VectorizedT vv(v);
93 + unroll_loop<int, N>([&](int i) { val[i] = vv; });
94 }
95
83 - // ASIMD does not support non-temporal loads
84 - explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
96 + FORCE_INLINE explicit NxVectorizedTVecReg(const void* ptr) { load(ptr); }
97 + explicit NxVectorizedTVecReg(const void* ptr, const int elem_num) {
98 + load(ptr, elem_num);
99 + }
100
86 - explicit FP16Vec16(const FP32Vec16& vec);
87 - void save(void* ptr) const {
88 - vst1q_f16(reinterpret_cast<__fp16*>(ptr), reg.val[0]);
89 - vst1q_f16(reinterpret_cast<__fp16*>(ptr) + 8, reg.val[1]);
101 + static constexpr int size() noexcept { return N * VectorizedT::size(); }
102 +
103 + FORCE_INLINE void save(void* ptr) const {
104 + value_t* base = reinterpret_cast<value_t*>(ptr);
105 + unroll_loop<int, N>(
106 + [&](int i) { val[i].store(base + i * VectorizedT::size()); });
107 + }
108 + FORCE_INLINE void load(const void* ptr) {
109 + const value_t* base = reinterpret_cast<const value_t*>(ptr);
110 + unroll_loop<int, N>([&](int i) {
111 + val[i] = VectorizedT::loadu(base + i * VectorizedT::size());
112 + });
113 }
114
92 - void save(void* ptr, const int elem_num) const {
93 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
94 - int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
115 + FORCE_INLINE void save(void* ptr, const int elem_num) const {
116 + value_t* base = reinterpret_cast<value_t*>(ptr);
117 + save_partial(base, elem_num);
118 + }
119
96 - if (full_blocks > 0) {
97 - vst1q_f16(reinterpret_cast<__fp16*>(ptr), reg.val[0]);
98 - if (full_blocks > 1) {
99 - vst1q_f16(reinterpret_cast<__fp16*>(ptr) + 8, reg.val[1]);
100 - }
101 - }
120 + FORCE_INLINE void load(const void* ptr, const int elem_num) {
121 + const value_t* base = reinterpret_cast<const value_t*>(ptr);
122 + load_partial(base, elem_num);
123 + }
124
103 - // Note: below is the unrolled version of the following code:
104 - //
105 - // for (int i = 0; i < remainder; ++i) {
106 - // reinterpret_cast<__fp16*>(ptr)[full_blocks * 8 + i] =
107 - // vgetq_lane_f16(temp, i);
108 - // }
109 - //
110 - // For macOS build (Clang), the arm/neon intrinsics function
111 - // `vgetq_lane_f16` needs the parameter `i` to be constant at compile
112 - // time.
125 + FORCE_INLINE void save_partial(value_t* base, int elem_num) const {
126 + const int w = VectorizedT::size();
127 + int full = elem_num / w;
128 + int rem = elem_num % w;
129 + for (int i = 0; i < full; i++) val[i].store(base + i * w);
130 + if (rem) val[full].store(base + full * w, rem);
131 + }
132
114 - if (remainder > 0) {
115 - float16x8_t temp = reg.val[full_blocks];
116 - __fp16* fp16_ptr = reinterpret_cast<__fp16*>(ptr);
117 - switch (remainder) {
118 - case 1:
119 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
120 - break;
121 - case 2:
122 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
123 - fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
124 - break;
125 - case 3:
126 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
127 - fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
128 - fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
129 - break;
130 - case 4:
131 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
132 - fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
133 - fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
134 - fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
135 - break;
136 - case 5:
137 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
138 - fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
139 - fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
140 - fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
141 - fp16_ptr[full_blocks * 8 + 4] = vgetq_lane_f16(temp, 4);
142 - break;
143 - case 6:
144 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
145 - fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
146 - fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
147 - fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
148 - fp16_ptr[full_blocks * 8 + 4] = vgetq_lane_f16(temp, 4);
149 - fp16_ptr[full_blocks * 8 + 5] = vgetq_lane_f16(temp, 5);
150 - break;
151 - case 7:
152 - fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
153 - fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
154 - fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
155 - fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
156 - fp16_ptr[full_blocks * 8 + 4] = vgetq_lane_f16(temp, 4);
157 - fp16_ptr[full_blocks * 8 + 5] = vgetq_lane_f16(temp, 5);
158 - fp16_ptr[full_blocks * 8 + 6] = vgetq_lane_f16(temp, 6);
159 - break;
133 + FORCE_INLINE void load_partial(const value_t* base, int elem_num) {
134 + const int w = VectorizedT::size();
135 + int full = elem_num / w;
136 + int rem = elem_num % w;
137 + for (int i = 0; i < full; i++) val[i] = VectorizedT::loadu(base + i * w);
138 + if (rem) val[full] = VectorizedT::loadu(base + full * w, rem);
139 + }
140
161 - default:
162 - break;
141 + template <VectorizedT (VectorizedT::*torch_vec_func)() const,
142 + value_t (*std_func)(value_t)>
143 + FORCE_INLINE NxVectorizedTVecReg opt_vec_func_impl() const {
144 + NxVectorizedTVecReg result;
145 +
146 + if constexpr (torch_vec_func != nullptr) {
147 + unroll_loop<int, N>(
148 + [&](int i) { result.val[i] = (val[i].*torch_vec_func)(); });
149 + } else {
150 + for (int i = 0; i < N; i++) {
151 + alignas(64) value_t buf[VectorizedT::size()];
152 + val[i].store(buf);
153 + for (int j = 0; j < VectorizedT::size(); ++j) {
154 + buf[j] = std_func(buf[j]);
155 + }
156 + result.val[i] = VectorizedT::loadu(buf);
157 }
158 }
159 + return result;
160 }
161 };
162
168 -#ifdef ARM_BF16_SUPPORT
169 -struct BF16Vec8 : public Vec<BF16Vec8> {
170 - constexpr static int VEC_ELEM_NUM = 8;
163 +template <typename DerivedClassT, int N, typename T>
164 +struct VectorizedRegWrapper {
165 + using ScalarT = T;
166 + using VectorizedT = Vectorized<T>;
167 + using NxVectorizedTArray = NxVectorizedTVecReg<N, T>;
168 +
169 + constexpr static int VEC_REG_NUM = N;
170 + constexpr static int VEC_ELEM_NUM = VEC_REG_NUM * VectorizedT::size();
171 + constexpr static int get_elem_num() { return VEC_ELEM_NUM; };
172 +
173 + NxVectorizedTArray reg;
174 +
175 + VectorizedRegWrapper() noexcept = default;
176 + explicit VectorizedRegWrapper(uninit_t) noexcept : reg{uninit} {};
177 + explicit VectorizedRegWrapper(T v) : reg(v) {};
178 + explicit VectorizedRegWrapper(const void* ptr) : reg(ptr) {};
179 + explicit VectorizedRegWrapper(const void* ptr, const int elem_num)
180 + : reg(ptr, elem_num) {};
181 + explicit VectorizedRegWrapper(const VectorizedT& r) : reg(r) {};
182 + explicit VectorizedRegWrapper(const NxVectorizedTArray& r) : reg(r) {};
183 +
184 + VectorizedRegWrapper(const VectorizedRegWrapper&) = default;
185 + VectorizedRegWrapper(VectorizedRegWrapper&&) = default;
186 + VectorizedRegWrapper& operator=(VectorizedRegWrapper&&) = default;
187 + VectorizedRegWrapper& operator=(const VectorizedRegWrapper&) = default;
188 +
189 + FORCE_INLINE void save(void* ptr) const { reg.save(ptr); }
190 + void save(void* ptr, const int elem_num) const { reg.save(ptr, elem_num); }
191 +
192 +// Define optimized functions using at::vec::Vectorized<T> where possible
193 +// Fallback to std:: functions when not available
194 +#define OPT_TORCH_IMPL(FUNC_NAME, STD_FUNC_NAME, TORCH_FUNC_NAME, ...) \
195 + FORCE_INLINE DerivedClassT FUNC_NAME() const { \
196 + if constexpr (is_one_of_v<T, __VA_ARGS__>) { \
197 + return DerivedClassT{ \
198 + reg.template opt_vec_func_impl<&VectorizedT::TORCH_FUNC_NAME, \
199 + std::STD_FUNC_NAME>()}; \
200 + } else { \
201 + return DerivedClassT{reg.template opt_vec_func_impl< \
202 + nullptr, static_cast<ScalarT (*)(ScalarT)>(&std::STD_FUNC_NAME)>()}; \
203 + } \
204 + }
205
172 - bfloat16x8_t reg;
206 + // Define optimized functions for datatypes passed in __VA_ARGS__
207 + OPT_TORCH_IMPL(abs, abs, abs, c10::Half, float)
208 + OPT_TORCH_IMPL(er, erf, erf, float)
209 + OPT_TORCH_IMPL(exp, exp, fexp_u20, float)
210 + OPT_TORCH_IMPL(exp_u20, exp, exp_u20, float)
211 + OPT_TORCH_IMPL(sin, sin, sin, float)
212 + OPT_TORCH_IMPL(sinh, sinh, sinh, float)
213 + OPT_TORCH_IMPL(cos, cos, cos, float)
214 + OPT_TORCH_IMPL(cosh, cosh, cosh, float)
215 + OPT_TORCH_IMPL(log, log, log, float)
216 + OPT_TORCH_IMPL(log10, log10, log10, float)
217 + OPT_TORCH_IMPL(sqrt, sqrt, sqrt, c10::Half, float)
218 + OPT_TORCH_IMPL(tan, tan, tan, float)
219 + OPT_TORCH_IMPL(tanh, tanh, tanh, float)
220 +
221 +#undef OPT_TORCH_IMPL
222 +};
223
174 - explicit BF16Vec8(const void* ptr)
175 - : reg(*reinterpret_cast<const bfloat16x8_t*>(ptr)) {};
224 +// forward declare vectorised dtypes
225 +struct FP32Vec8;
226 +struct FP32Vec16;
227 +struct FP16Vec8;
228 +struct FP16Vec16;
229 +struct BF16Vec8;
230 +struct BF16Vec16;
231
177 - explicit BF16Vec8(bfloat16x8_t data) : reg(data) {};
232 +struct INT8Vec16;
233 +struct INT32Vec16;
234
179 - explicit BF16Vec8(const FP32Vec8&);
235 +template <typename T>
236 +struct VecType {
237 + using vec_type = void;
238 +};
239
181 - explicit BF16Vec8(float32x4x2_t v)
182 - : reg(vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.val[0]), v.val[1])) {};
240 +template <typename T>
241 +using vec_t = typename VecType<T>::vec_type;
242
184 - void save(void* ptr) const { *reinterpret_cast<bfloat16x8_t*>(ptr) = reg; }
243 +template <>
244 +struct VecType<float> {
245 + using vec_type = FP32Vec8;
246 };
247
187 -struct BF16Vec16 : public Vec<BF16Vec16> {
188 - constexpr static int VEC_ELEM_NUM = 16;
189 -
190 - bfloat16x8x2_t reg;
248 +template <>
249 +struct VecType<c10::Half> {
250 + using vec_type = FP16Vec8;
251 +};
252
192 - explicit BF16Vec16(const void* ptr)
193 - : reg(*reinterpret_cast<const bfloat16x8x2_t*>(ptr)) {};
253 +template <>
254 +struct VecType<c10::BFloat16> {
255 + using vec_type = BF16Vec8;
256 +};
257
195 - // ASIMD does not support non-temporal loads
196 - explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
258 +struct FP16Vec8 : public VectorizedRegWrapper<FP16Vec8, 1, c10::Half> {
259 + using Base = VectorizedRegWrapper<FP16Vec8, 1, c10::Half>;
260 + using Base::Base;
261 + using Base::get_elem_num;
262 + using Base::VEC_ELEM_NUM;
263
198 - explicit BF16Vec16(bfloat16x8x2_t data) : reg(data) {};
264 + explicit FP16Vec8(const FP32Vec8&);
265 +};
266
200 - explicit BF16Vec16(const FP32Vec16&);
267 +struct FP16Vec16 : public VectorizedRegWrapper<FP16Vec16, 2, c10::Half> {
268 + using Base = VectorizedRegWrapper<FP16Vec16, 2, c10::Half>;
269 + using Base::Base;
270 + using Base::get_elem_num;
271 + using Base::VEC_ELEM_NUM;
272
202 - explicit BF16Vec16(float32x4x4_t v)
203 - : reg({vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.val[0]), v.val[1]),
204 - vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.val[2]), v.val[3])}) {};
273 + // ASIMD does not support non-temporal loads
274 + explicit FP16Vec16(bool, const void* ptr) : Base(ptr) {}
275
206 - void save(void* ptr) const { *reinterpret_cast<bfloat16x8x2_t*>(ptr) = reg; };
207 - void save(void* ptr, const int elem_num) const {
208 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
209 - int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
210 - for (int i = 0; i < full_blocks; i++)
211 - vst1q_bf16(
212 - reinterpret_cast<__bf16*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
213 - reg.val[i]);
214 - if (remainder > 0) {
215 - bfloat16x8_t temp = reg.val[full_blocks];
216 - bfloat16_t* base = reinterpret_cast<bfloat16_t*>(ptr) + full_blocks * 8;
217 - if (remainder > 0) base[0] = vgetq_lane_bf16(temp, 0);
218 - if (remainder > 1) base[1] = vgetq_lane_bf16(temp, 1);
219 - if (remainder > 2) base[2] = vgetq_lane_bf16(temp, 2);
220 - if (remainder > 3) base[3] = vgetq_lane_bf16(temp, 3);
221 - if (remainder > 4) base[4] = vgetq_lane_bf16(temp, 4);
222 - if (remainder > 5) base[5] = vgetq_lane_bf16(temp, 5);
223 - if (remainder > 6) base[6] = vgetq_lane_bf16(temp, 6);
224 - }
225 - };
276 + explicit FP16Vec16(const FP32Vec16& vec);
277 };
278
228 -struct BF16Vec32 : public Vec<BF16Vec32> {
229 - constexpr static int VEC_ELEM_NUM = 32;
279 +struct BF16Vec8 : public VectorizedRegWrapper<BF16Vec8, 1, c10::BFloat16> {
280 + using Base = VectorizedRegWrapper<BF16Vec8, 1, c10::BFloat16>;
281 + using VectorizedT = typename Base::VectorizedT;
282 + using Base::Base;
283 + using Base::get_elem_num;
284 + using Base::VEC_ELEM_NUM;
285
231 - bfloat16x8x4_t reg;
286 + explicit BF16Vec8(at_bfloat16x8_t data) : Base(VectorizedT(data)) {};
287
233 - explicit BF16Vec32(const void* ptr)
234 - : reg(*reinterpret_cast<const bfloat16x8x4_t*>(ptr)) {};
288 + explicit BF16Vec8(float32x4x2_t v) {
289 + reg.val[0] = convert_float_bfloat16(v.val[0], v.val[1]);
290 + };
291
236 - explicit BF16Vec32(bfloat16x8x4_t data) : reg(data) {};
292 + explicit BF16Vec8(const FP32Vec8&);
293 +};
294
238 - explicit BF16Vec32(const BF16Vec8& vec8_data)
239 - : reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {};
295 +struct BF16Vec16 : public VectorizedRegWrapper<BF16Vec16, 2, c10::BFloat16> {
296 + using Base = VectorizedRegWrapper<BF16Vec16, 2, c10::BFloat16>;
297 + using VectorizedT = typename Base::VectorizedT;
298 + using Base::Base;
299 + using Base::get_elem_num;
300 + using Base::VEC_ELEM_NUM;
301
241 - void save(void* ptr) const { *reinterpret_cast<bfloat16x8x4_t*>(ptr) = reg; };
242 - void save(void* ptr, const int elem_num) const {
243 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
244 - int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
245 - for (int i = 0; i < full_blocks; i++)
246 - vst1q_bf16(
247 - reinterpret_cast<__bf16*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
248 - reg.val[i]);
249 - if (remainder > 0) {
250 - bfloat16x8_t temp = reg.val[full_blocks];
251 - bfloat16_t* base = reinterpret_cast<bfloat16_t*>(ptr) + full_blocks * 8;
252 - base[0] = vgetq_lane_bf16(temp, 0);
253 - if (remainder > 1) base[1] = vgetq_lane_bf16(temp, 1);
254 - if (remainder > 2) base[2] = vgetq_lane_bf16(temp, 2);
255 - if (remainder > 3) base[3] = vgetq_lane_bf16(temp, 3);
256 - if (remainder > 4) base[4] = vgetq_lane_bf16(temp, 4);
257 - if (remainder > 5) base[5] = vgetq_lane_bf16(temp, 5);
258 - if (remainder > 6) base[6] = vgetq_lane_bf16(temp, 6);
259 - }
302 + // ASIMD does not support non-temporal loads
303 + explicit BF16Vec16(bool, const void* ptr) : Base(ptr) {}
304 +
305 + explicit BF16Vec16(float32x4x4_t v) {
306 + reg.val[0] = convert_float_bfloat16(v.val[0], v.val[1]);
307 + reg.val[1] = convert_float_bfloat16(v.val[2], v.val[3]);
308 };
261 -};
262 -#endif
309
264 -struct FP32Vec4 : public Vec<FP32Vec4> {
265 - constexpr static int VEC_ELEM_NUM = 4;
310 + explicit BF16Vec16(const FP32Vec16&);
311 +};
312
267 - union AliasReg {
268 - float32x4_t reg;
269 - float values[VEC_ELEM_NUM];
313 +struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
314 + using Base = VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16>;
315 + using Base::Base;
316 + using Base::get_elem_num;
317 + using Base::VEC_ELEM_NUM;
318 +
319 + explicit BF16Vec32(const BF16Vec8& vec8_data) {
320 + reg.val[0] = vec8_data.reg.val[0];
321 + reg.val[1] = vec8_data.reg.val[0];
322 + reg.val[2] = vec8_data.reg.val[0];
323 + reg.val[3] = vec8_data.reg.val[0];
324 };
325 +};
326
272 - float32x4_t reg;
273 -
274 - explicit FP32Vec4(float v) : reg(vdupq_n_f32(v)) {};
327 +struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
328 + using Base = VectorizedRegWrapper<FP32Vec4, 1, float>;
329 + using Base::Base;
330 + using Base::get_elem_num;
331 + using Base::VEC_ELEM_NUM;
332
276 - explicit FP32Vec4() : reg(vdupq_n_f32(0.0f)) {};
333 + using VectorizedT = typename Base::VectorizedT;
334 + using Vectorized1x4f = typename Base::NxVectorizedTArray;
335
278 - explicit FP32Vec4(const float* ptr) : reg(vld1q_f32(ptr)) {};
336 + FP32Vec4() : Base() {};
337 + explicit FP32Vec4(float v) : Base(v) {};
338
280 - explicit FP32Vec4(float32x4_t data) : reg(data) {};
339 + explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
340
282 - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
341 + explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
342 };
343
285 -struct FP32Vec8 : public Vec<FP32Vec8> {
286 - constexpr static int VEC_ELEM_NUM = 8;
287 - union AliasReg {
288 - float32x4x2_t reg;
289 - float values[VEC_ELEM_NUM];
290 - };
291 -
292 - float32x4x2_t reg;
344 +struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
345 + using Base = VectorizedRegWrapper<FP32Vec8, 2, float>;
346 + using Base::Base;
347 + using Base::get_elem_num;
348 + using Base::VEC_ELEM_NUM;
349 + using Base::VEC_REG_NUM;
350
294 - explicit FP32Vec8(float v) : reg({vmovq_n_f32(v), vmovq_n_f32(v)}) {};
351 + using VectorizedT = typename Base::VectorizedT;
352 + using Vectorized2x4f = typename Base::NxVectorizedTArray;
353
296 - explicit FP32Vec8() : reg({vmovq_n_f32(0.0), vmovq_n_f32(0.0)}) {};
354 + FP32Vec8() : Base() {};
355 + FP32Vec8(const FP32Vec8& data) : Base(data) {};
356
357 + explicit FP32Vec8(float v) : Base(v) {};
358 explicit FP32Vec8(const float* ptr)
299 - : reg({vld1q_f32(ptr), vld1q_f32(ptr + 4)}) {};
359 + : Base(reinterpret_cast<const void*>(ptr)) {};
360 + explicit FP32Vec8(const float* ptr, const int elem_num)
361 + : Base(reinterpret_cast<const void*>(ptr), elem_num) {};
362
301 - explicit FP32Vec8(float32x4x2_t data) : reg(data) {};
302 -
303 - explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
363 + explicit FP32Vec8(const Vectorized2x4f& data) {
364 + reg.val[0] = data.val[0];
365 + reg.val[1] = data.val[1];
366 + };
367
368 + explicit FP32Vec8(const BF16Vec8& v) {
369 + std::tie(reg.val[0], reg.val[1]) = convert_bfloat16_float(v.reg.val[0]);
370 + };
371 explicit FP32Vec8(const FP16Vec8& v) {
306 - reg.val[0] = vcvt_f32_f16(vget_low_f16(v.reg));
307 - reg.val[1] = vcvt_f32_f16(vget_high_f16(v.reg));
372 + reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[0])));
373 + reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[0])));
374 };
375 + explicit FP32Vec8(float16x8_t v) {
376 + reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v)));
377 + reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v)));
378 + };
379 + explicit FP32Vec8(at_bfloat16x8_t v) {
380 + std::tie(reg.val[0], reg.val[1]) =
381 + convert_bfloat16_float(Vectorized<c10::BFloat16>(v));
382 + };
383 + explicit FP32Vec8(float32x4x2_t data) {
384 + reg.val[0] = Vectorized<float>(data.val[0]);
385 + reg.val[1] = Vectorized<float>(data.val[1]);
386 + }
387
310 - explicit FP32Vec8(float16x8_t v)
311 - : reg({vcvt_f32_f16(vget_low_f16(v)), vcvt_f32_f16(vget_high_f16(v))}) {};
312 -
313 -#ifdef ARM_BF16_SUPPORT
314 -
315 - explicit FP32Vec8(bfloat16x8_t v)
316 - : reg({vcvtq_low_f32_bf16(v), vcvtq_high_f32_bf16(v)}) {};
317 -
318 - explicit FP32Vec8(const BF16Vec8& v)
319 - : reg({vcvtq_low_f32_bf16(v.reg), vcvtq_high_f32_bf16(v.reg)}) {};
320 -
321 -#endif
322 -
323 - float reduce_sum() const {
324 - AliasReg ar;
325 - ar.reg = reg;
388 + FORCE_INLINE float reduce_sum() const noexcept {
389 float answer = 0;
327 - unroll_loop<int, VEC_ELEM_NUM>(
328 - [&answer, &ar](int i) { answer += ar.values[i]; });
390 + std::plus<VectorizedT> add;
391
392 + unroll_loop<int, VEC_REG_NUM>([&](int i) {
393 + answer += at::vec::vec_reduce_all<float, std::plus<VectorizedT>>(
394 + add, reg.val[i]);
395 + });
396 return answer;
397 }
398
333 - FP32Vec8 exp() const {
334 - AliasReg ar;
335 - ar.reg = reg;
336 -
337 - float32x2_t exp_vec0 = {expf(ar.values[0]), expf(ar.values[1])};
338 - float32x2_t exp_vec1 = {expf(ar.values[2]), expf(ar.values[3])};
339 - float32x2_t exp_vec2 = {expf(ar.values[4]), expf(ar.values[5])};
340 - float32x2_t exp_vec3 = {expf(ar.values[6]), expf(ar.values[7])};
341 -
342 - float32x4_t result0 = vcombine_f32(exp_vec0, exp_vec1);
343 - float32x4_t result1 = vcombine_f32(exp_vec2, exp_vec3);
344 -
345 - float32x4x2_t result;
346 - result.val[0] = result0;
347 - result.val[1] = result1;
348 -
349 - return FP32Vec8(result);
399 + FORCE_INLINE FP32Vec8 operator+(const FP32Vec8& b) const noexcept {
400 + FP32Vec8 r(uninit);
401 + r.reg.val[0] = reg.val[0] + b.reg.val[0];
402 + r.reg.val[1] = reg.val[1] + b.reg.val[1];
403 + return r;
404 }
405
352 - FP32Vec8 tanh() const {
353 - AliasReg ar;
354 - ar.reg = reg;
355 -
356 - float32x2_t tanh_vec0 = {tanhf(ar.values[0]), tanhf(ar.values[1])};
357 - float32x2_t tanh_vec1 = {tanhf(ar.values[2]), tanhf(ar.values[3])};
358 - float32x2_t tanh_vec2 = {tanhf(ar.values[4]), tanhf(ar.values[5])};
359 - float32x2_t tanh_vec3 = {tanhf(ar.values[6]), tanhf(ar.values[7])};
360 -
361 - float32x4_t result0 = vcombine_f32(tanh_vec0, tanh_vec1);
362 - float32x4_t result1 = vcombine_f32(tanh_vec2, tanh_vec3);
363 -
364 - float32x4x2_t result;
365 - result.val[0] = result0;
366 - result.val[1] = result1;
367 -
368 - return FP32Vec8(result);
406 + FORCE_INLINE FP32Vec8 operator-(const FP32Vec8& b) const noexcept {
407 + FP32Vec8 r(uninit);
408 + r.reg.val[0] = reg.val[0] - b.reg.val[0];
409 + r.reg.val[1] = reg.val[1] - b.reg.val[1];
410 + return r;
411 }
412
371 - FP32Vec8 er() const {
372 - AliasReg ar;
373 - ar.reg = reg;
374 -
375 - float32x2_t er_vec0 = {static_cast<float32_t>(erf(ar.values[0])),
376 - static_cast<float32_t>(erf(ar.values[1]))};
377 - float32x2_t er_vec1 = {static_cast<float32_t>(erf(ar.values[2])),
378 - static_cast<float32_t>(erf(ar.values[3]))};
379 - float32x2_t er_vec2 = {static_cast<float32_t>(erf(ar.values[4])),
380 - static_cast<float32_t>(erf(ar.values[5]))};
381 - float32x2_t er_vec3 = {static_cast<float32_t>(erf(ar.values[6])),
382 - static_cast<float32_t>(erf(ar.values[7]))};
383 -
384 - float32x4_t result0 = vcombine_f32(er_vec0, er_vec1);
385 - float32x4_t result1 = vcombine_f32(er_vec2, er_vec3);
386 -
387 - float32x4x2_t result;
388 - result.val[0] = result0;
389 - result.val[1] = result1;
390 -
391 - return FP32Vec8(result);
392 - }
393 -
394 - FP32Vec8 operator*(const FP32Vec8& b) const {
395 - return FP32Vec8(float32x4x2_t({vmulq_f32(reg.val[0], b.reg.val[0]),
396 - vmulq_f32(reg.val[1], b.reg.val[1])}));
397 - }
398 -
399 - FP32Vec8 operator+(const FP32Vec8& b) const {
400 - return FP32Vec8(float32x4x2_t({vaddq_f32(reg.val[0], b.reg.val[0]),
401 - vaddq_f32(reg.val[1], b.reg.val[1])}));
402 - }
403 -
404 - FP32Vec8 operator-(const FP32Vec8& b) const {
405 - return FP32Vec8(float32x4x2_t({vsubq_f32(reg.val[0], b.reg.val[0]),
406 - vsubq_f32(reg.val[1], b.reg.val[1])}));
407 - }
408 -
409 - FP32Vec8 operator/(const FP32Vec8& b) const {
410 - return FP32Vec8(float32x4x2_t({vdivq_f32(reg.val[0], b.reg.val[0]),
411 - vdivq_f32(reg.val[1], b.reg.val[1])}));
413 + FORCE_INLINE FP32Vec8 operator*(const FP32Vec8& b) const noexcept {
414 + FP32Vec8 r(uninit);
415 + r.reg.val[0] = reg.val[0] * b.reg.val[0];
416 + r.reg.val[1] = reg.val[1] * b.reg.val[1];
417 + return r;
418 }
419
414 - void save(float* ptr) const {
415 - vst1q_f32(ptr, reg.val[0]);
416 - vst1q_f32(ptr + 4, reg.val[1]);
420 + FORCE_INLINE FP32Vec8 operator/(const FP32Vec8& b) const noexcept {
421 + FP32Vec8 r(uninit);
422 + r.reg.val[0] = reg.val[0] / b.reg.val[0];
423 + r.reg.val[1] = reg.val[1] / b.reg.val[1];
424 + return r;
425 }
426 };
427
420 -struct INT32Vec16 : public Vec<INT32Vec16> {
421 - constexpr static int VEC_ELEM_NUM = 16;
422 - union AliasReg {
423 - int32x4x4_t reg;
424 - int32_t values[VEC_ELEM_NUM];
425 - };
426 - int32x4x4_t reg;
428 +struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
429 + using Base = VectorizedRegWrapper<FP32Vec16, 4, float>;
430 + using Base::Base;
431 + using Base::get_elem_num;
432 + using Base::VEC_ELEM_NUM;
433
428 - explicit INT32Vec16(const void* ptr) {
429 - reg.val[0] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr));
430 - reg.val[1] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 4);
431 - reg.val[2] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 8);
432 - reg.val[3] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 12);
433 - }
434 + using ScalarT = typename Base::ScalarT;
435 + using VectorizedT = typename Base::VectorizedT;
436 + using Vectorized4x4f = typename Base::NxVectorizedTArray;
437
435 - void save(int32_t* ptr) const {
436 - vst1q_s32(ptr, reg.val[0]);
437 - vst1q_s32(ptr + 4, reg.val[1]);
438 - vst1q_s32(ptr + 8, reg.val[2]);
439 - vst1q_s32(ptr + 12, reg.val[3]);
438 + FP32Vec16() : Base() {};
439 + FP32Vec16(const FP32Vec16& data) : Base(data) {};
440 + explicit FP32Vec16(float v) : Base(v) {};
441 + explicit FP32Vec16(const float* ptr)
442 + : Base(reinterpret_cast<const void*>(ptr)) {};
443 + explicit FP32Vec16(const float* ptr, const int elem_num)
444 + : Base(reinterpret_cast<const void*>(ptr), elem_num) {};
445 + explicit FP32Vec16(const Vectorized4x4f& data) {
446 + reg.val[0] = data.val[0];
447 + reg.val[1] = data.val[1];
448 + reg.val[2] = data.val[2];
449 + reg.val[3] = data.val[3];
450 };
451
442 - void save(int32_t* ptr, const int elem_num) const {
443 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
444 - int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
445 -
446 - for (int i = 0; i < full_blocks; i++)
447 - vst1q_s32(
448 - reinterpret_cast<__int32_t*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
449 - reg.val[i]);
450 -
451 - if (remainder > 0) {
452 - int32x4_t temp = reg.val[full_blocks];
453 - int32_t* base = reinterpret_cast<int32_t*>(ptr) + full_blocks * 4;
454 - if (remainder > 0) base[0] = vgetq_lane_s32(temp, 0);
455 - if (remainder > 1) base[1] = vgetq_lane_s32(temp, 1);
456 - if (remainder > 2) base[2] = vgetq_lane_s32(temp, 2);
457 - if (remainder > 3) base[3] = vgetq_lane_s32(temp, 3);
458 - }
459 - }
460 -};
452 + // ASIMD does not support non-temporal loads
453 + explicit FP32Vec16(bool, const float* ptr) : Base(ptr) {}
454
462 -struct FP32Vec16 : public Vec<FP32Vec16> {
463 - constexpr static int VEC_ELEM_NUM = 16;
464 - union AliasReg {
465 - float32x4x4_t reg;
466 - float values[VEC_ELEM_NUM];
455 + explicit FP32Vec16(float32x4x4_t data) {
456 + reg.val[0] = data.val[0];
457 + reg.val[1] = data.val[1];
458 + reg.val[2] = data.val[2];
459 + reg.val[3] = data.val[3];
460 };
461
469 - float32x4x4_t reg;
470 -
471 - explicit FP32Vec16(float v)
472 - : reg({vmovq_n_f32(v), vmovq_n_f32(v), vmovq_n_f32(v), vmovq_n_f32(v)}) {}
473 -
474 - explicit FP32Vec16()
475 - : reg({vmovq_n_f32(0.0), vmovq_n_f32(0.0), vmovq_n_f32(0.0),
476 - vmovq_n_f32(0.0)}) {}
477 -
478 - explicit FP32Vec16(const float* ptr)
479 - : reg({vld1q_f32(ptr), vld1q_f32(ptr + 4), vld1q_f32(ptr + 8),
480 - vld1q_f32(ptr + 12)}) {}
481 -
482 - // ASIMD does not support non-temporal loads
483 - explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
484 -
485 - explicit FP32Vec16(float32x4x4_t data) : reg(data) {}
462 + explicit FP32Vec16(const FP32Vec4& data) {
463 + reg.val[0] = data.reg.val[0];
464 + reg.val[1] = data.reg.val[0];
465 + reg.val[2] = data.reg.val[0];
466 + reg.val[3] = data.reg.val[0];
467 + };
468
469 explicit FP32Vec16(const FP32Vec8& data) {
470 reg.val[0] = data.reg.val[0];
471 reg.val[1] = data.reg.val[1];
472 reg.val[2] = data.reg.val[0];
473 reg.val[3] = data.reg.val[1];
492 - }
493 -
494 - explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}
495 -
496 - explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v.reg)) {}
497 -
498 -#ifdef ARM_BF16_SUPPORT
499 - explicit FP32Vec16(bfloat16x8x2_t v)
500 - : reg({vcvtq_low_f32_bf16(v.val[0]), vcvtq_high_f32_bf16(v.val[0]),
501 - vcvtq_low_f32_bf16(v.val[1]), vcvtq_high_f32_bf16(v.val[1])}) {};
502 -#endif
503 -
504 - explicit FP32Vec16(const FP32Vec4& data) {
505 - reg.val[0] = data.reg;
506 - reg.val[1] = data.reg;
507 - reg.val[2] = data.reg;
508 - reg.val[3] = data.reg;
474 };
475
511 -#ifdef ARM_BF16_SUPPORT
512 - explicit FP32Vec16(const BF16Vec16& v)
513 - : reg({vcvtq_low_f32_bf16(v.reg.val[0]),
514 - vcvtq_high_f32_bf16(v.reg.val[0]),
515 - vcvtq_low_f32_bf16(v.reg.val[1]),
516 - vcvtq_high_f32_bf16(v.reg.val[1])}) {};
476 + explicit FP32Vec16(const BF16Vec16& v) {
477 + std::tie(reg.val[0], reg.val[1]) = convert_bfloat16_float(v.reg.val[0]);
478 + std::tie(reg.val[2], reg.val[3]) = convert_bfloat16_float(v.reg.val[1]);
479 + };
480
481 explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
519 -#endif
482
483 explicit FP32Vec16(const FP16Vec16& v) {
522 - reg.val[0] = vcvt_f32_f16(vget_low_f16(v.reg.val[0]));
523 - reg.val[1] = vcvt_f32_f16(vget_high_f16(v.reg.val[0]));
524 - reg.val[2] = vcvt_f32_f16(vget_low_f16(v.reg.val[1]));
525 - reg.val[3] = vcvt_f32_f16(vget_high_f16(v.reg.val[1]));
526 - };
527 - explicit FP32Vec16(const INT32Vec16& v) {
528 - reg.val[0] = vcvtq_f32_s32(v.reg.val[0]);
529 - reg.val[1] = vcvtq_f32_s32(v.reg.val[1]);
530 - reg.val[2] = vcvtq_f32_s32(v.reg.val[2]);
531 - reg.val[3] = vcvtq_f32_s32(v.reg.val[3]);
532 - };
533 - FP32Vec16 operator+(const FP32Vec16& b) const {
534 - return FP32Vec16(float32x4x4_t({vaddq_f32(reg.val[0], b.reg.val[0]),
535 - vaddq_f32(reg.val[1], b.reg.val[1]),
536 - vaddq_f32(reg.val[2], b.reg.val[2]),
537 - vaddq_f32(reg.val[3], b.reg.val[3])}));
484 + reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[0])));
485 + reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[0])));
486 + reg.val[2] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[1])));
487 + reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
488 };
489
540 - FP32Vec16 operator*(const FP32Vec16& b) const {
541 - return FP32Vec16(float32x4x4_t({vmulq_f32(reg.val[0], b.reg.val[0]),
542 - vmulq_f32(reg.val[1], b.reg.val[1]),
543 - vmulq_f32(reg.val[2], b.reg.val[2]),
544 - vmulq_f32(reg.val[3], b.reg.val[3])}));
545 - };
490 + FORCE_INLINE FP32Vec16 operator+(const FP32Vec16& b) const noexcept {
491 + FP32Vec16 r(uninit);
492 + r.reg.val[0] = reg.val[0] + b.reg.val[0];
493 + r.reg.val[1] = reg.val[1] + b.reg.val[1];
494 + r.reg.val[2] = reg.val[2] + b.reg.val[2];
495 + r.reg.val[3] = reg.val[3] + b.reg.val[3];
496 + return r;
497 + }
498
547 - FP32Vec16 operator-(const FP32Vec16& b) const {
548 - return FP32Vec16(float32x4x4_t({vsubq_f32(reg.val[0], b.reg.val[0]),
549 - vsubq_f32(reg.val[1], b.reg.val[1]),
550 - vsubq_f32(reg.val[2], b.reg.val[2]),
551 - vsubq_f32(reg.val[3], b.reg.val[3])}));
552 - };
499 + FORCE_INLINE FP32Vec16 operator-(const FP32Vec16& b) const noexcept {
500 + FP32Vec16 r(uninit);
501 + r.reg.val[0] = reg.val[0] - b.reg.val[0];
502 + r.reg.val[1] = reg.val[1] - b.reg.val[1];
503 + r.reg.val[2] = reg.val[2] - b.reg.val[2];
504 + r.reg.val[3] = reg.val[3] - b.reg.val[3];
505 + return r;
506 + }
507
554 - FP32Vec16 operator/(const FP32Vec16& b) const {
555 - return FP32Vec16(float32x4x4_t({vdivq_f32(reg.val[0], b.reg.val[0]),
556 - vdivq_f32(reg.val[1], b.reg.val[1]),
557 - vdivq_f32(reg.val[2], b.reg.val[2]),
558 - vdivq_f32(reg.val[3], b.reg.val[3])}));
508 + FORCE_INLINE FP32Vec16 operator*(const FP32Vec16& b) const noexcept {
509 + FP32Vec16 r(uninit);
510 + r.reg.val[0] = reg.val[0] * b.reg.val[0];
511 + r.reg.val[1] = reg.val[1] * b.reg.val[1];
512 + r.reg.val[2] = reg.val[2] * b.reg.val[2];
513 + r.reg.val[3] = reg.val[3] * b.reg.val[3];
514 + return r;
515 + }
516 +
517 + FORCE_INLINE FP32Vec16 operator/(const FP32Vec16& b) const noexcept {
518 + FP32Vec16 r(uninit);
519 + r.reg.val[0] = reg.val[0] / b.reg.val[0];
520 + r.reg.val[1] = reg.val[1] / b.reg.val[1];
521 + r.reg.val[2] = reg.val[2] / b.reg.val[2];
522 + r.reg.val[3] = reg.val[3] / b.reg.val[3];
523 + return r;
524 + }
525 +
526 + FORCE_INLINE FP32Vec16 clamp(const FP32Vec16& min,
527 + const FP32Vec16& max) const {
528 + FP32Vec16 r(uninit);
529 + r.reg.val[0] = at::vec::clamp(reg.val[0], min.reg.val[0], max.reg.val[0]);
530 + r.reg.val[1] = at::vec::clamp(reg.val[1], min.reg.val[1], max.reg.val[1]);
531 + r.reg.val[2] = at::vec::clamp(reg.val[2], min.reg.val[2], max.reg.val[2]);
532 + r.reg.val[3] = at::vec::clamp(reg.val[3], min.reg.val[3], max.reg.val[3]);
533 + return r;
534 };
535
561 - FP32Vec16 clamp(const FP32Vec16& min, const FP32Vec16& max) const {
562 - return FP32Vec16(float32x4x4_t(
563 - {vminq_f32(max.reg.val[0], vmaxq_f32(min.reg.val[0], reg.val[0])),
564 - vminq_f32(max.reg.val[1], vmaxq_f32(min.reg.val[1], reg.val[1])),
565 - vminq_f32(max.reg.val[2], vmaxq_f32(min.reg.val[2], reg.val[2])),
566 - vminq_f32(max.reg.val[3], vmaxq_f32(min.reg.val[3], reg.val[3]))}));
536 + FORCE_INLINE FP32Vec16 min(const FP32Vec16& b) const {
537 + FP32Vec16 r(uninit);
538 + r.reg.val[0] = minimum(b.reg.val[0], reg.val[0]),
539 + r.reg.val[1] = minimum(b.reg.val[1], reg.val[1]);
540 + r.reg.val[2] = minimum(b.reg.val[2], reg.val[2]);
541 + r.reg.val[3] = minimum(b.reg.val[3], reg.val[3]);
542 + return r;
543 };
544
569 - FP32Vec16 max(const FP32Vec16& b) const {
570 - return FP32Vec16(float32x4x4_t({vmaxq_f32(b.reg.val[0], reg.val[0]),
571 - vmaxq_f32(b.reg.val[1], reg.val[1]),
572 - vmaxq_f32(b.reg.val[2], reg.val[2]),
573 - vmaxq_f32(b.reg.val[3], reg.val[3])}));
545 + FORCE_INLINE FP32Vec16 max(const FP32Vec16& b) const {
546 + FP32Vec16 r(uninit);
547 + r.reg.val[0] = maximum(b.reg.val[0], reg.val[0]);
548 + r.reg.val[1] = maximum(b.reg.val[1], reg.val[1]);
549 + r.reg.val[2] = maximum(b.reg.val[2], reg.val[2]);
550 + r.reg.val[3] = maximum(b.reg.val[3], reg.val[3]);
551 + return r;
552 };
553
576 - FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
577 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
578 - int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
579 - float32x4x4_t temp;
554 + FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
555 + size_t num_elements = reg.val[0].size();
556 +
557 + if (elem_num == VEC_ELEM_NUM) {
558 + return FP32Vec16::min(b);
559 + }
560 +
561 + int full_blocks = elem_num / num_elements;
562 + const int remainder = elem_num % num_elements;
563
564 + FP32Vec16 res(uninit);
565 for (int i = 0; i < full_blocks; i++)
582 - temp.val[i] = vmaxq_f32(b.reg.val[i], reg.val[i]);
566 + res.reg.val[i] = minimum(b.reg.val[i], reg.val[i]);
567
568 if (remainder > 0) {
585 - float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 0),
569 + float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 0),
570 vgetq_lane_f32(b.reg.val[full_blocks], 0));
587 - temp.val[full_blocks] = vsetq_lane_f32(max_v, temp.val[full_blocks], 0);
571 + res.reg.val[full_blocks] =
572 + vsetq_lane_f32(min_v, res.reg.val[full_blocks], 0);
573 }
574 if (remainder > 1) {
590 - float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 1),
575 + float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 1),
576 vgetq_lane_f32(b.reg.val[full_blocks], 1));
592 - temp.val[full_blocks] = vsetq_lane_f32(max_v, temp.val[full_blocks], 1);
577 + res.reg.val[full_blocks] =
578 + vsetq_lane_f32(min_v, res.reg.val[full_blocks], 1);
579 }
580 if (remainder > 2) {
595 - float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 2),
581 + float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 2),
582 vgetq_lane_f32(b.reg.val[full_blocks], 2));
597 - temp.val[full_blocks] = vsetq_lane_f32(max_v, temp.val[full_blocks], 2);
583 + res.reg.val[full_blocks] =
584 + vsetq_lane_f32(min_v, res.reg.val[full_blocks], 2);
585 }
599 - return FP32Vec16(temp);
600 - };
586
602 - FP32Vec16 min(const FP32Vec16& b) const {
603 - return FP32Vec16(float32x4x4_t({
604 - vminq_f32(b.reg.val[0], reg.val[0]),
605 - vminq_f32(b.reg.val[1], reg.val[1]),
606 - vminq_f32(b.reg.val[2], reg.val[2]),
607 - vminq_f32(b.reg.val[3], reg.val[3]),
608 - }));
587 + return res;
588 };
610 - FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
611 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
612 - const int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
613 - float32x4x4_t temp;
589 +
590 + FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
591 + size_t num_elements = reg.val[0].size();
592 +
593 + if (elem_num == VEC_ELEM_NUM) {
594 + return FP32Vec16::max(b);
595 + }
596 +
597 + int full_blocks = elem_num / num_elements;
598 + int remainder = elem_num % num_elements;
599 +
600 + FP32Vec16 res(uninit);
601 +
602 for (int i = 0; i < full_blocks; i++)
615 - temp.val[i] = vminq_f32(b.reg.val[i], reg.val[i]);
603 + res.reg.val[i] = maximum(b.reg.val[i], reg.val[i]);
604
605 if (remainder > 0) {
618 - float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 0),
606 + float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 0),
607 vgetq_lane_f32(b.reg.val[full_blocks], 0));
620 - temp.val[full_blocks] = vsetq_lane_f32(min_v, temp.val[full_blocks], 0);
608 + res.reg.val[full_blocks] =
609 + vsetq_lane_f32(max_v, res.reg.val[full_blocks], 0);
610 }
611 if (remainder > 1) {
623 - float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 1),
612 + float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 1),
613 vgetq_lane_f32(b.reg.val[full_blocks], 1));
625 - temp.val[full_blocks] = vsetq_lane_f32(min_v, temp.val[full_blocks], 1);
614 + res.reg.val[full_blocks] =
615 + vsetq_lane_f32(max_v, res.reg.val[full_blocks], 1);
616 }
617 if (remainder > 2) {
628 - float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 2),
618 + float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 2),
619 vgetq_lane_f32(b.reg.val[full_blocks], 2));
630 - temp.val[full_blocks] = vsetq_lane_f32(min_v, temp.val[full_blocks], 2);
620 + res.reg.val[full_blocks] =
621 + vsetq_lane_f32(max_v, res.reg.val[full_blocks], 2);
622 }
632 -
633 - return FP32Vec16(temp);
634 - };
635 - FP32Vec16 abs() const {
636 - return FP32Vec16(
637 - float32x4x4_t({vabsq_f32(reg.val[0]), vabsq_f32(reg.val[1]),
638 - vabsq_f32(reg.val[2]), vabsq_f32(reg.val[3])}));
639 - }
640 - float reduce_sum() const {
641 - AliasReg ar;
642 - ar.reg = reg;
643 - float answer = 0;
644 - unroll_loop<int, VEC_ELEM_NUM>(
645 - [&answer, &ar](int i) { answer += ar.values[i]; });
646 -
647 - return answer;
623 + return res;
624 };
625
626 float reduce_max() const {
651 - AliasReg ar;
652 - ar.reg = reg;
653 - float max_v = std::numeric_limits<float>::lowest();
654 - unroll_loop<int, VEC_ELEM_NUM>(
655 - [&max_v, &ar](int i) { max_v = std::max(max_v, ar.values[i]); });
656 - return max_v;
627 + VectorizedT max_vec = reg.val[0];
628 + unroll_loop<int, VEC_REG_NUM>([&](int i) {
629 + if (i > 0) max_vec = maximum(max_vec, reg.val[i]);
630 + });
631 +
632 + return vmaxvq_f32(max_vec);
633 }
634
635 float reduce_min() const {
660 - AliasReg ar;
661 - ar.reg = reg;
662 - float min_v = std::numeric_limits<float>::max();
663 - unroll_loop<int, VEC_ELEM_NUM>(
664 - [&min_v, &ar](int i) { min_v = std::min(min_v, ar.values[i]); });
665 - return min_v;
636 + VectorizedT min_vec = reg.val[0];
637 + unroll_loop<int, VEC_REG_NUM>([&](int i) {
638 + if (i > 0) min_vec = minimum(min_vec, reg.val[i]);
639 + });
640 +
641 + return vminvq_f32(min_vec);
642 }
643
644 template <int group_size>
645 float reduce_sub_sum(int idx) {
646 static_assert(VEC_ELEM_NUM % group_size == 0);
647
672 - AliasReg ar;
673 - ar.reg = reg;
648 + AliasReg<NxVectorizedTArray, ScalarT, VEC_ELEM_NUM> ar{reg};
649 float answer = 0;
650 const int start = idx * group_size;
651 unroll_loop<int, group_size>(
677 - [&answer, &start, ar](int i) { answer += ar.values[start + i]; });
652 + [&](int i) { answer += ar.values[start + i]; });
653
654 return answer;
655 };
656
682 - void save(float* ptr) const {
683 - vst1q_f32(ptr, reg.val[0]);
684 - vst1q_f32(ptr + 4, reg.val[1]);
685 - vst1q_f32(ptr + 8, reg.val[2]);
686 - vst1q_f32(ptr + 12, reg.val[3]);
687 - };
688 -
689 - void save(float* ptr, const int elem_num) const {
690 - int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
691 - int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
692 -
693 - for (int i = 0; i < full_blocks; i++)
694 - vst1q_f32(
695 - reinterpret_cast<float32_t*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
696 - reg.val[i]);
657 + float reduce_sum() const {
658 + float answer = 0;
659 + std::plus<VectorizedT> add;
660 + unroll_loop<int, VEC_REG_NUM>([&](int i) {
661 + answer += at::vec::vec_reduce_all<float>(add, reg.val[i]);
662 + });
663
698 - if (remainder > 0) {
699 - float32x4_t temp = reg.val[full_blocks];
700 - float* base = reinterpret_cast<float32_t*>(ptr) +
701 - full_blocks * NUM_ELEMENTS_REG(reg.val[0]);
702 - if (remainder > 0) base[0] = vgetq_lane_f32(temp, 0);
703 - if (remainder > 1) base[1] = vgetq_lane_f32(temp, 1);
704 - if (remainder > 2) base[2] = vgetq_lane_f32(temp, 2);
705 - }
664 + return answer;
665 }
666 };
667
668 +// Only used for int types for now could be replaced when
669 +// int8/32 vectorised ops are added in ATen
670 +template <typename T>
671 +struct Vec {
672 + constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
673 +};
674 +
675 struct INT8Vec16 : public Vec<INT8Vec16> {
676 constexpr static int VEC_ELEM_NUM = 16;
677 union AliasReg {
@@ -854,30 +820,47 @@ struct INT8Vec64 : public Vec<INT8Vec64> {
820 void nt_save(int8_t* ptr) const { save(ptr); }
821 }; // INT8Vec64
822
857 -template <typename T>
858 -struct VecType {
859 - using vec_type = void;
860 -};
823 +struct INT32Vec16 : public Vec<INT32Vec16> {
824 + constexpr static int VEC_ELEM_NUM = 16;
825 + union AliasReg {
826 + int32x4x4_t reg;
827 + int32_t values[VEC_ELEM_NUM];
828 + };
829 + int32x4x4_t reg;
830
862 -template <typename T>
863 -using vec_t = typename VecType<T>::vec_type;
831 + explicit INT32Vec16(const void* ptr) {
832 + reg.val[0] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr));
833 + reg.val[1] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 4);
834 + reg.val[2] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 8);
835 + reg.val[3] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 12);
836 + }
837
865 -template <>
866 -struct VecType<float> {
867 - using vec_type = FP32Vec8;
868 -};
838 + void save(int32_t* ptr) const {
839 + vst1q_s32(ptr, reg.val[0]);
840 + vst1q_s32(ptr + 4, reg.val[1]);
841 + vst1q_s32(ptr + 8, reg.val[2]);
842 + vst1q_s32(ptr + 12, reg.val[3]);
843 + };
844
870 -template <>
871 -struct VecType<c10::Half> {
872 - using vec_type = FP16Vec8;
873 -};
845 + void save(int32_t* ptr, const int elem_num) const {
846 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
847 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
848
875 -#ifdef ARM_BF16_SUPPORT
876 -template <>
877 -struct VecType<c10::BFloat16> {
878 - using vec_type = BF16Vec8;
849 + for (int i = 0; i < full_blocks; i++)
850 + vst1q_s32(
851 + reinterpret_cast<__int32_t*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
852 + reg.val[i]);
853 +
854 + if (remainder > 0) {
855 + int32x4_t temp = reg.val[full_blocks];
856 + int32_t* base = reinterpret_cast<int32_t*>(ptr) + full_blocks * 4;
857 + if (remainder > 0) base[0] = vgetq_lane_s32(temp, 0);
858 + if (remainder > 1) base[1] = vgetq_lane_s32(temp, 1);
859 + if (remainder > 2) base[2] = vgetq_lane_s32(temp, 2);
860 + if (remainder > 3) base[3] = vgetq_lane_s32(temp, 3);
861 + }
862 + }
863 };
880 -#endif
864
865 template <typename T>
866 void storeFP32(float v, T* ptr) {
@@ -889,66 +872,55 @@ inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
872 *reinterpret_cast<__fp16*>(ptr) = v;
873 }
874
892 -inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
893 - float16x4_t low_0 = vcvt_f16_f32(v.reg.val[0]);
894 - float16x4_t high_0 = vcvt_f16_f32(v.reg.val[1]);
895 - float16x4_t low_1 = vcvt_f16_f32(v.reg.val[2]);
896 - float16x4_t high_1 = vcvt_f16_f32(v.reg.val[3]);
875 +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
876 + reg.val[0] = convert_float_half(v.reg.val[0], v.reg.val[1]);
877 +};
878
898 - reg.val[0] = vcombine_f16(low_0, high_0);
899 - reg.val[1] = vcombine_f16(low_1, high_1);
879 +inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
880 + reg.val[0] = convert_float_half(v.reg.val[0], v.reg.val[1]);
881 + reg.val[1] = convert_float_half(v.reg.val[2], v.reg.val[3]);
882 };
883
902 -inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
903 - float16x4_t lower_half = vcvt_f16_f32(v.reg.val[0]);
904 - float16x4_t upper_half = vcvt_f16_f32(v.reg.val[1]);
884 +inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
885 + fmadd(acc.reg.val[0], a.reg.val[0], b.reg.val[0]);
886 + fmadd(acc.reg.val[1], a.reg.val[1], b.reg.val[1]);
887 + fmadd(acc.reg.val[2], a.reg.val[2], b.reg.val[2]);
888 + fmadd(acc.reg.val[3], a.reg.val[3], b.reg.val[3]);
889 +};
890
906 - reg = vcombine_f16(lower_half, upper_half);
891 +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
892 + reg.val[0] = convert_float_bfloat16(v.reg.val[0], v.reg.val[1]);
893 };
894
909 -inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
910 - acc.reg.val[0] = vfmaq_f32(acc.reg.val[0], a.reg.val[0], b.reg.val[0]);
911 - acc.reg.val[1] = vfmaq_f32(acc.reg.val[1], a.reg.val[1], b.reg.val[1]);
912 - acc.reg.val[2] = vfmaq_f32(acc.reg.val[2], a.reg.val[2], b.reg.val[2]);
913 - acc.reg.val[3] = vfmaq_f32(acc.reg.val[3], a.reg.val[3], b.reg.val[3]);
895 +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
896 + reg.val[0] = convert_float_bfloat16(v.reg.val[0], v.reg.val[1]);
897 + reg.val[1] = convert_float_bfloat16(v.reg.val[2], v.reg.val[3]);
898 };
899
916 -#ifdef ARM_BF16_SUPPORT
900 inline void fma(FP32Vec16& acc, BF16Vec32& a, BF16Vec32& b) {
918 - float32x4_t a0_low = vcvt_f32_bf16(vget_low_bf16(a.reg.val[0]));
919 - float32x4_t a0_high = vcvt_f32_bf16(vget_high_bf16(a.reg.val[0]));
920 - float32x4_t a1_low = vcvt_f32_bf16(vget_low_bf16(a.reg.val[1]));
921 - float32x4_t a1_high = vcvt_f32_bf16(vget_high_bf16(a.reg.val[1]));
922 -
923 - float32x4_t b0_low = vcvt_f32_bf16(vget_low_bf16(b.reg.val[0]));
924 - float32x4_t b0_high = vcvt_f32_bf16(vget_high_bf16(b.reg.val[0]));
925 - float32x4_t b1_low = vcvt_f32_bf16(vget_low_bf16(b.reg.val[1]));
926 - float32x4_t b1_high = vcvt_f32_bf16(vget_high_bf16(b.reg.val[1]));
927 -
928 - acc.reg.val[0] = vfmaq_f32(acc.reg.val[0], a0_low, b0_low);
929 - acc.reg.val[1] = vfmaq_f32(acc.reg.val[1], a0_high, b0_high);
930 - acc.reg.val[2] = vfmaq_f32(acc.reg.val[2], a1_low, b1_low);
931 - acc.reg.val[3] = vfmaq_f32(acc.reg.val[3], a1_high, b1_high);
901 + Vectorized<float> a0_low, a0_high, a1_low, a1_high, b0_low, b0_high, b1_low,
902 + b1_high;
903 +
904 + std::tie(a0_low, a0_high) = convert_bfloat16_float(a.reg.val[0]);
905 + std::tie(a1_low, a1_high) = convert_bfloat16_float(a.reg.val[1]);
906 + std::tie(b0_low, b0_high) = convert_bfloat16_float(b.reg.val[0]);
907 + std::tie(b1_low, b1_high) = convert_bfloat16_float(b.reg.val[1]);
908 +
909 + fmadd(acc.reg.val[0], a0_low, b0_low);
910 + fmadd(acc.reg.val[1], a0_high, b0_high);
911 + fmadd(acc.reg.val[2], a1_low, b1_low);
912 + fmadd(acc.reg.val[3], a1_high, b1_high);
913 };
933 -#endif
914
915 +template <>
916 +inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
917 #ifdef ARM_BF16_SUPPORT
936 -inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
937 - : reg(vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.reg.val[0]), v.reg.val[1])) {
938 - };
939 -
940 -inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
941 - : reg({vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.reg.val[0]), v.reg.val[1]),
942 - vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.reg.val[2]),
943 - v.reg.val[3])}) {};
918 + *reinterpret_cast<__bf16*>(ptr) = vcvth_bf16_f32(v);
919 +#else
920 + *ptr = static_cast<c10::BFloat16>(v);
921 #endif
922 +};
923
924 inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); };
925
948 -#ifdef ARM_BF16_SUPPORT
949 -template <>
950 -inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
951 - *reinterpret_cast<__bf16*>(ptr) = vcvth_bf16_f32(v);
952 -};
953 -#endif
926 }; // namespace vec_op
\ No newline at end of file
csrc/cpu/dnnl_kernels.cpp
-2
@@ -14,13 +14,11 @@ struct KernelVecType<float> {
14 using cvt_vec_type = vec_op::FP32Vec16;
15 };
16
17 -#if !defined(__aarch64__) || defined(ARM_BF16_SUPPORT)
17 template <>
18 struct KernelVecType<c10::BFloat16> {
19 using load_vec_type = vec_op::BF16Vec16;
20 using cvt_vec_type = vec_op::FP32Vec16;
21 };
23 -#endif
22
23 template <>
24 struct KernelVecType<c10::Half> {
csrc/cpu/mla_decode.cpp
+1 -3
@@ -38,9 +38,7 @@ struct KernelVecType<c10::BFloat16> {
38 using qk_vec_type = vec_op::BF16Vec32;
39 using v_load_vec_type = vec_op::BF16Vec16;
40 };
41 -#elif defined(__aarch64__) && !defined(ARM_BF16_SUPPORT)
42 -// pass
43 -#else
41 +#elif defined(__aarch64__)
42 template <>
43 struct KernelVecType<c10::BFloat16> {
44 using qk_load_vec_type = vec_op::BF16Vec16;
csrc/cpu/utils.hpp
-2
@@ -30,12 +30,10 @@ struct VecTypeTrait<float> {
30 using vec_t = vec_op::FP32Vec16;
31 };
32
33 -#if !defined(__aarch64__) || defined(ARM_BF16_SUPPORT)
33 template <>
34 struct VecTypeTrait<c10::BFloat16> {
35 using vec_t = vec_op::BF16Vec16;
36 };
38 -#endif
37
38 #if !defined(__powerpc__)
39 template <>