copy_quda.cu

Go to the documentation of this file.

#include <blas_quda.h>
#include <tune_quda.h>
#include <float_vector.h>
#include <register_traits.h>

// For kernels with precision conversion built in
#define checkSpinorLength(a, b)           \
  {                 \
    if (a.Length() != b.Length())         \
      errorQuda("lengths do not match: %lu %lu", a.Length(), b.Length()); \
    if (a.Stride() != b.Stride())         \
      errorQuda("strides do not match: %d %d", a.Stride(), b.Stride()); \
    if (a.GammaBasis() != b.GammaBasis())       \
      errorQuda("gamma basis does not match: %d %d", a.GammaBasis(), b.GammaBasis()); \
  }

namespace quda {

  namespace blas {
    cudaStream_t* getStream();
    
    namespace copy_ns {

#include <texture.h>

    static struct {
      const char *vol_str;
      const char *aux_str;      
    } blasStrings;

    template <typename FloatN, int N, typename Output, typename Input>
    __global__ void copyKernel(Output Y, Input X, int length) {
      unsigned int i = blockIdx.x*(blockDim.x) + threadIdx.x;
      unsigned int parity = blockIdx.y;
      unsigned int gridSize = gridDim.x*blockDim.x;

      while (i < length) {
  FloatN x[N];
  X.load(x, i, parity);
  Y.save(x, i, parity);
  i += gridSize;
      }
    }
      
      template <typename FloatN, int N, typename Output, typename Input>
      class CopyCuda : public Tunable {

      private:
  Input &X;
  Output &Y;
  const int length;
  const int nParity;

  unsigned int sharedBytesPerThread() const { return 0; }
  unsigned int sharedBytesPerBlock(const TuneParam &param) const { return 0; }

  virtual bool advanceSharedBytes(TuneParam &param) const
  {
    TuneParam next(param);
    advanceBlockDim(next); // to get next blockDim
    int nthreads = next.block.x * next.block.y * next.block.z;
    param.shared_bytes = sharedBytesPerThread()*nthreads > sharedBytesPerBlock(param) ?
      sharedBytesPerThread()*nthreads : sharedBytesPerBlock(param);
    return false;
  }

    public:
  CopyCuda(Output &Y, Input &X, int length, int nParity)
    : X(X), Y(Y), length(length/nParity), nParity(nParity) { }
      virtual ~CopyCuda() { ; }

      inline TuneKey tuneKey() const {
  return TuneKey(blasStrings.vol_str, "copyKernel", blasStrings.aux_str); 
      }

      inline void apply(const cudaStream_t &stream) {
  TuneParam tp = tuneLaunch(*this, getTuning(), getVerbosity());
  copyKernel<FloatN, N><<<tp.grid, tp.block, tp.shared_bytes, stream>>>(Y, X, length);
      }

  void preTune() { ; } // no need to save state for copy kernels
  void postTune() { ; } // no need to restore state for copy kernels

  void initTuneParam(TuneParam &param) const {
    Tunable::initTuneParam(param);
    param.grid.y = nParity;
  }

  void defaultTuneParam(TuneParam &param) const {
    Tunable::defaultTuneParam(param);
    param.grid.y = nParity;
  }

  long long flops() const { return 0; }
  long long bytes() const { 
    const int Ninternal = (sizeof(FloatN)/sizeof(((FloatN*)0)->x))*N;
    size_t bytes = (X.Precision() + Y.Precision())*Ninternal;
    if (X.Precision() == QUDA_HALF_PRECISION) bytes += sizeof(float);
    if (Y.Precision() == QUDA_HALF_PRECISION) bytes += sizeof(float);
    return bytes*length*nParity;
  }
  int tuningIter() const { return 3; }
      };

