contract.cu

Go to the documentation of this file.

namespace quda {

  namespace dslash_aux {
    #include <dslash_constants.h>
    #include <dslash_textures.h>
    #include <io_spinor.h>
  }

  using namespace dslash_aux;

#ifdef GPU_CONTRACT
#ifdef READ_SPINOR_SINGLE
#undef READ_SPINOR_SINGLE
#endif

#include "contract_core.h"
#include "contract_core_plus.h"
#include "contract_core_minus.h"

#ifndef _TWIST_QUDA_CONTRACT
#error  "Contraction core undefined"
#endif

#ifndef _TWIST_QUDA_CONTRACT_PLUS
#error  "Contraction core (plus) undefined"
#endif

#ifndef _TWIST_QUDA_CONTRACT_MINUS
#error  "Contraction core (minus) undefined"
#endif

#define checkSpinor(a, b)           \
  {                 \
    if  (a.Precision() != b.Precision())        \
      errorQuda("precisions do not match: %d %d", a.Precision(), b.Precision()); \
    if  (a.Length() != b.Length())          \
      errorQuda("lengths do not match: %d %d", a.Length(), b.Length()); \
    if  (a.Stride() != b.Stride())          \
      errorQuda("strides do not match: %d %d", a.Stride(), b.Stride()); \
  }

  template <typename Float2, typename rFloat>
  class ContractCuda : public Tunable {

  private:
    DslashParam dslashParam;
    const cudaColorSpinorField x;   // Spinor to be contracted
    const cudaColorSpinorField y;   // Spinor to be contracted
    const QudaParity parity;    // Parity of the field, actual kernels act on parity spinors
    const QudaContractType contract_type; // Type of contraction, to be detailed later

    void *result;       // The output array with the result of the contraction

    const int nTSlice;      // Time-slice in case of time-dilution

    char aux[16][TuneKey::aux_n];     // For tuning purposes

    unsigned int sharedBytesPerThread() const { return 16*sizeof(rFloat); }
    unsigned int sharedBytesPerBlock(const TuneParam &param) const { return 0; }
    bool tuneGridDim() const { return false; } // Don't tune the grid dimensions.
    unsigned int minThreads() const { return x.X(0) * x.X(1) * x.X(2) * x.X(3); }

    char *saveOut, *saveOutNorm;

    void fillAux(QudaContractType contract_type, const char *contract_str) { strcpy(aux[contract_type], contract_str); }

  public:
    ContractCuda(const cudaColorSpinorField &x, const cudaColorSpinorField &y, void *result, const QudaParity parity, const QudaContractType contract_type) :
      x(x), y(y), result(result), parity(parity), contract_type(contract_type), nTSlice(-1) {
      fillAux(QUDA_CONTRACT, "type=plain");
      fillAux(QUDA_CONTRACT_PLUS, "type=plain-plus");
      fillAux(QUDA_CONTRACT_MINUS, "type=plain-minus");
      fillAux(QUDA_CONTRACT_GAMMA5, "type=gamma5");
      fillAux(QUDA_CONTRACT_GAMMA5_PLUS, "type=gamma5-plus");
      fillAux(QUDA_CONTRACT_GAMMA5_MINUS, "type=gamma5-minus");
      fillAux(QUDA_CONTRACT_TSLICE, "type=tslice");
      fillAux(QUDA_CONTRACT_TSLICE_PLUS, "type=tslice-plus");
      fillAux(QUDA_CONTRACT_TSLICE_MINUS, "type=tslice-minus");

      dslashParam.threads = x.Volume();
      dslashParam.dc = y.getDslashConstant();
      bindSpinorTex<Float2>(&x, &y);
    }

