gauge_random.cu

Go to the documentation of this file.

#include <quda_internal.h>
#include <quda_matrix.h>
#include <tune_quda.h>
#include <gauge_field.h>
#include <gauge_field_order.h>
#include <launch_kernel.cuh>
#include <atomic.cuh>
#include <cub_helper.cuh>
#include <index_helper.cuh>
#include <random_quda.h>

namespace quda {

  template <typename Float, QudaReconstructType recon, bool group_> struct GaugeGaussArg {
    using Gauge = typename gauge_mapper<Float, recon>::type;
    using real = typename mapper<Float>::type;
    static constexpr bool group = group_;
    int threads; // number of active threads required
    int E[4]; // extended grid dimensions
    int X[4]; // true grid dimensions
    int border[4];
    Gauge U;
    RNG rngstate;
    real sigma; // where U = exp(sigma * H)

    GaugeGaussArg(const GaugeField &U, RNG &rngstate, double sigma) : U(U), rngstate(rngstate), sigma(sigma)
    {
      int R = 0;
      for (int dir = 0; dir < 4; ++dir) {
        border[dir] = U.R()[dir];
        E[dir] = U.X()[dir];
        X[dir] = U.X()[dir] - border[dir] * 2;
        R += border[dir];
      }
      threads = X[0]*X[1]*X[2]*X[3]/2;
    }
  };

  template <typename real, typename Link> __device__ __host__ Link gauss_su3(cuRNGState &localState)
  {
    Link ret;
    real rand1[4], rand2[4], phi[4], radius[4], temp1[4], temp2[4];

    for (int i = 0; i < 4; ++i) {
      rand1[i] = Random<real>(localState);
      rand2[i] = Random<real>(localState);
    }

    for (int i = 0; i < 4; ++i) {
      phi[i] = 2.0 * M_PI * rand1[i];
      radius[i] = sqrt(-log(rand2[i]));
      sincos(phi[i], &temp2[i], &temp1[i]);
      temp1[i] *= radius[i];
      temp2[i] *= radius[i];
    }

    // construct Anti-Hermitian matrix
    ret(0, 0) = complex<real>(0.0, temp1[2] + rsqrt(3.0) * temp2[3]);
    ret(1, 1) = complex<real>(0.0, -temp1[2] + rsqrt(3.0) * temp2[3]);
    ret(2, 2) = complex<real>(0.0, -2.0 * rsqrt(3.0) * temp2[3]);
    ret(0, 1) = complex<real>(temp1[0], temp1[1]);
    ret(1, 0) = complex<real>(-temp1[0], temp1[1]);
    ret(0, 2) = complex<real>(temp1[3], temp2[0]);
    ret(2, 0) = complex<real>(-temp1[3], temp2[0]);
    ret(1, 2) = complex<real>(temp2[1], temp2[2]);
    ret(2, 1) = complex<real>(-temp2[1], temp2[2]);

    return ret;
  }

  template <typename Float, typename Arg> __global__ void computeGenGauss(Arg arg)
  {
    using real = typename mapper<Float>::type;
    using Link = Matrix<complex<real>, 3>;
    int x_cb = threadIdx.x + blockIdx.x * blockDim.x;
    int parity = threadIdx.y + blockIdx.y * blockDim.y;
    if (x_cb >= arg.threads) return;

    int x[4];
    getCoords(x, x_cb, arg.X, parity);
    for (int dr = 0; dr < 4; ++dr) x[dr] += arg.border[dr]; // extended grid coordinates

    if (arg.group && arg.sigma == 0.0) {
      // if sigma = 0 then we just set the output matrix to the identity and finish
      Link I;
      setIdentity(&I);
      for (int mu = 0; mu < 4; mu++) arg.U(mu, linkIndex(x, arg.E), parity) = I;
    } else {
      for (int mu = 0; mu < 4; mu++) {
        cuRNGState localState = arg.rngstate.State()[parity * arg.threads + x_cb];

        // generate Gaussian distributed su(n) fiueld
        Link u = gauss_su3<real, Link>(localState);
        if (arg.group) {
          u = arg.sigma * u;
          expsu3<real>(u);
        }
        arg.U(mu, linkIndex(x, arg.E), parity) = u;

        arg.rngstate.State()[parity * arg.threads + x_cb] = localState;
      }
    }
  }

