Eigen/src/Core/arch/AVX512/Reductions.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2025 Charlie Schlosser <cs.schlosser@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#ifndef EIGEN_REDUCTIONS_AVX512_H
#define EIGEN_REDUCTIONS_AVX512_H

// IWYU pragma: private
#include "../../InternalHeaderCheck.h"

namespace Eigen {

namespace internal {

// Preserve this backend's any-bit semantics by testing 32-bit chunks for every scalar width.
EIGEN_STRONG_INLINE bool avx512_predux_any(const Packet16i& bits) { return _mm512_test_epi32_mask(bits, bits) != 0; }

/* -- -- -- -- -- -- -- -- -- -- -- -- Packet16i -- -- -- -- -- -- -- -- -- -- -- -- */

template <>
EIGEN_STRONG_INLINE int predux(const Packet16i& a) {
  return _mm512_reduce_add_epi32(a);
}

template <>
EIGEN_STRONG_INLINE int predux_mul(const Packet16i& a) {
  return _mm512_reduce_mul_epi32(a);
}

template <>
EIGEN_STRONG_INLINE int predux_min(const Packet16i& a) {
  return _mm512_reduce_min_epi32(a);
}

template <>
EIGEN_STRONG_INLINE int predux_max(const Packet16i& a) {
  return _mm512_reduce_max_epi32(a);
}

template <>
EIGEN_STRONG_INLINE bool predux_any(const Packet16i& a) {
  return avx512_predux_any(a);
}

template <>
EIGEN_STRONG_INLINE bool predux_all(const Packet16i& a) {
  return _mm512_cmp_epi32_mask(a, _mm512_setzero_epi32(), _MM_CMPINT_EQ) == 0;
}

/* -- -- -- -- -- -- -- -- -- -- -- -- Packet8l -- -- -- -- -- -- -- -- -- -- -- -- */

template <>
EIGEN_STRONG_INLINE int64_t predux(const Packet8l& a) {
  return _mm512_reduce_add_epi64(a);
}

#if EIGEN_COMP_MSVC
// MSVC's _mm512_reduce_mul_epi64 is borked, at least up to and including 1939.
//    alignas(64) int64_t data[] = { 1,1,-1,-1,1,-1,-1,-1 };
//    int64_t out = _mm512_reduce_mul_epi64(_mm512_load_epi64(data));
// produces garbage: 4294967295.  This occurs when the result should be negative.
// Fall back to a manual approach:
template <>
EIGEN_STRONG_INLINE int64_t predux_mul(const Packet8l& a) {
  Packet4l lane0 = _mm512_extracti64x4_epi64(a, 0);
  Packet4l lane1 = _mm512_extracti64x4_epi64(a, 1);
  return predux_mul(pmul(lane0, lane1));
}
#else
template <>
EIGEN_STRONG_INLINE int64_t predux_mul<Packet8l>(const Packet8l& a) {
  return _mm512_reduce_mul_epi64(a);
}
#endif

template <>
EIGEN_STRONG_INLINE int64_t predux_min(const Packet8l& a) {
  return _mm512_reduce_min_epi64(a);
}

template <>
EIGEN_STRONG_INLINE int64_t predux_max(const Packet8l& a) {
  return _mm512_reduce_max_epi64(a);
}

template <>
EIGEN_STRONG_INLINE bool predux_any(const Packet8l& a) {
  return avx512_predux_any(a);
}

template <>
EIGEN_STRONG_INLINE bool predux_all(const Packet8l& a) {
  return _mm512_cmp_epi64_mask(a, _mm512_setzero_si512(), _MM_CMPINT_EQ) == 0;
}

/* -- -- -- -- -- -- -- -- -- -- -- -- Packet16f -- -- -- -- -- -- -- -- -- -- -- -- */

template <>
EIGEN_STRONG_INLINE float predux(const Packet16f& a) {
  return _mm512_reduce_add_ps(a);
}

template <>
EIGEN_STRONG_INLINE float predux_mul(const Packet16f& a) {
  return _mm512_reduce_mul_ps(a);
}

template <>
EIGEN_STRONG_INLINE float predux_min(const Packet16f& a) {
  return _mm512_reduce_min_ps(a);
}

template <>
EIGEN_STRONG_INLINE float predux_min<PropagateNumbers>(const Packet16f& a) {
  Packet8f lane0 = _mm512_extractf32x8_ps(a, 0);
  Packet8f lane1 = _mm512_extractf32x8_ps(a, 1);
  return predux_min<PropagateNumbers>(pmin<PropagateNumbers>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE float predux_min<PropagateNaN>(const Packet16f& a) {
  Packet8f lane0 = _mm512_extractf32x8_ps(a, 0);
  Packet8f lane1 = _mm512_extractf32x8_ps(a, 1);
  return predux_min<PropagateNaN>(pmin<PropagateNaN>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE float predux_max(const Packet16f& a) {
  return _mm512_reduce_max_ps(a);
}

template <>
EIGEN_STRONG_INLINE float predux_max<PropagateNumbers>(const Packet16f& a) {
  Packet8f lane0 = _mm512_extractf32x8_ps(a, 0);
  Packet8f lane1 = _mm512_extractf32x8_ps(a, 1);
  return predux_max<PropagateNumbers>(pmax<PropagateNumbers>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE float predux_max<PropagateNaN>(const Packet16f& a) {
  Packet8f lane0 = _mm512_extractf32x8_ps(a, 0);
  Packet8f lane1 = _mm512_extractf32x8_ps(a, 1);
  return predux_max<PropagateNaN>(pmax<PropagateNaN>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE bool predux_any(const Packet16f& a) {
  return avx512_predux_any(_mm512_castps_si512(a));
}

template <>
EIGEN_STRONG_INLINE bool predux_all(const Packet16f& a) {
  return _mm512_cmp_ps_mask(a, _mm512_setzero_ps(), _CMP_EQ_OQ) == 0;
}

template <>
EIGEN_STRONG_INLINE Index predux_count(const Packet16f& a) {
  return Index(popcount(static_cast<unsigned int>(_mm512_cmp_ps_mask(a, _mm512_setzero_ps(), _CMP_NEQ_UQ))));
}

/* -- -- -- -- -- -- -- -- -- -- -- -- Packet8d -- -- -- -- -- -- -- -- -- -- -- -- */

template <>
EIGEN_STRONG_INLINE double predux(const Packet8d& a) {
  return _mm512_reduce_add_pd(a);
}

template <>
EIGEN_STRONG_INLINE double predux_mul(const Packet8d& a) {
  return _mm512_reduce_mul_pd(a);
}

template <>
EIGEN_STRONG_INLINE double predux_min(const Packet8d& a) {
  return _mm512_reduce_min_pd(a);
}

template <>
EIGEN_STRONG_INLINE double predux_min<PropagateNumbers>(const Packet8d& a) {
  Packet4d lane0 = _mm512_extractf64x4_pd(a, 0);
  Packet4d lane1 = _mm512_extractf64x4_pd(a, 1);
  return predux_min<PropagateNumbers>(pmin<PropagateNumbers>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE double predux_min<PropagateNaN>(const Packet8d& a) {
  Packet4d lane0 = _mm512_extractf64x4_pd(a, 0);
  Packet4d lane1 = _mm512_extractf64x4_pd(a, 1);
  return predux_min<PropagateNaN>(pmin<PropagateNaN>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE double predux_max(const Packet8d& a) {
  return _mm512_reduce_max_pd(a);
}

template <>
EIGEN_STRONG_INLINE double predux_max<PropagateNumbers>(const Packet8d& a) {
  Packet4d lane0 = _mm512_extractf64x4_pd(a, 0);
  Packet4d lane1 = _mm512_extractf64x4_pd(a, 1);
  return predux_max<PropagateNumbers>(pmax<PropagateNumbers>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE double predux_max<PropagateNaN>(const Packet8d& a) {
  Packet4d lane0 = _mm512_extractf64x4_pd(a, 0);
  Packet4d lane1 = _mm512_extractf64x4_pd(a, 1);
  return predux_max<PropagateNaN>(pmax<PropagateNaN>(lane0, lane1));
}

template <>
EIGEN_STRONG_INLINE bool predux_any(const Packet8d& a) {
  return avx512_predux_any(_mm512_castpd_si512(a));
}

template <>
EIGEN_STRONG_INLINE bool predux_all(const Packet8d& a) {
  return _mm512_cmp_pd_mask(a, _mm512_setzero_pd(), _CMP_EQ_OQ) == 0;
}

template <>
EIGEN_STRONG_INLINE Index predux_count(const Packet8d& a) {
  return Index(popcount(static_cast<unsigned int>(_mm512_cmp_pd_mask(a, _mm512_setzero_pd(), _CMP_NEQ_UQ))));
}

// Count 16-bit floating-point lanes through integer bits so fast-math cannot treat NaN masks as zero.
EIGEN_STRONG_INLINE Index predux_count_16bit(const __m256i& a) {
#if defined(EIGEN_VECTORIZE_AVX512VL) && defined(__AVX512BW__)
  const unsigned int nonzero_lanes = _cvtmask16_u32(_mm256_test_epi16_mask(a, _mm256_set1_epi16(0x7fff)));
  return Index(popcount(nonzero_lanes));
#else
  const __m256i magnitude = _mm256_and_si256(a, _mm256_set1_epi16(0x7fff));
  const __m256i zeros = _mm256_cmpeq_epi16(magnitude, _mm256_setzero_si256());
  const unsigned int zero_bytes = static_cast<unsigned int>(_mm256_movemask_epi8(zeros));
  return Index(16 - popcount(zero_bytes) / 2);
#endif
}

#ifndef EIGEN_VECTORIZE_AVX512FP16
/* -- -- -- -- -- -- -- -- -- -- -- -- Packet16h -- -- -- -- -- -- -- -- -- -- -- -- */

template <>
EIGEN_STRONG_INLINE half predux(const Packet16h& from) {
  return half(predux(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_mul(const Packet16h& from) {
  return half(predux_mul(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_min(const Packet16h& from) {
  return half(predux_min(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_min<PropagateNumbers>(const Packet16h& from) {
  return half(predux_min<PropagateNumbers>(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_min<PropagateNaN>(const Packet16h& from) {
  return half(predux_min<PropagateNaN>(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_max(const Packet16h& from) {
  return half(predux_max(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_max<PropagateNumbers>(const Packet16h& from) {
  return half(predux_max<PropagateNumbers>(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE half predux_max<PropagateNaN>(const Packet16h& from) {
  return half(predux_max<PropagateNaN>(half2float(from)));
}

template <>
EIGEN_STRONG_INLINE bool predux_any(const Packet16h& a) {
  return predux_any<Packet8i>(a.m_val);
}

template <>
EIGEN_STRONG_INLINE bool predux_all(const Packet16h& a) {
  return predux_count_16bit(a.m_val) == 16;
}

template <>
EIGEN_STRONG_INLINE Index predux_count(const Packet16h& a) {
  return predux_count_16bit(a.m_val);
}
#endif

/* -- -- -- -- -- -- -- -- -- -- -- -- Packet16bf -- -- -- -- -- -- -- -- -- -- -- -- */

template <>
EIGEN_STRONG_INLINE bfloat16 predux(const Packet16bf& from) {
  return static_cast<bfloat16>(predux<Packet16f>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_mul(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_mul<Packet16f>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_min(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_min<Packet16f>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_min<PropagateNumbers>(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_min<PropagateNumbers>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_min<PropagateNaN>(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_min<PropagateNaN>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_max(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_max(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_max<PropagateNumbers>(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_max<PropagateNumbers>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bfloat16 predux_max<PropagateNaN>(const Packet16bf& from) {
  return static_cast<bfloat16>(predux_max<PropagateNaN>(Bf16ToF32(from)));
}

template <>
EIGEN_STRONG_INLINE bool predux_any(const Packet16bf& a) {
  return predux_any<Packet8i>(a.m_val);
}

template <>
EIGEN_STRONG_INLINE bool predux_all(const Packet16bf& a) {
  return predux_count_16bit(a.m_val) == 16;
}

template <>
EIGEN_STRONG_INLINE Index predux_count(const Packet16bf& a) {
  return predux_count_16bit(a.m_val);
}

}  // end namespace internal
}  // end namespace Eigen

#endif  // EIGEN_REDUCTIONS_AVX512_H