reduce Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/ReduceOpsKernel.cpp lines 428–460

template <typename scalar_t, typename acc_t = uint64_t, typename out_t = acc_t>
struct XorSumOps {
  inline C10_DEVICE acc_t reduce(acc_t acc, scalar_t data, int64_t /*idx*/) const {
    if (std::is_same<scalar_t, bool>::value) {
      return acc ^ (data ? 1 : 0);
    } else if (
        std::is_same<scalar_t, float>::value ||
        std::is_same<scalar_t, double>::value ||
        std::is_same<scalar_t, at::BFloat16>::value ||
        std::is_same<scalar_t, at::Half>::value) {
      union {
        double d;
        uint64_t u;
      } converter;
      converter.d = static_cast<double>(data);
      return acc ^ converter.u;
    } else {
      return acc ^ static_cast<uint64_t>(data);
    }
  }

  inline C10_DEVICE acc_t combine(acc_t a, acc_t b) const {
    return a ^ b;
  }

  inline C10_DEVICE out_t project(acc_t a) const {
    return a;
  }

  static C10_DEVICE acc_t translate_idx(acc_t acc, int64_t /*base_idx*/) {
    return acc;
  }
};