inv_mr_quda.cpp

Go to the documentation of this file.

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
 
#include <complex>
 
#include <quda_internal.h>
#include <blas_quda.h>
#include <dslash_quda.h>
#include <invert_quda.h>
#include <util_quda.h>
#include <color_spinor_field.h>
 
namespace quda {
 
  MR::MR(const DiracMatrix &mat, const DiracMatrix &matSloppy, SolverParam &param, TimeProfile &profile) :
    Solver(mat, matSloppy, matSloppy, matSloppy, param, profile),
    rp(nullptr),
    r_sloppy(nullptr),
    Arp(nullptr),
    tmpp(nullptr),
    tmp_sloppy(nullptr),
    x_sloppy(nullptr),
    init(false)
  {
    if (param.schwarz_type == QUDA_MULTIPLICATIVE_SCHWARZ && param.Nsteps % 2 == 1) {
      errorQuda("For multiplicative Schwarz, number of solver steps %d must be even", param.Nsteps);
    }
  }
 
  MR::~MR() {
    if (!param.is_preconditioner) profile.TPSTART(QUDA_PROFILE_FREE);
    if (init) {
      if (x_sloppy) delete x_sloppy;
      if (tmp_sloppy) delete tmp_sloppy;
      if (tmpp) delete tmpp;
      if (Arp) delete Arp;
      if (r_sloppy) delete r_sloppy;
      if (rp) delete rp;
    }
    if (!param.is_preconditioner) profile.TPSTOP(QUDA_PROFILE_FREE);
  }
 
