Skip to content

Commit d0fb76e

Browse files
sterrettm2sterrettm2
committed
Moved more things into #if guards and cleaned up implementation
1 parent f8e6ca6 commit d0fb76e

1 file changed

Lines changed: 76 additions & 96 deletions

File tree

‎numpy/_core/src/umath/loops_hyperbolic.dispatch.cpp.src‎

Lines changed: 76 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,12 @@
1111
#include "loops_utils.h"
1212
#include "loops.h"
1313
#include "fast_loop_macros.h"
14-
#include <hwy/highway.h>
1514

1615

1716
#if NPY_SIMD_FMA3 && (NPY_SIMD_F32 || NPY_SIMD_F64) // requires native FMA
17+
18+
#include <hwy/highway.h>
19+
1820
/*
1921
* NOTE: The following implementation of tanh(f32, f64) have been converted from
2022
* Intel SVML to universal intrinsics, and the original code can be found in:
@@ -81,97 +83,18 @@
8183
namespace hn = hwy::HWY_NAMESPACE;
8284

8385
const hn::ScalableTag<float> f32;
84-
const hn::ScalableTag<double> f64;
8586
const hn::ScalableTag<int32_t> s32;
86-
const hn::ScalableTag<int64_t> s64;
8787
const hn::ScalableTag<uint32_t> u32;
88-
const hn::ScalableTag<uint64_t> u64;
8988
using vec_f32 = hn::Vec<decltype(f32)>;
90-
using vec_f64 = hn::Vec<decltype(f64)>;
9189
using vec_s32 = hn::Vec<decltype(s32)>;
92-
using vec_s64 = hn::Vec<decltype(s64)>;
9390
using vec_u32 = hn::Vec<decltype(u32)>;
94-
using vec_u64 = hn::Vec<decltype(u64)>;
95-
using opmask_t = hn::Mask<decltype(f32)>;
96-
97-
struct hwy_f32x2 {
98-
vec_f32 val[2];
99-
};
10091

101-
HWY_ATTR NPY_FINLINE hwy_f32x2 zip_f32(vec_f32 a, vec_f32 b){
102-
hwy_f32x2 res;
103-
res.val[0] = hn::InterleaveLower(f32, a, b);
104-
res.val[1] = hn::InterleaveUpper(f32, a, b);
105-
return res;
106-
}
107-
108-
HWY_ATTR NPY_FINLINE vec_f64 lut_16_f64(const double * lut, vec_u64 idx){
109-
if constexpr(hn::Lanes(f64) == 8){
110-
const vec_f64 lut0 = hn::Load(f64, lut);
111-
const vec_f64 lut1 = hn::Load(f64, lut + 8);
112-
return hn::TwoTablesLookupLanes(f64, lut0, lut1, hn::IndicesFromVec(f64, idx));
113-
}else if constexpr (hn::Lanes(f64) == 4){
114-
const vec_f64 lut0 = hn::Load(f64, lut);
115-
const vec_f64 lut1 = hn::Load(f64, lut + 4);
116-
const vec_f64 lut2 = hn::Load(f64, lut + 8);
117-
const vec_f64 lut3 = hn::Load(f64, lut + 12);
118-
119-
const auto high_mask = hn::Ne(hn::ShiftRight<3>(idx), hn::Zero(u64));
120-
const auto load_mask = hn::And(idx, hn::Set(u64, 0b111));
121-
122-
const vec_f64 lut_low = hn::TwoTablesLookupLanes(f64, lut0, lut1, hn::IndicesFromVec(f64, load_mask));
123-
const vec_f64 lut_high = hn::TwoTablesLookupLanes(f64, lut2, lut3, hn::IndicesFromVec(f64, load_mask));
124-
125-
return hn::IfThenElse(hn::RebindMask(f64, high_mask), lut_high, lut_low);
126-
}else{
127-
return hn::GatherIndex(f64, lut, hn::BitCast(s64, idx));
128-
}
129-
}
130-
131-
HWY_ATTR NPY_FINLINE vec_f32 lut_32_f32(const float * lut, vec_u32 idx){
132-
if constexpr(hn::Lanes(f32) == 16){
133-
const vec_f32 lut0 = hn::Load(f32, lut);
134-
const vec_f32 lut1 = hn::Load(f32, lut + 16);
135-
return hn::TwoTablesLookupLanes(f32, lut0, lut1, hn::IndicesFromVec(f32, idx));
136-
}else if constexpr (hn::Lanes(f32) == 8){
137-
const vec_f32 lut0 = hn::Load(f32, lut);
138-
const vec_f32 lut1 = hn::Load(f32, lut + 8);
139-
const vec_f32 lut2 = hn::Load(f32, lut + 16);
140-
const vec_f32 lut3 = hn::Load(f32, lut + 24);
141-
142-
const auto high_mask = hn::Ne(hn::ShiftRight<4>(idx), hn::Zero(u32));
143-
const auto load_mask = hn::And(idx, hn::Set(u32, 0b1111));
144-
145-
const vec_f32 lut_low = hn::TwoTablesLookupLanes(f32, lut0, lut1, hn::IndicesFromVec(f32, load_mask));
146-
const vec_f32 lut_high = hn::TwoTablesLookupLanes(f32, lut2, lut3, hn::IndicesFromVec(f32, load_mask));
147-
148-
return hn::IfThenElse(hn::RebindMask(f32, high_mask), lut_high, lut_low);
149-
}else{
150-
return hn::GatherIndex(f32, lut, hn::BitCast(s32, idx));
151-
}
152-
}
153-
154-
template <typename vtype>
155-
HWY_ATTR NPY_FINLINE vtype
156-
combine_lower(vtype a, vtype b){
157-
vec_f64 a64 = hn::BitCast(f64, a);
158-
vec_f64 b64 = hn::BitCast(f64, b);
159-
160-
auto D = hn::DFromV<vtype>();
161-
162-
return hn::BitCast(D, hn::InterleaveLower(f64, a64, b64));
163-
}
164-
165-
template <typename vtype>
166-
HWY_ATTR NPY_FINLINE vtype
167-
combine_upper(vtype a, vtype b){
168-
vec_f64 a64 = hn::BitCast(f64, a);
169-
vec_f64 b64 = hn::BitCast(f64, b);
170-
171-
auto D = hn::DFromV<vtype>();
172-
173-
return hn::BitCast(D, hn::InterleaveUpper(f64, a64, b64));
174-
}
92+
const hn::ScalableTag<double> f64;
93+
const hn::ScalableTag<int64_t> s64;
94+
const hn::ScalableTag<uint64_t> u64;
95+
using vec_f64 = hn::Vec<decltype(f64)>;
96+
using vec_s64 = hn::Vec<decltype(s64)>;
97+
using vec_u64 = hn::Vec<decltype(u64)>;
17598

17699
template <typename vtype, typename type_t>
177100
HWY_ATTR NPY_FINLINE vtype
@@ -243,6 +166,29 @@ store_vector(vtype vec, type_t* dst, npy_intp sdst, npy_intp len){
243166

244167
#if NPY_SIMD_F64
245168

169+
HWY_ATTR NPY_FINLINE vec_f64 lut_16_f64(const double * lut, vec_u64 idx){
170+
if constexpr(hn::Lanes(f64) == 8){
171+
const vec_f64 lut0 = hn::Load(f64, lut);
172+
const vec_f64 lut1 = hn::Load(f64, lut + 8);
173+
return hn::TwoTablesLookupLanes(f64, lut0, lut1, hn::IndicesFromVec(f64, idx));
174+
}else if constexpr (hn::Lanes(f64) == 4){
175+
const vec_f64 lut0 = hn::Load(f64, lut);
176+
const vec_f64 lut1 = hn::Load(f64, lut + 4);
177+
const vec_f64 lut2 = hn::Load(f64, lut + 8);
178+
const vec_f64 lut3 = hn::Load(f64, lut + 12);
179+
180+
const auto high_mask = hn::Ne(hn::ShiftRight<3>(idx), hn::Zero(u64));
181+
const auto load_mask = hn::And(idx, hn::Set(u64, 0b111));
182+
183+
const vec_f64 lut_low = hn::TwoTablesLookupLanes(f64, lut0, lut1, hn::IndicesFromVec(f64, load_mask));
184+
const vec_f64 lut_high = hn::TwoTablesLookupLanes(f64, lut2, lut3, hn::IndicesFromVec(f64, load_mask));
185+
186+
return hn::IfThenElse(hn::RebindMask(f64, high_mask), lut_high, lut_low);
187+
}else{
188+
return hn::GatherIndex(f64, lut, hn::BitCast(s64, idx));
189+
}
190+
}
191+
246192
// For architectures without efficient gather / scatter instructions, it is
247193
// better to use a transposed LUT where we can load all coefficients for an
248194
// index linearly. In order to keep the same vertical calculation, we
@@ -482,8 +428,8 @@ simd_tanh_f64(const double *src, npy_intp ssrc, double *dst, npy_intp sdst, npy_
482428
*/
483429
e0e1[@lane@] = hn::BitCast(f64, hn::LoadU(u64, lut18x16 + index[@lane@] * 18 + @off@));
484430
/**end repeat1**/
485-
vec_f64 @e0@ = combine_lower<vec_f64>(e0e1[0], e0e1[1]);
486-
vec_f64 @e1@ = combine_upper<vec_f64>(e0e1[0], e0e1[1]);
431+
vec_f64 @e0@ = hn::ConcatLowerLower(f64, e0e1[1], e0e1[0]);
432+
vec_f64 @e1@ = hn::ConcatUpperUpper(f64, e0e1[1], e0e1[0]);
487433
/**end repeat**/
488434
#else
489435
vec_f64 b = lut_16_f64((const double*)lut16x18 + 16*0, idx);
@@ -542,6 +488,40 @@ simd_tanh_f64(const double *src, npy_intp ssrc, double *dst, npy_intp sdst, npy_
542488

543489
#if NPY_SIMD_F32
544490

491+
struct hwy_f32x2 {
492+
vec_f32 val[2];
493+
};
494+
495+
HWY_ATTR NPY_FINLINE hwy_f32x2 zip_f32(vec_f32 a, vec_f32 b){
496+
hwy_f32x2 res;
497+
res.val[0] = hn::InterleaveLower(f32, a, b);
498+
res.val[1] = hn::InterleaveUpper(f32, a, b);
499+
return res;
500+
}
501+
502+
HWY_ATTR NPY_FINLINE vec_f32 lut_32_f32(const float * lut, vec_u32 idx){
503+
if constexpr(hn::Lanes(f32) == 16){
504+
const vec_f32 lut0 = hn::Load(f32, lut);
505+
const vec_f32 lut1 = hn::Load(f32, lut + 16);
506+
return hn::TwoTablesLookupLanes(f32, lut0, lut1, hn::IndicesFromVec(f32, idx));
507+
}else if constexpr (hn::Lanes(f32) == 8){
508+
const vec_f32 lut0 = hn::Load(f32, lut);
509+
const vec_f32 lut1 = hn::Load(f32, lut + 8);
510+
const vec_f32 lut2 = hn::Load(f32, lut + 16);
511+
const vec_f32 lut3 = hn::Load(f32, lut + 24);
512+
513+
const auto high_mask = hn::Ne(hn::ShiftRight<4>(idx), hn::Zero(u32));
514+
const auto load_mask = hn::And(idx, hn::Set(u32, 0b1111));
515+
516+
const vec_f32 lut_low = hn::TwoTablesLookupLanes(f32, lut0, lut1, hn::IndicesFromVec(f32, load_mask));
517+
const vec_f32 lut_high = hn::TwoTablesLookupLanes(f32, lut2, lut3, hn::IndicesFromVec(f32, load_mask));
518+
519+
return hn::IfThenElse(hn::RebindMask(f32, high_mask), lut_high, lut_low);
520+
}else{
521+
return hn::GatherIndex(f32, lut, hn::BitCast(s32, idx));
522+
}
523+
}
524+
545525
// For architectures without efficient gather / scatter instructions, it is
546526
// better to use a transposed LUT where we can load all coefficients for an
547527
// index linearly. In order to keep the same vertical calculation, we
@@ -696,17 +676,17 @@ simd_tanh_f32(const float *src, npy_intp ssrc, float *dst, npy_intp sdst, npy_in
696676

697677
hwy_f32x2 c6543_l01 = zip_f32(c6543[0], c6543[1]);
698678
hwy_f32x2 c6543_l23 = zip_f32(c6543[2], c6543[3]);
699-
vec_f32 c6 = combine_lower<vec_f32>(c6543_l01.val[0], c6543_l23.val[0]);
700-
vec_f32 c5 = combine_upper<vec_f32>(c6543_l01.val[0], c6543_l23.val[0]);
701-
vec_f32 c4 = combine_lower<vec_f32>(c6543_l01.val[1], c6543_l23.val[1]);
702-
vec_f32 c3 = combine_upper<vec_f32>(c6543_l01.val[1], c6543_l23.val[1]);
679+
vec_f32 c6 = hn::ConcatLowerLower(f32, c6543_l23.val[0], c6543_l01.val[0]);
680+
vec_f32 c5 = hn::ConcatUpperUpper(f32, c6543_l23.val[0], c6543_l01.val[0]);
681+
vec_f32 c4 = hn::ConcatLowerLower(f32, c6543_l23.val[1], c6543_l01.val[1]);
682+
vec_f32 c3 = hn::ConcatUpperUpper(f32, c6543_l23.val[1], c6543_l01.val[1]);
703683

704684
hwy_f32x2 c210b_l01 = zip_f32(c210b[0], c210b[1]);
705685
hwy_f32x2 c210b_l23 = zip_f32(c210b[2], c210b[3]);
706-
vec_f32 c2 = combine_lower<vec_f32>(c210b_l01.val[0], c210b_l23.val[0]);
707-
vec_f32 c1 = combine_upper<vec_f32>(c210b_l01.val[0], c210b_l23.val[0]);
708-
vec_f32 c0 = combine_lower<vec_f32>(c210b_l01.val[1], c210b_l23.val[1]);
709-
vec_f32 b = combine_upper<vec_f32>(c210b_l01.val[1], c210b_l23.val[1]);
686+
vec_f32 c2 = hn::ConcatLowerLower(f32, c210b_l23.val[0], c210b_l01.val[0]);
687+
vec_f32 c1 = hn::ConcatUpperUpper(f32, c210b_l23.val[0], c210b_l01.val[0]);
688+
vec_f32 c0 = hn::ConcatLowerLower(f32, c210b_l23.val[1], c210b_l01.val[1]);
689+
vec_f32 b = hn::ConcatUpperUpper(f32, c210b_l23.val[1], c210b_l01.val[1]);
710690
#else
711691
vec_f32 b = lut_32_f32((const float*)lut32x8 + 32*0, idx);
712692
vec_f32 c0 = lut_32_f32((const float*)lut32x8 + 32*1, idx);

0 commit comments

Comments
 (0)