    ContractCuda(const cudaColorSpinorField &x, const cudaColorSpinorField &y, void *result, const QudaParity parity, const QudaContractType contract_type, const int tSlice) :
      x(x), y(y), result(result), parity(parity), contract_type(contract_type), nTSlice(tSlice) {
      fillAux(QUDA_CONTRACT, "type=plain");
      fillAux(QUDA_CONTRACT_PLUS, "type=plain-plus");
      fillAux(QUDA_CONTRACT_MINUS, "type=plain-minus");
      fillAux(QUDA_CONTRACT_GAMMA5, "type=gamma5");
      fillAux(QUDA_CONTRACT_GAMMA5_PLUS, "type=gamma5-plus");
      fillAux(QUDA_CONTRACT_GAMMA5_MINUS, "type=gamma5-minus");
      fillAux(QUDA_CONTRACT_TSLICE, "type=tslice");
      fillAux(QUDA_CONTRACT_TSLICE_PLUS, "type=tslice-plus");
      fillAux(QUDA_CONTRACT_TSLICE_MINUS, "type=tslice-minus");

      DslashParam dslashParam;
      dslashParam.threads = x.X(0)*x.X(1)*x.X(2);
      dslashParam.Vsh = (x.X(0)*x.X(1)*x.X(2)) / x.SiteSubset();
      dslashParam.dc = y.getDslashConstant();
    }

    virtual ~ContractCuda() { unbindSpinorTex<Float2>(&x, &y); } // if (tSlice != NULL) { cudaFreeHost(tSlice); } }

    QudaContractType ContractType() const { return contract_type; }

    TuneKey tuneKey() const
    {
      return TuneKey(x.VolString(), typeid(*this).name(), aux[contract_type]);
    }

    void apply(const cudaStream_t &stream)
    {
      TuneParam tp = tuneLaunch(*this, getTuning(), getVerbosity());
      switch (contract_type)
  {
  default:
  case QUDA_CONTRACT_GAMMA5:    // Calculates the volume contraction (x^+ g5)_\mu y_\nu and stores it in result
    contractGamma5Kernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), parity, dslashParam);
    break;
  
  case QUDA_CONTRACT_GAMMA5_PLUS: // Calculates the volume contraction (x^+ g5)_\mu y_\nu and adds it to result
    contractGamma5PlusKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), parity, dslashParam);
    break;

  case QUDA_CONTRACT_GAMMA5_MINUS:  // Calculates the volume contraction (x^+ g5)_\mu y_\nu and substracts it from result
    contractGamma5MinusKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), parity, dslashParam);
    break;

  case QUDA_CONTRACT:     // Calculates the volume contraction x^+_\mu y_\nu and stores it in result
    contractKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), parity, dslashParam);
    break;                                                  
                                                                                  
  case QUDA_CONTRACT_PLUS:    // Calculates the volume contraction x^+_\mu y_\nu and adds it to result
    contractPlusKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), parity, dslashParam);
    break;                                                  
                                                                                        
  case QUDA_CONTRACT_MINUS:   // Calculates the volume contraction x^+_\mu y_\nu and substracts it from result
    contractMinusKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), parity, dslashParam);
    break;

  case QUDA_CONTRACT_TSLICE:    // Calculates the time-slice contraction x^+_\mu y_\nu and stores it in result
    contractTsliceKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), nTSlice, parity, dslashParam);
    break;                                                  
                                                                                  
  case QUDA_CONTRACT_TSLICE_PLUS: // Calculates the time-slice contraction x^+_\mu y_\nu and adds it to result
    contractTslicePlusKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), nTSlice, parity, dslashParam);
    break;                                                  
                                                                                        
  case QUDA_CONTRACT_TSLICE_MINUS:  // Calculates the time-slice contraction x^+_\mu y_\nu and substracts it from result
    contractTsliceMinusKernel<<<tp.grid, tp.block, tp.shared_bytes>>>((rFloat*)result, (Float2*)x.V(), (Float2*)y.V(), x.Stride(), nTSlice, parity, dslashParam);
    break;
  }
    }

    void preTune()      {
      saveOut = new char[dslashParam.threads*sizeof(Float2)*32];
      cudaMemcpy(saveOut, result, dslashParam.threads*sizeof(Float2)*32, cudaMemcpyDeviceToHost);
    }

    void postTune()     {
      cudaMemcpy(result, saveOut, dslashParam.threads*sizeof(Float2)*32, cudaMemcpyHostToDevice);
      delete[] saveOut;
    }

    long long flops() const { return 120ll * x.VolumeCB(); }
    long long bytes() const { return x.Bytes() + x.NormBytes() + y.Bytes() + y.NormBytes(); }
  };