      void copy(cudaColorSpinorField &dst, const cudaColorSpinorField &src) {
  if (&src == &dst) return; // aliasing fields
  
  if (src.SiteSubset() != dst.SiteSubset())
    errorQuda("Spinor fields do not have matching subsets dst=%d src=%d\n", src.SiteSubset(), dst.SiteSubset());

  checkSpinorLength(dst, src);

  blasStrings.vol_str = src.VolString();
  char tmp[256];
  strcpy(tmp, "dst=");
  strcat(tmp, dst.AuxString());
  strcat(tmp, ",src=");
  strcat(tmp, src.AuxString());
  blasStrings.aux_str = tmp;
  
  if (dst.Nspin() != src.Nspin())
    errorQuda("Spins (%d,%d) do not match", dst.Nspin(), src.Nspin());

  // For a given dst precision, there are two non-trivial possibilities for the
  // src precision.

  blas::bytes += (unsigned long long)src.RealLength()*(src.Precision() + dst.Precision());

  int partitions = (src.IsComposite() ? src.CompositeDim() : 1) * (src.SiteSubset());

  if (dst.Precision() == src.Precision()) {
    if (src.Bytes() != dst.Bytes()) errorQuda("Precisions match, but bytes do not");
    qudaMemcpy(dst.V(), src.V(), dst.Bytes(), cudaMemcpyDeviceToDevice);
    if (dst.Precision() == QUDA_HALF_PRECISION) {
      qudaMemcpy(dst.Norm(), src.Norm(), dst.NormBytes(), cudaMemcpyDeviceToDevice);
      blas::bytes += 2*(unsigned long long)dst.RealLength()*sizeof(float);
    }
  } else if (dst.Precision() == QUDA_DOUBLE_PRECISION && src.Precision() == QUDA_SINGLE_PRECISION) {
    if (src.Nspin() == 4){
      Spinor<float4, float4, 6, 0, 0> src_tex(src);
      Spinor<float4, double2, 6, 1> dst_spinor(dst);
      CopyCuda<float4, 6, Spinor<float4, double2, 6, 1>, Spinor<float4, float4, 6, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream()); 
    } else if (src.Nspin() == 2) {
      if (src.Length() != src.RealLength() || dst.Length() != dst.RealLength())
        errorQuda("Non-zero stride not supported"); // we need to know how many colors to set "M" (requires JIT)
      Spinor<float2, float2, 1, 0, 0> src_tex(src);
      Spinor<float2, double2, 1, 1> dst_spinor(dst);
      CopyCuda<float2, 1, Spinor<float2, double2, 1, 1>, Spinor<float2, float2, 1, 0, 0> >
        copy(dst_spinor, src_tex, src.Length()/2, partitions);
      copy.apply(*blas::getStream());
    } else if (src.Nspin() == 1) {
      Spinor<float2, float2, 3, 0, 0> src_tex(src);
      Spinor<float2, double2, 3, 1> dst_spinor(dst);
      CopyCuda<float2, 3, Spinor<float2, double2, 3, 1>,  Spinor<float2, float2, 3, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream()); 
    } else {
      errorQuda("Nspin(%d) is not supported", src.Nspin());
    }
  } else if (dst.Precision() == QUDA_SINGLE_PRECISION && src.Precision() == QUDA_DOUBLE_PRECISION) {
    if (src.Nspin() == 4){
      Spinor<float4, double2, 6, 0, 0> src_tex(src);
      Spinor<float4, float4, 6, 1> dst_spinor(dst);
      CopyCuda<float4, 6, Spinor<float4, float4, 6, 1>, Spinor<float4, double2, 6, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
        copy.apply(*blas::getStream());
    } else if (src.Nspin() == 2) {
      if (src.Length() != src.RealLength() || dst.Length() != dst.RealLength())
        errorQuda("Non-zero stride not supported"); // we need to know how many colors to set "M" (requires JIT)
      Spinor<float2, double2, 1, 0, 0> src_tex(src);
      Spinor<float2, float2, 1, 1> dst_spinor(dst);
      CopyCuda<float2, 1, Spinor<float2, float2, 1, 1>, Spinor<float2, double2, 1, 0, 0> >
        copy(dst_spinor, src_tex, src.Length()/2, partitions);
      copy.apply(*blas::getStream());
    } else if (src.Nspin() == 1) {
      Spinor<float2, double2, 3, 0, 0> src_tex(src);
      Spinor<float2, float2, 3, 1> dst_spinor(dst);
      CopyCuda<float2, 3, Spinor<float2, float2, 3, 1>, Spinor<float2, double2, 3, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream()); 
    } else {
      errorQuda("Nspin(%d) is not supported", src.Nspin());
    }
  } else if (dst.Precision() == QUDA_SINGLE_PRECISION && src.Precision() == QUDA_HALF_PRECISION) {
    blas::bytes += (unsigned long long)src.Volume()*sizeof(float);
    if (src.Nspin() == 4){      
      Spinor<float4, short4, 6, 0, 0> src_tex(src);
      Spinor<float4, float4, 6, 1> dst_spinor(dst);
      CopyCuda<float4, 6, Spinor<float4, float4, 6, 1>, Spinor<float4, short4, 6, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
        copy.apply(*blas::getStream());
    } else if (src.Nspin() == 1) {
      Spinor<float2, short2, 3, 0, 0> src_tex(src);
      Spinor<float2, float2, 3, 1> dst_spinor(dst);
      CopyCuda<float2, 3, Spinor<float2, float2, 3, 1>, Spinor<float2, short2, 3, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream());
    } else {
      errorQuda("Nspin(%d) is not supported", src.Nspin());
    }
  } else if (dst.Precision() == QUDA_HALF_PRECISION && src.Precision() == QUDA_SINGLE_PRECISION) {
    blas::bytes += (unsigned long long)dst.Volume()*sizeof(float);
    if (src.Nspin() == 4){
      Spinor<float4, float4, 6, 0, 0> src_tex(src);
      Spinor<float4, short4, 6, 1> dst_spinor(dst);
      CopyCuda<float4, 6, Spinor<float4, short4, 6, 1>,  Spinor<float4, float4, 6, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream());
    } else if (src.Nspin() == 1) {
      Spinor<float2, float2, 3, 0, 0> src_tex(src);
      Spinor<float2, short2, 3, 1> dst_spinor(dst);
      CopyCuda<float2, 3, Spinor<float2, short2, 3, 1>, Spinor<float2, float2, 3, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream()); 
    } else {
      errorQuda("Nspin(%d) is not supported", src.Nspin());
    }
  } else if (dst.Precision() == QUDA_DOUBLE_PRECISION && src.Precision() == QUDA_HALF_PRECISION) {
    blas::bytes += (unsigned long long)src.Volume()*sizeof(float);
    if (src.Nspin() == 4){
      Spinor<double2, short4, 12, 0, 0> src_tex(src);
      Spinor<double2, double2, 12, 1> dst_spinor(dst);
      CopyCuda<double2, 12, Spinor<double2, double2, 12, 1>, Spinor<double2, short4, 12, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream());
    } else if (src.Nspin() == 1) {
      Spinor<double2, short2, 3, 0, 0> src_tex(src);
      Spinor<double2, double2, 3, 1> dst_spinor(dst);
      CopyCuda<double2, 3, Spinor<double2, double2, 3, 1>, Spinor<double2, short2, 3, 0, 0> >
      copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream());
    } else {
      errorQuda("Nspin(%d) is not supported", src.Nspin());
    }
  } else if (dst.Precision() == QUDA_HALF_PRECISION && src.Precision() == QUDA_DOUBLE_PRECISION) {
    blas::bytes += (unsigned long long)dst.Volume()*sizeof(float);
    if (src.Nspin() == 4){
      Spinor<double2, double2, 12, 0, 0> src_tex(src);
      Spinor<double2, short4, 12, 1> dst_spinor(dst);
      CopyCuda<double2, 12, Spinor<double2, short4, 12, 1>, Spinor<double2, double2, 12, 0, 0> >
        copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream());
    } else if (src.Nspin() == 1) {
      Spinor<double2, double2, 3, 0, 0> src_tex(src);
      Spinor<double2, short2, 3, 1> dst_spinor(dst);
      CopyCuda<double2, 3, Spinor<double2, short2, 3, 1>, Spinor<double2, double2, 3, 0, 0> >
      copy(dst_spinor, src_tex, src.Volume(), partitions);
      copy.apply(*blas::getStream());
    } else {
      errorQuda("Nspin(%d) is not supported", src.Nspin());
    }
  } else {
    errorQuda("Invalid precision combination dst=%d and src=%d", dst.Precision(), src.Precision());
  }
      
  checkCudaError();
      }

    } // namespace copy_nw

    void copy(ColorSpinorField &dst, const ColorSpinorField &src) {
      if (dst.Location() == QUDA_CUDA_FIELD_LOCATION &&
    src.Location() == QUDA_CUDA_FIELD_LOCATION) {
  copy_ns::copy(static_cast<cudaColorSpinorField&>(dst), 
          static_cast<const cudaColorSpinorField&>(src));
      } else {
  dst = src;
      }
    }
  
  } // namespace blas
} // namespace quda