Team-based BLAS unnecessarily redefines real & imag from ArithTraits

[mhoemmen]

@vqd8a , your recent changes to #215 defined the RealPart and ImagPart functions, for taking the real resp. imaginary part of a complex number (either Kokkos::complex<T>, or SIMD<Kokkos::complex<T>>). Kokkos::ArithTraits already has real and imag functions. For example:

kokkos-kernels/src/Kokkos_ArithTraits.hpp

Line 442 in 6e8e97a

static KOKKOS_FORCEINLINE_FUNCTION mag_type real (const T& x);

Thus, you don't need these extra RealPart and ImagPart functions. If you can instantiate Kokkos::ArithTraits for the SIMD<Kokkos::complex<T>> type, then you should be able to call the function in ArithTraits, and it should return the right thing. If it doesn't return the right thing, then we should fix ArithTraits to do the right thing.

Please remove these extra functions, and instead define real and imag for SIMD<Kokkos::complex<T>> in ArithTraits. I would be happy to review your PR.

@crtrott requested that I file this as a separate issue.

[mhoemmen]

Here is an example code showing a possible solution. (Replace std::complex with Kokkos::complex, etc.)

#include <complex>
#include <type_traits>

template<class T> class ArithTraits {};
  
template<>
struct ArithTraits<double> {
  using mag_type = double;   
  static mag_type real (const double x) { return x; }
  static mag_type imag (const double x) { return 0.0; }
};

template<>
struct ArithTraits<std::complex<double>> {
  using mag_type = double;
  static mag_type real (const std::complex<double> x) { return std::real (x); }
  static mag_type imag (const std::complex<double> x) { return std::imag (x); }
};

template<class T>
struct SIMD {
  using value_type = T;
};

template<class T>
struct AVX {
  using value_type = T;
};

template<class X, int L> struct Vector {
  using value_type = typename X::value_type;
  value_type data_[L];
};

template<template<typename> typename X, class T, int L>
struct ArithTraits< Vector<X<T>, L> > {
  using value_type = Vector<X<T>, L>;
private:
  using inner_value_type = typename X<T>::value_type;
  using inner_mag_type = typename ArithTraits<inner_value_type>::mag_type;
public:
  using mag_type = Vector<X<inner_mag_type>, L>;

  static mag_type real (const value_type z) {
      mag_type w;
      for (int k = 0; k < L; ++k) {
          w.data_[k] = ArithTraits<inner_value_type>::real (z.data_[k]);
      }
      return w;
  }

  static mag_type imag (const value_type z) {
      mag_type w;
      for (int k = 0; k < L; ++k) {
          w.data_[k] = ArithTraits<inner_value_type>::imag (z.data_[k]);
      }
      return w;
  }
};

int main( int argc, char *argv[] ) {
  using complex_simd_type = Vector<SIMD<std::complex<double>>, 4>;
  complex_simd_type v_complex;

  using real_simd_type = typename ArithTraits<complex_simd_type>::mag_type;
  real_simd_type v_real = ArithTraits<complex_simd_type>::real (v_complex);
  real_simd_type v_imag = ArithTraits<complex_simd_type>::imag (v_complex);
  
  return 0;
}

[mhoemmen]

This is just an illustration. You will have to modify it. Start with

template<template<typename> typename X, class T, int L>
struct ArithTraits< Vector<X<T>, L> > { /* ... */ };

[mhoemmen]

The class X in the above example stands for either SIMD or AVX (or any other tag type in kokkos-kernels representing a vectorization implementation option). @kyungjoo-kim , does this look like a good solution?

[mhoemmen]

@vqd8a FYI, I tested that this compiles, using godbolt.org. I highly recommend this. It lets you try out short chunks of code with different compilers and compiler flags.

[mhoemmen]

Once this compiles and passes tests, it may pay to think about specializing for vector implementations other than the generic SIMD<T> implementation. For example, you can use shuffle / swizzle instructions to vectorize extraction of real or imaginary parts.

[kyungjoo-kim]

@crtrott @mhoemmen

Initially, we do not want to put user defiend types (e.g., simd) into arith traits and it does not need to support all the trait functions either. I don't mind if we put the real, imag into arith tratis. However, once we decide to put it in the arith tratis, it should stay there for good.

[mhoemmen]

@kyungjoo-kim What's the plan for ArithTraits? Do we want to get rid of it, or leave it? I need something like ArithTraits in order to write generic solver code, and I would definitely need real and imag in it. If kokkos-kernels doesn't want to support it, that's fine.

[kyungjoo-kim]

@mhoemmen

I do not have a plan nor control for the ArithTraits. Based on a discussion with @crtrott a while ago, the arith traits is supposed to live in kokkos. That is why it is named Kokkos_ArithTraits. Assuming that the trait class lives in kokkos, kokkos does not support user defined types as far as I know. I could add a trait specialization in a separate file but at that time I just want a simple code that does not cause a trouble with the traits class.

[mhoemmen]

@kyungjoo-kim Even though ArithTraits will live in Kokkos, is it OK for downstream non-Kokkos code (say kokkos-kernels or Tpetra) to specialize ArithTraits? Tpetra currently depends on this behavior for e.g., Stokhos types.