  template <typename Float, typename Arg> class GaugeGauss : TunableVectorY
  {
    Arg &arg;
    const GaugeField &meta;

private:
    unsigned int minThreads() const { return arg.threads; }
    bool tuneGridDim() const { return false; } // Don't tune the grid dimensions.

public:
    GaugeGauss(Arg &arg, GaugeField &meta) : TunableVectorY(2), arg(arg), meta(meta) {}
    ~GaugeGauss() {}

    void apply(const cudaStream_t &stream)
    {
      if (meta.Location() == QUDA_CUDA_FIELD_LOCATION) {
        TuneParam tp = tuneLaunch(*this, getTuning(), getVerbosity());
        computeGenGauss<Float><<<tp.grid, tp.block, tp.shared_bytes>>>(arg);
      } else {
        errorQuda("Randomize GaugeFields on CPU not supported yet\n");
      }
    }

    TuneKey tuneKey() const { return TuneKey(meta.VolString(), typeid(*this).name(), meta.AuxString()); }

    long long flops() const { return 0; }
    long long bytes() const { return meta.Bytes(); }

    void preTune() { arg.rngstate.backup(); }
    void postTune() { arg.rngstate.restore(); }
  };

  template <typename Float, QudaReconstructType recon, bool group>
  void genGauss(GaugeField &U, RNG &rngstate, double sigma)
  {
    GaugeGaussArg<Float, recon, group> arg(U, rngstate, sigma);
    GaugeGauss<Float, decltype(arg)> gaugeGauss(arg, U);
    gaugeGauss.apply(0);
  }

  template <typename Float> void gaugeGauss(GaugeField &U, RNG &rngstate, double sigma)
  {
    if (U.LinkType() == QUDA_SU3_LINKS) { // generate Gaussian distributed SU(3) field
      if (getVerbosity() >= QUDA_SUMMARIZE)
        printfQuda("Creating Gaussian distrbuted gauge field with sigma = %e\n", sigma);
      switch (U.Reconstruct()) {
      case QUDA_RECONSTRUCT_NO: genGauss<Float, QUDA_RECONSTRUCT_NO, true>(U, rngstate, sigma); break;
      case QUDA_RECONSTRUCT_13: genGauss<Float, QUDA_RECONSTRUCT_13, true>(U, rngstate, sigma); break;
      case QUDA_RECONSTRUCT_12: genGauss<Float, QUDA_RECONSTRUCT_12, true>(U, rngstate, sigma); break;
      case QUDA_RECONSTRUCT_9: genGauss<Float, QUDA_RECONSTRUCT_9, true>(U, rngstate, sigma); break;
      case QUDA_RECONSTRUCT_8: genGauss<Float, QUDA_RECONSTRUCT_8, true>(U, rngstate, sigma); break;
      default: errorQuda("Reconstruction type %d of gauge field not supported", U.Reconstruct());
      }
    } else if (U.LinkType() == QUDA_MOMENTUM_LINKS) { // generate Gaussian distributed su(3) field
      if (getVerbosity() >= QUDA_SUMMARIZE) printfQuda("Creating Gaussian distrbuted momentum field\n");
      switch (U.Reconstruct()) {
      case QUDA_RECONSTRUCT_NO: genGauss<Float, QUDA_RECONSTRUCT_NO, false>(U, rngstate, sigma); break;
      case QUDA_RECONSTRUCT_10: genGauss<Float, QUDA_RECONSTRUCT_10, false>(U, rngstate, sigma); break;
      default: errorQuda("Reconstruction type %d of gauge field not supported", U.Reconstruct());
      }
    }
  }

  void gaugeGauss(GaugeField &U, RNG &rngstate, double sigma)
  {
    if (!U.isNative()) errorQuda("Order %d with %d reconstruct not supported", U.Order(), U.Reconstruct());
    if (U.Ncolor() != 3) errorQuda("Nc = %d not supported", U.Ncolor());

    switch (U.Precision()) {
    case QUDA_DOUBLE_PRECISION: gaugeGauss<double>(U, rngstate, sigma); break;
    case QUDA_SINGLE_PRECISION: gaugeGauss<float>(U, rngstate, sigma); break;
    default: errorQuda("Precision %d not supported", U.Precision());
    }

    // ensure multi-gpu consistency if required
    if (U.GhostExchange() == QUDA_GHOST_EXCHANGE_EXTENDED) {
      U.exchangeExtendedGhost(U.R());
    } else if (U.GhostExchange() == QUDA_GHOST_EXCHANGE_PAD) {
      U.exchangeGhost();
    }
  }

  void gaugeGauss(GaugeField &U, unsigned long long seed, double sigma)
  {
    RNG *randstates = new RNG(U, seed);
    randstates->Init();
    quda::gaugeGauss(U, *randstates, sigma);
    randstates->Release();
    delete randstates;
  }
}