PyTorch includes avx_mathfun.h, which includes an 8-way vectorized implementation of exp (exp256_ps). We'd like to use that to implement a vectorized sigmoid function (accessible from torch using torch.sigmoid). Currently the sigmoid function is not vectorized.
To get started, look into how other vectorized functions, like cadd, are implemented:
Actual implementation:
https://github.com/pytorch/pytorch/blob/77c792ec276ee8bf9e279ce34ecb8dac5ecbf472/aten/src/TH/vector/AVX2.c
Code to figure out whether AVX2 is availiable dynamically and dispatch:
|
static void (*THVector_(cadd_DISPATCHPTR))(real *, const real *, const real *, const real, const ptrdiff_t) = &THVector_(cadd_DEFAULT); |
Code that implements non-vectorized default:
|
void THVector_(cadd_DEFAULT)(real *z, const real *x, const real *y, const real c, const ptrdiff_t n) |
Current place where sigmoid (non-vectorized) is created (you will need to modify this to do correct AVX2 dispatch):
|
VECTOR_IMPLEMENT_FUNCTION(sigmoid,TH_MATH_NAME(TH_sigmoid)) |
@vedanuj
PyTorch includes avx_mathfun.h, which includes an 8-way vectorized implementation of
exp(exp256_ps). We'd like to use that to implement a vectorized sigmoid function (accessible from torch usingtorch.sigmoid). Currently the sigmoid function is not vectorized.To get started, look into how other vectorized functions, like cadd, are implemented:
Actual implementation:
https://github.com/pytorch/pytorch/blob/77c792ec276ee8bf9e279ce34ecb8dac5ecbf472/aten/src/TH/vector/AVX2.c
Code to figure out whether AVX2 is availiable dynamically and dispatch:
pytorch/aten/src/TH/generic/THVectorDispatch.c
Line 47 in 77c792e
Code that implements non-vectorized default:
pytorch/aten/src/TH/generic/THVectorDefault.c
Line 37 in 77c792e
Current place where sigmoid (non-vectorized) is created (you will need to modify this to do correct AVX2 dispatch):
pytorch/aten/src/TH/generic/THVectorDefault.c
Line 238 in 77c792e
@vedanuj