/usr/local/lib64/python3.6/site-packages/torch/include/ATen/cpu/vec
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/cpu/vec/vec_base.h (36871B)
#pragma once
// DO NOT DEFINE STATIC DATA IN THIS HEADER!
// See Note [Do not compile initializers with AVX]
//
// Note [Do not compile initializers with AVX]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// If you define a static initializer in this file, the initialization will use
// AVX instructions because these object files are compiled with AVX enabled.
// We need to avoid non-trivial global data in these architecture specific files
// because there's no way to guard the global initializers with CPU capability
// detection.
//
// See https://github.com/pytorch/pytorch/issues/37577 for an instance
// of this bug in the past.
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
// These macros helped us unify vec_base.h
#ifdef CPU_CAPABILITY_AVX512
#if defined(__GNUC__)
#define __at_align__ __attribute__((aligned(64)))
#elif defined(_WIN32)
#define __at_align__ __declspec(align(64))
#else
#define __at_align__
#endif
#define VECTOR_WIDTH 64
#define int_vector __m512i
#else // CPU_CAPABILITY_AVX512
#if defined(__GNUC__)
#define __at_align__ __attribute__((aligned(32)))
#elif defined(_WIN32)
#define __at_align__ __declspec(align(32))
#else
#define __at_align__
#endif
#define VECTOR_WIDTH 32
#define int_vector __m256i
#endif // CPU_CAPABILITY_AVX512
namespace at {
namespace vec {
// See Note [Acceptable use of anonymous namespace in header]
namespace {
// at::Half and at::BFloat16 should be treated as floating point
template
struct is_floating_point:
std::integral_constant::value ||
std::is_same::value ||
std::is_same::value> {
};
template struct int_of_size;
#define DEFINE_INT_OF_SIZE(int_t) \
template<> struct int_of_size { using type = int_t; }
DEFINE_INT_OF_SIZE(int64_t);
DEFINE_INT_OF_SIZE(int32_t);
DEFINE_INT_OF_SIZE(int16_t);
DEFINE_INT_OF_SIZE(int8_t);
#undef DEFINE_INT_OF_SIZE
template
using int_same_size_t = typename int_of_size::type;
// NOTE: If you specialize on a type, you must define all operations!
// emulates Vectorized types
template
struct Vectorized {
private:
__at_align__ T values[VECTOR_WIDTH / sizeof(T)];
public:
using value_type = T;
using size_type = int;
// Note [constexpr static function to avoid odr-usage compiler bug]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Why, you might ask, is size defined to be a static constexpr function,
// rather than a more ordinary 'static constexpr int size;' variable?
// The problem lies within ODR rules for static constexpr members versus
// static constexpr functions. First, recall that this class (along with all
// of its derivations) live in an anonymous namespace: they are intended to be
// *completely* inlined at their use-sites, because we need to compile it
// multiple times for different instruction sets.
//
// Because of this constraint, we CANNOT provide a single definition for
// any static members in this class; since we want to compile the class
// multiple times, there wouldn't actually be any good place to put the
// definition. Now here is the problem: if we ODR-use a static constexpr
// member, we are *obligated* to provide a definition. Without the
// definition, you get a compile error like:
//
// relocation R_X86_64_PC32 against undefined symbol
// `_ZN2at6vec25612_GLOBAL__N_16VectorizedIdE4sizeE' can not be used when making
// a shared object; recompile with -fPIC
//
// If this were C++17, we could replace a static constexpr variable with
// an inline variable which doesn't require one definition. But we are not
// C++17. So the next best thing is to replace the member with a static
// constexpr (and therefore inline) function, which does not require ODR
// either.
//
// Also, technically according to the C++ standard, we don't have to define
// a constexpr variable if we never odr-use it. But it seems that some
// versions GCC/Clang have buggy determinations on whether or not an
// identifier is odr-used or not, and in any case it's hard to tell if
// a variable is odr-used or not. So best to just cut the problem at the root.
static constexpr size_type size() {
return VECTOR_WIDTH / sizeof(T);
}
Vectorized() : values{0} {}
Vectorized(T val) {
for (int i = 0; i != size(); i++) {
values[i] = val;
}
}
template>
Vectorized(Args... vals) : values{vals...}{
}
// This also implies const T& operator[](int idx) const
inline operator const T*() const {
return values;
}
// This also implies T& operator[](int idx)
inline operator T*() {
return values;
}
// Return the values as char* for type punning
auto as_bytes() const -> const char* {
return reinterpret_cast(values);
}
template
static Vectorized blend(const Vectorized& a, const Vectorized& b) {
int64_t mask = mask_;
Vectorized vector;
for (int64_t i = 0; i < size(); i++) {
if (mask & 0x01) {
vector[i] = b[i];
} else {
vector[i] = a[i];
}
mask = mask >> 1;
}
return vector;
}
static Vectorized blendv(const Vectorized& a, const Vectorized& b,
const Vectorized& mask) {
Vectorized vector;
int_same_size_t buffer[size()];
mask.store(buffer);
for (int64_t i = 0; i < size(); i++) {
if (buffer[i] & 0x01)
{
vector[i] = b[i];
} else {
vector[i] = a[i];
}
}
return vector;
}
template // step sometimes requires a higher precision type (e.g., T=int, step_t=double)
static Vectorized arange(T base = static_cast(0), step_t step = static_cast(1)) {
Vectorized vector;
for (int64_t i = 0; i < size(); i++) {
vector.values[i] = base + i * step;
}
return vector;
}
static Vectorized set(const Vectorized& a, const Vectorized& b, int64_t count = size()) {
Vectorized vector;
for (int64_t i = 0; i < size(); i++) {
if (i < count) {
vector[i] = b[i];
} else {
vector[i] = a[i];
}
}
return vector;
}
static Vectorized loadu(const void* ptr) {
Vectorized vector;
std::memcpy(vector.values, ptr, VECTOR_WIDTH);
return vector;
}
static Vectorized loadu(const void* ptr, int64_t count) {
Vectorized vector;
std::memcpy(vector.values, ptr, count * sizeof(T));
return vector;
}
void store(void* ptr, int count = size()) const {
std::memcpy(ptr, values, count * sizeof(T));
}
int zero_mask() const {
// returns an integer mask where all zero elements are translated to 1-bit and others are translated to 0-bit
int mask = 0;
for (int i = 0; i < size(); ++ i) {
if (values[i] == static_cast(0)) {
mask |= (1 << i);
}
}
return mask;
}
Vectorized isnan() const {
Vectorized vector;
for (int64_t i = 0; i != size(); i++) {
if (_isnan(values[i])) {
std::memset(static_cast(vector.values + i), 0xFF, sizeof(T));
} else {
std::memset(static_cast(vector.values + i), 0, sizeof(T));
}
}
return vector;
}
Vectorized map(T (*const f)(T)) const {
Vectorized ret;
for (int64_t i = 0; i != size(); i++) {
ret[i] = f(values[i]);
}
return ret;
}
Vectorized map(T (*const f)(const T &)) const {
Vectorized ret;
for (int64_t i = 0; i != size(); i++) {
ret[i] = f(values[i]);
}
return ret;
}
template ::value && !c10::is_complex::value, int>::type = 0>
Vectorized abs() const {
// other_t_abs is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "other_t_abs must be T");
return map([](T x) -> T { return x < static_cast(0) ? -x : x; });
}
template ::value, int>::type = 0>
Vectorized abs() const {
// float_t_abs is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "float_t_abs must be T");
// Specifically deal with floating-point because the generic code above won't handle -0.0 (which should result in
// 0.0) properly.
return map([](T x) -> T { return std::abs(x); });
}
template ::value, int>::type = 0>
Vectorized abs() const {
// complex_t_abs is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "complex_t_abs must be T");
// Specifically map() does not perform the type conversion needed by abs.
return map([](T x) { return static_cast(std::abs(x)); });
}
template ::value, int>::type = 0>
Vectorized sgn() const {
return map(at::native::sgn_impl);
}
template ::value, int>::type = 0>
Vectorized angle() const {
// other_t_angle is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "other_t_angle must be T");
return map(at::native::angle_impl); // compiler is unable to resolve the overload without
}
template ::value, int>::type = 0>
Vectorized angle() const {
// complex_t_angle is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "complex_t_angle must be T");
return map([](T x) { return static_cast(std::arg(x)); });
}
template ::value, int>::type = 0>
Vectorized real() const {
// other_t_real is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "other_t_real must be T");
return *this;
}
template ::value, int>::type = 0>
Vectorized real() const {
// complex_t_real is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "complex_t_real must be T");
return map([](T x) { return static_cast(x.real()); });
}
template ::value, int>::type = 0>
Vectorized imag() const {
// other_t_imag is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "other_t_imag must be T");
return Vectorized(0);
}
template ::value, int>::type = 0>
Vectorized imag() const {
// complex_t_imag is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "complex_t_imag must be T");
return map([](T x) { return static_cast(x.imag()); });
}
template ::value, int>::type = 0>
Vectorized conj() const {
// other_t_conj is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "other_t_conj must be T");
return *this;
}
template ::value, int>::type = 0>
Vectorized conj() const {
// complex_t_conj is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "complex_t_conj must be T");
return map([](T x) { return static_cast(std::conj(x)); });
}
Vectorized acos() const {
return map(std::acos);
}
Vectorized asin() const {
return map(std::asin);
}
Vectorized atan() const {
return map(std::atan);
}
Vectorized atan2(const Vectorized &exp) const {
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = std::atan2(values[i], exp[i]);
}
return ret;
}
template <
typename U = T,
typename std::enable_if_t::value, int> = 0>
Vectorized copysign(const Vectorized &sign) const {
Vectorized ret;
for (size_type i = 0; i < size(); i++) {
ret[i] = c10::copysign(values[i], sign[i]);
}
return ret;
}
Vectorized erf() const {
return map(std::erf);
}
Vectorized erfc() const {
return map(std::erfc);
}
Vectorized erfinv() const {
return map(calc_erfinv);
}
Vectorized exp() const {
return map(std::exp);
}
Vectorized expm1() const {
return map(std::expm1);
}
Vectorized frac() const {
return *this - this->trunc();
}
template <
typename U = T,
typename std::enable_if_t::value, int> = 0>
Vectorized fmod(const Vectorized& q) const {
// U is for SFINAE purposes only. Make sure it is not changed.
static_assert(std::is_same::value, "U must be T");
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = std::fmod(values[i], q[i]);
}
return ret;
}
Vectorized log() const {
return map(std::log);
}
Vectorized log10() const {
return map(std::log10);
}
Vectorized log1p() const {
return map(std::log1p);
}
template ::value, int>::type = 0>
Vectorized log2() const {
// other_t_log2 is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "other_t_log2 must be T");
return map(std::log2);
}
template ::value, int>::type = 0>
Vectorized log2() const {
// complex_t_log2 is for SFINAE and clarity. Make sure it is not changed.
static_assert(std::is_same::value, "complex_t_log2 must be T");
const T log_2 = T(std::log(2.0));
return Vectorized(map(std::log))/Vectorized(log_2);
}
Vectorized ceil() const {
return map(at::native::ceil_impl);
}
Vectorized cos() const {
return map(std::cos);
}
Vectorized cosh() const {
return map(std::cosh);
}
Vectorized floor() const {
return map(at::native::floor_impl);
}
Vectorized hypot(const Vectorized &b) const {
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = std::hypot(values[i], b[i]);
}
return ret;
}
Vectorized i0() const {
return map(calc_i0);
}
Vectorized i0e() const {
return map(calc_i0e);
}
Vectorized igamma(const Vectorized &x) const {
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = calc_igamma(values[i], x[i]);
}
return ret;
}
Vectorized igammac(const Vectorized &x) const {
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = calc_igammac(values[i], x[i]);
}
return ret;
}
Vectorized neg() const {
// NB: the trailing return type is needed because we need to coerce the
// return value back to T in the case of unary operator- incuring a
// promotion
return map([](T x) -> T { return -x; });
}
Vectorized nextafter(const Vectorized &b) const {
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = std::nextafter(values[i], b[i]);
}
return ret;
}
Vectorized round() const {
// We do not use std::round because we would like to round midway numbers to the nearest even integer.
return map(at::native::round_impl);
}
Vectorized sin() const {
return map(std::sin);
}
Vectorized sinh() const {
return map(std::sinh);
}
Vectorized tan() const {
return map(std::tan);
}
Vectorized tanh() const {
return map(std::tanh);
}
Vectorized trunc() const {
return map(at::native::trunc_impl);
}
Vectorized lgamma() const {
return map(std::lgamma);
}
Vectorized sqrt() const {
return map(std::sqrt);
}
Vectorized reciprocal() const {
return map([](T x) { return (T)(1) / x; });
}
Vectorized rsqrt() const {
return map([](T x) { return (T)1 / std::sqrt(x); });
}
Vectorized pow(const Vectorized &exp) const {
Vectorized ret;
for (int64_t i = 0; i < size(); i++) {
ret[i] = std::pow(values[i], exp[i]);
}
return ret;
}
private:
template
inline Vectorized binary_pred(const Vectorized& other, Op op) const {
// All bits are set to 1 if the pred is true, otherwise 0.
Vectorized vector;
for (int64_t i = 0; i != size(); i++) {
if (op(values[i], other.values[i])) {
std::memset(static_cast(vector.values + i), 0xFF, sizeof(T));
} else {
std::memset(static_cast(vector.values + i), 0, sizeof(T));
}
}
return vector;
}
public:
Vectorized operator==(const Vectorized& other) const { return binary_pred(other, std::equal_to()); }
Vectorized operator!=(const Vectorized& other) const { return binary_pred(other, std::not_equal_to()); }
Vectorized operator>=(const Vectorized& other) const { return binary_pred(other, std::greater_equal()); }
Vectorized operator<=(const Vectorized& other) const { return binary_pred(other, std::less_equal()); }
Vectorized operator>(const Vectorized& other) const { return binary_pred(other, std::greater()); }
Vectorized operator<(const Vectorized& other) const { return binary_pred(other, std::less()); }
private:
template
inline Vectorized binary_pred_bool(const Vectorized& other, Op op) const {
// 1 if the pred is true, otherwise 0.
Vectorized vector;
for (int i = 0; i != size(); ++ i) {
vector[i] = bool(op(values[i], other.values[i]));
}
return vector;
}
public:
Vectorized eq(const Vectorized& other) const { return binary_pred_bool(other, std::equal_to()); }
Vectorized ne(const Vectorized& other) const { return binary_pred_bool(other, std::not_equal_to()); }
Vectorized gt(const Vectorized& other) const { return binary_pred_bool(other, std::greater()); }
Vectorized ge(const Vectorized& other) const { return binary_pred_bool(other, std::greater_equal()); }
Vectorized lt(const Vectorized& other) const { return binary_pred_bool(other, std::less()); }
Vectorized le(const Vectorized& other) const { return binary_pred_bool(other, std::less_equal()); }
};
template Vectorized inline operator+(const Vectorized &a, const Vectorized &b) {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = a[i] + b[i];
}
return c;
}
template Vectorized inline operator-(const Vectorized &a, const Vectorized &b) {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = a[i] - b[i];
}
return c;
}
template Vectorized inline operator*(const Vectorized &a, const Vectorized &b) {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = a[i] * b[i];
}
return c;
}
template Vectorized inline operator/(const Vectorized &a, const Vectorized &b) __ubsan_ignore_float_divide_by_zero__ {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = a[i] / b[i];
}
return c;
}
template Vectorized inline operator||(
const Vectorized &a, const Vectorized &b) {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = a[i] || b[i];
}
return c;
}
// Implements the IEEE 754 201X `maximum` operation, which propagates NaN if
// either input is a NaN.
template ::value, int>::type = 0>
Vectorized inline maximum(const Vectorized &a, const Vectorized &b) {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = (a[i] > b[i]) ? a[i] : b[i];
if (_isnan(a[i])) {
// If either input is NaN, propagate a NaN.
// NOTE: The case where b[i] was NaN is handled correctly by the naive
// ternary operator above.
c[i] = a[i];
}
}
return c;
}
template ::value, int>::type = 0>
Vectorized inline maximum(const Vectorized &a, const Vectorized &b) {
Vectorized c;
for (int i = 0; i != Vectorized::size(); i++) {
c[i] = (std::abs(a[i]) > std::abs(b[i])) ? a[i] : b[i];
if (_isnan(a[i])) {
// If either input is NaN, propagate a NaN.
// NOTE: The case where b[i] was NaN is handled correctly by the naive
// ternary operator above.
c[i] = a[i];
}
}
return c;
}
// Implements the IEEE 754 201X `minimum` operation, which propagates NaN if
// either input is a NaN.
template ::value, int>::type = 0>
Vectorized inline minimum(const Vectorized &a, const Vectorized &b) {
Vectorized