#endif

  void contractCuda(const cudaColorSpinorField &x, const cudaColorSpinorField &y, void *result, const QudaContractType contract_type, const QudaParity parity, TimeProfile &profile)
  {
#ifdef GPU_CONTRACT
    if  ((contract_type == QUDA_CONTRACT_TSLICE) || (contract_type == QUDA_CONTRACT_TSLICE_PLUS) || (contract_type == QUDA_CONTRACT_TSLICE_MINUS)) {
      errorQuda("No time-slice specified for contraction\n");
      return;
    }

    profile.TPSTART(QUDA_PROFILE_TOTAL);
    profile.TPSTART(QUDA_PROFILE_INIT);

    Tunable *contract = 0;

    if (x.Precision() == QUDA_DOUBLE_PRECISION) {
      contract = new ContractCuda<double2,double2>(x, y, result, parity, contract_type);
    } else if (x.Precision() == QUDA_SINGLE_PRECISION) {
      contract = new ContractCuda<float4,float2>(x, y, result, parity, contract_type);
    } else if (x.Precision() == QUDA_HALF_PRECISION) {
      errorQuda("Half precision not supported for gamma5 kernel yet");
    }
    profile.TPSTOP(QUDA_PROFILE_INIT);

    profile.TPSTART(QUDA_PROFILE_COMPUTE);
    contract->apply(streams[Nstream-1]);
    qudaStreamSynchronize(streams[Nstream-1]);
    profile.TPSTOP(QUDA_PROFILE_COMPUTE);

    profile.TPSTART(QUDA_PROFILE_EPILOGUE);
    checkCudaError();

    delete contract;

    profile.TPSTOP(QUDA_PROFILE_EPILOGUE);
    profile.TPSTOP(QUDA_PROFILE_TOTAL);
#else
    errorQuda("Contraction code has not been built");
#endif
  }

  void contractCuda(const cudaColorSpinorField &x, const cudaColorSpinorField &y, void *result, const QudaContractType contract_type,
        const int nTSlice, const QudaParity parity, TimeProfile &profile)
  {
#ifdef GPU_CONTRACT
    if  ((contract_type != QUDA_CONTRACT_TSLICE) || (contract_type != QUDA_CONTRACT_TSLICE_PLUS) || (contract_type != QUDA_CONTRACT_TSLICE_MINUS)) {
      errorQuda("No time-slice input allowed for volume contractions\n");
      return;
    }

    profile.TPSTART(QUDA_PROFILE_TOTAL);
    profile.TPSTART(QUDA_PROFILE_INIT);

    Tunable *contract = 0;

    if (x.Precision() == QUDA_DOUBLE_PRECISION) {
      contract = new ContractCuda<double2,double2>(x, y, result, parity, contract_type, nTSlice);
    } else if (x.Precision() == QUDA_SINGLE_PRECISION) {
      contract = new ContractCuda<float4,float2>(x, y, result, parity, contract_type, nTSlice);
    } else if (x.Precision() == QUDA_HALF_PRECISION) {
      errorQuda("Half precision not supported for gamma5 kernel yet");
    }
    profile.TPSTOP(QUDA_PROFILE_INIT);

    profile.TPSTART(QUDA_PROFILE_COMPUTE);
    contract->apply(streams[Nstream-1]);
    qudaStreamSynchronize(streams[Nstream-1]);
    profile.TPSTOP(QUDA_PROFILE_COMPUTE);

    profile.TPSTART(QUDA_PROFILE_EPILOGUE);
    checkCudaError();
    delete contract;

    profile.TPSTOP(QUDA_PROFILE_EPILOGUE);
    profile.TPSTOP(QUDA_PROFILE_TOTAL);
#else
    errorQuda("Contraction code has not been built");
#endif
  }
  
} // namespace quda