  void MR::operator()(ColorSpinorField &x, ColorSpinorField &b)
  {
    if (checkPrecision(x,b) != param.precision) errorQuda("Precision mismatch %d %d", checkPrecision(x,b), param.precision);
 
    if (param.maxiter == 0 || param.Nsteps == 0) {
      if (param.use_init_guess == QUDA_USE_INIT_GUESS_NO) blas::zero(x);
      return;
    }
 
    if (!init) {
      bool mixed = param.precision != param.precision_sloppy;
 
      ColorSpinorParam csParam(x);
      csParam.create = QUDA_NULL_FIELD_CREATE;
 
      // Source needs to be preserved if we're computing the true residual
      rp = (param.use_init_guess == QUDA_USE_INIT_GUESS_YES || param.preserve_source == QUDA_PRESERVE_SOURCE_YES
            || param.Nsteps > 1 || param.compute_true_res == 1) ?
        ColorSpinorField::Create(csParam) : nullptr;
 
      tmpp = (param.use_init_guess == QUDA_USE_INIT_GUESS_YES || param.Nsteps > 1 || param.compute_true_res) ?
        ColorSpinorField::Create(csParam) : nullptr;
 
      // now allocate sloppy fields
      csParam.setPrecision(param.precision_sloppy);
 
      r_sloppy = mixed ? ColorSpinorField::Create(csParam) : nullptr;  // we need a separate sloppy residual vector
      Arp = ColorSpinorField::Create(csParam);
 
      //sloppy temporary for mat-vec
      tmp_sloppy = (!tmpp || mixed) ? ColorSpinorField::Create(csParam) : nullptr;
 
      //  iterated sloppy solution vector
      x_sloppy = ColorSpinorField::Create(csParam);
 
      init = true;
    } // init
 
    ColorSpinorField &r = rp ? *rp : b;
    ColorSpinorField &rSloppy = r_sloppy ? *r_sloppy : r;
    ColorSpinorField &Ar = *Arp;
    ColorSpinorField &tmp = tmpp ? *tmpp : b;
    ColorSpinorField &tmpSloppy = tmp_sloppy ? *tmp_sloppy : tmp;
    ColorSpinorField &xSloppy = *x_sloppy;
 
    if (!param.is_preconditioner) {
      blas::flops = 0;
      profile.TPSTART(QUDA_PROFILE_COMPUTE);
    }
 
    double b2 = blas::norm2(b);  //Save norm of b
    double r2 = 0.0; // if zero source then we will exit immediately doing no work
    if (param.use_init_guess == QUDA_USE_INIT_GUESS_YES) {
      mat(r, x, tmp);
      r2 = blas::xmyNorm(b, r);   //r = b - Ax0
    } else {
      r2 = b2;
      blas::copy(r, b);
      blas::zero(x); // needed?
    }
    blas::copy(rSloppy, r);
 
    // if invalid residual then convergence is set by iteration count only
    double stop = param.residual_type == QUDA_INVALID_RESIDUAL ? 0.0 : b2*param.tol*param.tol;
    int step = 0;
 
    if (getVerbosity() >= QUDA_VERBOSE) printfQuda("MR: Initial residual = %e\n", sqrt(r2));
 
    bool converged = false;
    while (!converged) {
 
      double scale = 1.0;
      if ((node_parity+step)%2 == 0 && param.schwarz_type == QUDA_MULTIPLICATIVE_SCHWARZ) {
        // for multiplicative Schwarz we alternate updates depending on node parity
      } else {
 
        commGlobalReductionSet(param.global_reduction); // use local reductions for DD solver
 
        blas::zero(xSloppy);  // can get rid of this for a special first update kernel
        double c2 = param.global_reduction == QUDA_BOOLEAN_TRUE ? r2 : blas::norm2(r);  // c2 holds the initial r2
        scale = c2 > 0.0 ? sqrt(c2) : 1.0;
 
        // domain-wise normalization of the initial residual to prevent underflow
        if (c2 > 0.0) {
          blas::ax(1/scale, rSloppy); // can merge this with the prior copy
          r2 = 1.0; // by definition by this is now true
        }
 
        int k = 0;
        if (getVerbosity() >= QUDA_VERBOSE) printfQuda("MR: %d cycle, %d iterations, r2 = %e\n", step, k, r2);
 
        double3 Ar3;
        while (k < param.maxiter && r2 > 0.0) {
    
          matSloppy(Ar, rSloppy, tmpSloppy);
 
          if (param.global_reduction) {
            Ar3 = blas::cDotProductNormA(Ar, rSloppy);
            Complex alpha = Complex(Ar3.x, Ar3.y) / Ar3.z;
 
            // x += omega*alpha*r, r -= omega*alpha*Ar, r2 = blas::norm2(r)
            //r2 = blas::caxpyXmazNormX(omega*alpha, r, x, Ar);
            blas::caxpyXmaz(param.omega*alpha, rSloppy, xSloppy, Ar);
 
            if (getVerbosity() >= QUDA_VERBOSE)
              printfQuda("MR: %d cycle, %d iterations, <r|A|r> = (%e, %e)\n", step, k+1, Ar3.x, Ar3.y);
          } else {
            // doing local reductions so can make it asynchronous
            commAsyncReductionSet(true);
            Ar3 = blas::cDotProductNormA(Ar, rSloppy);
 
            // omega*alpha is done in the kernel
            blas::caxpyXmazMR(param.omega, rSloppy, xSloppy, Ar);
            commAsyncReductionSet(false);
          }
          k++;
 
        }
 
        // Scale and sum to accumulator
        blas::axpy(scale,xSloppy,x);
 
        commGlobalReductionSet(true); // renable global reductions for outer solver
 
      }
      step++;
 
      // FIXME - add over/under relaxation in outer loop
      if (param.compute_true_res || param.Nsteps > 1) {
        mat(r, x, tmp);
        r2 = blas::xmyNorm(b, r);
        param.true_res = sqrt(r2 / b2);
 
        converged = (step < param.Nsteps && r2 > stop) ? false : true;
 
        // if not preserving source and finished then overide source with residual
        if (param.preserve_source == QUDA_PRESERVE_SOURCE_NO && converged) blas::copy(b, r);
        else blas::copy(rSloppy, r);
 
        if (getVerbosity() >= QUDA_SUMMARIZE) {
          printfQuda("MR: %d cycle, Converged after %d iterations, relative residual: true = %e\n",
                     step, param.maxiter, sqrt(r2));
        }
      } else {
 
        blas::ax(scale, rSloppy);
        r2 = blas::norm2(rSloppy);
 
        converged = (step < param.Nsteps) ? false : true;
 
        // if not preserving source and finished then overide source with residual
        if (param.preserve_source == QUDA_PRESERVE_SOURCE_NO && converged) blas::copy(b, rSloppy);
        else blas::copy(r, rSloppy);
 
        if (getVerbosity() >= QUDA_SUMMARIZE) {
          printfQuda("MR: %d cycle, Converged after %d iterations, relative residual: iterated = %e\n",
                     step, param.maxiter, sqrt(r2));
        }
      }
 
    }
 
    if (!param.is_preconditioner) {
      profile.TPSTOP(QUDA_PROFILE_COMPUTE);
      profile.TPSTART(QUDA_PROFILE_EPILOGUE);
      param.secs += profile.Last(QUDA_PROFILE_COMPUTE);
 
      // store flops and reset counters
      double gflops = (blas::flops + mat.flops() + matSloppy.flops())*1e-9;
 
      param.gflops += gflops;
      param.iter += param.Nsteps * param.maxiter;
      blas::flops = 0;
 
      profile.TPSTOP(QUDA_PROFILE_EPILOGUE);
    }
 
    return;
  }
 
} // namespace quda