Feature/distance-preconditioning

[SaltyChiang]

We have been working on the bottom quark with the clover action recently. We cannot obtain a stable plateau in the effective mass diagram even though we have set the residual to be 1e-15. I implement the distance preconditioning algorithm proposed in https://arxiv.org/pdf/1006.4028.pdf.

Two parameters distance_pc_alpha0 and distance_pc_t0 are added to QudaInverParam, which follow Eq.(9) in the paper. I treated the preconditioning as a special solver instead of a new Dirac action, so I just modified the Wilson Dslash and didn't make a new Dirac class. The preconditioning is only enabled when calling inverQuda.

I use a point source to solve propagators with different distance_pc_alpha0 and draw the effective mass diagram generated from the pseudoscalar meson correlation functions. The effect of the preconditioning is shown below.

The effective mass plot becomes much better with alpha0=1.5.

The preconditioning is only enabled for CG inverter now due to a lack of testing. This could be enabled in more inverters later.

I created wilson_distance.cu, wilson_clover_distance.cu and wilson_clover_distance_preconditioned.cu, because the compilation time will be extremely long if I add distance and non-distance functions in one file.

[maddyscientist]

Thanks @SaltyChiang for this PR, and for making aware of this method (which I was not aware before your work 😄 ).

It's great to have this feature added to QUDA, but one issue is that there is a lot of code duplication here. A few ideas:

• I wonder if we can have a single applyWilson function, that simply takes an extra template parameter distance_pc to avoid having two different definitions of applyWilson.
• Then instead of

if (d == 3) {
  out += ratio_fwd * (U * in.project(d, proj_dir)).reconstruct(d, proj_dir);
} else {
  out += (U * in.project(d, proj_dir)).reconstruct(d, proj_dir);
}

we could use something like

out += ratio_fwd[d] * (U * in.project(d, proj_dir)).reconstruct(d, proj_dir);

where ratio_fwd would be a constexpr array that would trivially evaluate as 1.0 when distance_pc = 1 or if d < 3. I think that would allow us to have a single applyWilson definition.

• For the different .cu files: perhaps we could have a single .hpp file that contains all the boilerplate and have this file included in the stub files dslash_wilson.cu and dslash_wilson_distance.cu, for example, where the only different between the two is a booled template parameter. E.g., like we have dslash_coarse.cu and dslash_coarse_dagger.cu, etc.

Have you given any consideration to the light quark preconditioning technique outlined in the paper?

[SaltyChiang]

• I wonder if we can have a single applyWilson function, that simply takes an extra template parameter distance_pc to avoid having two different definitions of applyWilson.
• Then instead of

if (d == 3) {
  out += ratio_fwd * (U * in.project(d, proj_dir)).reconstruct(d, proj_dir);
} else {
  out += (U * in.project(d, proj_dir)).reconstruct(d, proj_dir);
}

we could use something like

out += ratio_fwd[d] * (U * in.project(d, proj_dir)).reconstruct(d, proj_dir);

where ratio_fwd would be a constexpr array that would trivially evaluate as 1.0 when distance_pc = 1 or if d < 3. I think that would allow us to have a single applyWilson definition.

@maddyscientist That's sounds better than my implementation. Thanks and I'll try to do that these days.

• For the different .cu files: perhaps we could have a single .hpp file that contains all the boilerplate and have this file included in the stub files dslash_wilson.cu and dslash_wilson_distance.cu, for example, where the only different between the two is a booled template parameter. E.g., like we have dslash_coarse.cu and dslash_coarse_dagger.cu, etc.

I'll check these files, thank you for your advice.

Have you given any consideration to the light quark preconditioning technique outlined in the paper?

Sure. The light quark preconditioning described in the paper requires a 4D $\alpha(x)$. Notice the coefficient is $\frac{1}{\cosh[\alpha_0(x_\mu-L_\mu/2)]}$, which is the inverse of that in the heavy quark case. Preconditioning in space dimensions seems to make sense only if we solve the propagator on a point source, which is not a common situation. Also, you might notice that I force the alpha0 parameter to be positive here:

quda/lib/interface_quda.cpp

Lines 2742 to 2745 in b9be584

	// Force the alpha0 to be positive.
	// A negative alpha0 matches something like Eq.(12) in arXiv:1006.4028, but the effect doesn't
	// seem to be good. Disable the negative situation as QUDA already has multigrid for light quarks.
	const double alpha0 = abs(param->distance_pc_alpha0);

In fact, leaving the alpha0 negative will trigger a t-direction light quark preconditioning. (I use the sign of alpha0 to decide whether to apply the coefficient or the inverse) I've tried this with the CG solver, but don't see a significant improvement in iteration numbers. Using multigrid for light quarks should be a much better strategy in my opinion.

[mathiaswagner]

just reviewed cmake but that looks uncontroversial.

[SaltyChiang]

@maddyscientist I move some common parts of dslash subroutine into public hpp files for Wilson and clover. I tried your advice about setting coefficients in Wilson kernel and the performance without distance preconditioning is the same as before. Do you think these modifications fit the QUDA's standard? I'll add more comments and tests for distance precondition later.

[maddyscientist]

Hi @SaltyChiang, I see your changes and this looks like a great improvement, thanks for your efforts here. I think we'll be able to get this merged with a bit more cleanup. I've left some comments, but some other outstanding things:

• Make the distance preconditioning an optional CMake parameter, so that a user can disable it if they don't need it to reduce compilation time and binary size
• Provide a sensible error message if distance preconditioning is requested for a Dirac operator type that it's not supported on
• Can you think of a good test for the distance preconditioning? Ideally we'd want to augment the invert_test suite to test this.
  • Also good to test that apply the reweighting to a quark field and removing it recovers the original field.
• Apply clang format

[SaltyChiang]

• Make the distance preconditioning an optional CMake parameter, so that a user can disable it if they don't need it to reduce compilation time and binary size

A CMake parameter called QUDA_WILSON_DISTANCE is added to enable the distance preconditioning code.

• Provide a sensible error message if distance preconditioning is requested for a Dirac operator type that it's not supported on

I show that only Wilson and clover dslash support the distance preconditioning in the error message, not sure if this is enough.

• Can you think of a good test for the distance preconditioning? Ideally we'd want to augment the invert_test suite to test this.

  • Also good to test that apply the reweighting to a quark field and removing it recovers the original field.

Two parameters --distance-pc-alpha0 and --distance-pc-t0 added to invert_test cli app. The verification code should be compatible with distance preconditioning now, as the solution vector should not differ much from the normal one. A function named verifySpinorDistanceReweight will be called if distance preconditioning is required. I'm wondering if the test implementation fits the QUDA style.

There are some other modifications:

• Applying and removing distance reweighting is performed in MatQuda, MatDagMatQuda and dslashQuda.
• invertMultiShiftQuda will raise an error if distance preconditioning parameters are set.
• Requiring distance preconditioning with single or half cuda_prec and cuda_sloppy_prec raises a warning due to the poor convergence in this situation.
• I'm not sure about the implementation of invertMultiSrcQuda, do you think it's safe to enable distance precondition here?
• We are testing more solver types.

[maddyscientist]

A CMake parameter called QUDA_WILSON_DISTANCE is added to enable the distance preconditioning code.

Maybe make this parameter QUDA_DIRAC_DISTANCE_PRECONDITIONING, since it might be extended to other Dirac operators in the future, and the more descriptive name will help explain it better.

I show that only Wilson and clover dslash support the distance preconditioning in the error message, not sure if this is enough.

That's fine, thank you.

Two parameters --distance-pc-alpha0 and --distance-pc-t0 added to invert_test cli app. The verification code should be compatible with distance preconditioning now, as the solution vector should not differ much from the normal one. A function named verifySpinorDistanceReweight will be called if distance preconditioning is required. I'm wondering if the test implementation fits the QUDA style.

For testing, I was thinking you should add some specific distance preconditioning to invert_test_gtest.hpp. That way, if the Dirac operator is being tested is the Wilson or Clover operator, and distance preconditioning is enabled, then we would test distance preconditioning with the CG solver. We have a bunch of conditional tests in invert_test_gtest at the moment, so it shouldn't be hard to add this. If we include a distance preconditioning test in invert_test_gtest, then it would be automatically run when ctest is run. If you're not sure how to do this, I can add it after we merge in your branch.

• Applying and removing distance reweighting is performed in MatQuda, MatDagMatQuda and dslashQuda.
• invertMultiShiftQuda will raise an error if distance preconditioning parameters are set.

Does distance preconditioning break the multi-shift solver? I haven't thought too hard about this, but naively the distance preconditioning shouldn't break the shifted nature of the Krylov space.

• Requiring distance preconditioning with single or half cuda_prec and cuda_sloppy_prec raises a warning due to the poor convergence in this situation.

Understood, that's fine.

• I'm not sure about the implementation of invertMultiSrcQuda, do you think it's safe to enable distance precondition here?

I think this should be fine, and just work.

[SaltyChiang]

Maybe make this parameter QUDA_DIRAC_DISTANCE_PRECONDITIONING, since it might be extended to other Dirac operators in the future, and the more descriptive name will help explain it better.

The parameter is renamed to QUDA_DIRAC_DISTANCE_PRECONDITIONING now.

Does distance preconditioning break the multi-shift solver? I haven't thought too hard about this, but naively the distance preconditioning shouldn't break the shifted nature of the Krylov space.

There shouldn't be any mathematical problem with the multi-shift solver. I considered that the multi-shift solver is usually used in RHMC algorithm, where the source is a spinor in which all time slices are filled with random values. This indicates that the solution will not have a small magnitude at t far from distance_pc_t0, where the distance preconditioning does not give any improvement in precision.

For testing, I was thinking you should add some specific distance preconditioning to invert_test_gtest.hpp. That way, if the Dirac operator is being tested is the Wilson or Clover operator, and distance preconditioning is enabled, then we would test distance preconditioning with the CG solver. We have a bunch of conditional tests in invert_test_gtest at the moment, so it shouldn't be hard to add this. If we include a distance preconditioning test in invert_test_gtest, then it would be automatically run when ctest is run. If you're not sure how to do this, I can add it after we merge in your branch.

Some tests are skipped if distance preconditioning is enabled with --distance-pc-alpha0 and --distance-pc-t0. Some tests with distance preconditioning enabled are added to tests/CMakeLists.txt. Do you think the implementation is correct?

[SaltyChiang]

Enable distance preconditioning for other solvers except for multigrid. I tried some tests and they passed:

./invert_test --dslash-type wilson --dim 4 4 4 8 --niter 1000 --ngcrkrylov 8 --matpc even-even --distance-pc-alpha0 0.1 --distance-pc-t0 1 --enable-testing true
./invert_test --dslash-type clover --compute-clover true --dim 4 4 4 8 --niter 1000 --ngcrkrylov 8 --matpc even-even --distance-pc-alpha0 0.1 --distance-pc-t0 1 --enable-testing true
./invert_test --dslash-type clover --compute-clover true --dim 4 4 4 8 --niter 1000 --ngcrkrylov 8 --matpc even-even-asym --distance-pc-alpha0 0.1 --distance-pc-t0 1 --enable-testing true

[maddyscientist]

Thanks for all the work you've done on this @SaltyChiang. The PR is looking great. Just a couple of things left to do before merging this:

• replicate the additions to QudaInvertParam in the Fortran interface
• move a function definition into spinor_reweight.cu

Once you've done this, then we can get this merged.

[maddyscientist]

Also need this CI build error to be fixed

/home/ghrunner/runners/lattice/quda/_work/quda/quda/lib/interface_quda.cpp:1788:80: error: data argument not used by format string [-Werror,-Wformat-extra-args]
      errorQuda("Multigrid solver doesn't support distance preconditioning\n", param.inv_type);

[maddyscientist]

@Jenkins ok to test