Fast, efficient implementation of SOC by camelto2 · Pull Request #4933 · QMCPACK/qmcpack

camelto2 · 2024-02-20T01:34:13Z

Please review the developer documentation
on the wiki of this project that contains help and requirements.

Proposed changes

This PR fixes a longstanding problem with SOC. Traditionally, due to the nature of the TrialWaveFunction which doesn't really exploit the fact that we only really ever use Jastrows * Determinants (or multidets), I had to do the spin integration numerically. This was quite inefficient, since we were essentially doing 8x the work since the default number of spin quadrature points was 8.

The spin integration can be done exactly if you can exploit the fact that the wavefunction is a exp(J) * Det(spinors), which can be further simplified into the ratio jastrows * inverse transpose slater matrix * spinors. Since the spinors are phi_up * exp(is) + phi_dn * exp(-is) you can pull the spin integration inside and do those terms exactly. That removes the need for the spin quadrature all together.

This is done by utilizing the TWFFastForceWrapper which gives convenient access to the appropriate quantities. A simple test problem with 16 walkers, 1000 blocks, 1 thread on my laptop shows the expected speedup.

Traditional slow evaluation: 103.8466 s
New implementation: 12.3576 s
speedup: 8.4

Right now, this only works with legacy I think. Need to look at mw_ APIs in the future. Had to add some new features to the TWFFastDerivWrapper and the SpinorSet, but everything is unit tested.

What type(s) of changes does this code introduce?

Delete the items that do not apply

New feature

Does this introduce a breaking change?

No

What systems has this change been tested on?

M1 mac

Checklist

Update the following with a yes where the items apply. If you're unsure about any of them, don't hesitate to ask. This is
simply a reminder of what we are going to look for before merging your code.

Yes. This PR is up to date with current the current state of 'develop'
Yes. Code added or changed in the PR has been clang-formatted
Yes. This PR adds tests to cover any new code, or to catch a bug that is being fixed
Yes. Documentation has been added (if appropriate)

ye-luo · 2024-02-20T02:05:42Z

Is it possible to move this optimization to NLPP batched code path where NLPP ratios of dets (without Jastrow) are batched?
I mean specialize

qmcpack/src/QMCWaveFunctions/SPOSet.h

Line 205 in a82117f

virtual void evaluateDetRatios(const VirtualParticleSet& VP,

ye-luo · 2024-02-20T02:11:30Z

src/QMCHamiltonians/SOECPComponent.cpp

+    W.makeMove(iel, deltaV_[iq]);
+    psi.getRowMSpinDecomposed(W, iel, up_row, dn_row);
+    RealType jratio = psi.evaluateJastrowRatio(W, iel);
+    W.rejectMove(iel);


To me, any NLPP code paths that require accept and reject are big no no. It can always be reformulated without need of accept/reject.

Yea I agree, but I wanted to get it working first.

Actually, on second thought we aren't doing anything costly like updating the inverse or anything like that by accepting/rejecting moves. the TWFwrapper being used here just grabs the jastrow ratio at each quadrature point and grabs the SPOSet values at the quadrature points. So I'm not sure there is much benefit from moving to using VirtualParticles

camelto2 · 2024-02-20T16:34:52Z

Is it possible to move this optimization to NLPP batched code path where NLPP ratios of dets (without Jastrow) are batched? I mean specialize

qmcpack/src/QMCWaveFunctions/SPOSet.h

Line 205 in a82117f

virtual void evaluateDetRatios(const VirtualParticleSet& VP,

evaluateDetRatios is used to give you psi(q)/psi(r). But I need it in a slightly different form in order to remove the spin integration and get what I want. Will send you some notes to take a look

rcclay

LGTM. This is a very straightforward generalization of what was already done in NonLocalECPotential for the fast forces. There are probably opportunities for optimization (better use of batching, multiwalker interfaces, etc.), but those should probably be left for a future, more targeted PR.

prckent · 2024-02-20T18:10:15Z

Test this please

prckent · 2024-02-20T19:15:45Z

FYI, the more modern compilers are complaining e.g.

src/QMCHamiltonians/tests/test_SOECPotential.cpp:258:54: error: non-const lvalue reference to type 'Real' (aka 'float') cannot bind to a value of unrelated type 'Return_t' (aka 'double')
  testing::TestSOECPotential::evalFast(so_ecp, elec, value)***

ye-luo · 2024-02-20T19:52:34Z

src/QMCHamiltonians/ECPotentialBuilder.cpp

  pAttrib.add(pbc, "pbc", {"yes", "no"});
  pAttrib.add(forces, "forces", {"no", "yes"});
  pAttrib.add(physicalSO, "physicalSO", {"yes", "no"});
+  pAttrib.add(fast_so, "fastSO", {"yes", "no"});


"fastSO" is hard to understand what is fast. Could you find a more understandable name?

Switched to spin_integrator. Can choose between "exact" and "simpson" which will use the old simpsos rule to integrate

camelto2 · 2024-02-20T20:50:18Z

FYI, the more modern compilers are complaining e.g.

src/QMCHamiltonians/tests/test_SOECPotential.cpp:258:54: error: non-const lvalue reference to type 'Real' (aka 'float') cannot bind to a value of unrelated type 'Return_t' (aka 'double')
  testing::TestSOECPotential::evalFast(so_ecp, elec, value)***

This should be fixed now. There was a type mismatch in my unit test

camelto2 · 2024-02-20T22:05:19Z

Not sure I understand why linux (Clang14-NoMPI-UBSan-Real) is failing. Looks like test_parser, which is something I didn't touch

ye-luo · 2024-02-21T03:13:09Z

Test this please

ye-luo

Temporary blocking merge. Will review this PR in 2 days.

rcclay · 2024-02-27T14:59:05Z

What's the status of this?

prckent · 2024-02-27T16:20:57Z

A huge improvement and good to get in. Need to identify what is needed for batched code execution in followup work.

prckent · 2024-03-26T16:49:27Z

Just checking on the status of this. I assume the rework is being done with the limited time available. Any gotchas?

camelto2 · 2024-03-26T17:14:11Z

Just no time. I reworked it using Ye's approach on my end, and got the unit test passing for the SOECPotential. Just need to finish adapting the unit test for the changes to the SpinorSet. Will hopefully get to it today

camelto2 · 2024-04-04T16:37:17Z

Sorry for taking so long, this should be up to date with Ye's recommendations

PDoakORNL · 2024-04-04T18:38:00Z

There is a great deal of per walker data stashed in SpinorSet and the force wrapper. In fact none of this really looks like it took into account the batched design. These should really have been crowd scope memory resources i.e setup through resource sets (which will also allow use to track memory use much more easily). Evaluations should be over these whole blocks and not broken up by walker unless necessary. Certainly the end goal should be a single kernel launch to deal with the force calculations over a block of walkers.

The way this is setup right now assuming a Wrapper object per walker and a SpinsorSet object per walker that a bunch of workspace or additional state can be stashed in really isn't great. Extracting pointers from each walker and its inheritance pile is not a good design and was just a hack to avoid refactoring too much at once.

prckent · 2024-04-05T14:01:55Z

Thanks Cody. I am hopeful that this gets merged after the CI passes. The real build needs fixing and there are sanitizer failures.

Peter: I think we should see where we end up with e.g. real world memory usage and performance for observables before tackling design issues here.

ye-luo · 2024-04-05T16:24:14Z

src/QMCHamiltonians/ECPotentialBuilder.cpp

+    if (spin_integrator == "exact")
+    {
+      app_log() << "    Using fast SOECP evaluation. Spin integration is exact" << std::endl;
+      apot->setFastEvaluation(true);


Do you need fast and slow code paths both active in a single object?
If not, Can this option be passed via constructor?

will fix this next, want to make sure tests pass first.

I fixed the complex build. When you remove this setting, please make sure

qmcpack/src/QMCHamiltonians/SOECPComponent.cpp

Line 67 in 5921824

void SOECPComponent::resize_warrays(int n, int m, int s)

being reactive to the fast algorithm selection and set total_knots_ properly.

This looks to be taken care of now from what I can tell. Should fix the failing test as well

src/QMCHamiltonians/SOECPComponent.h

src/QMCHamiltonians/SOECPotential.h

ye-luo · 2024-04-05T16:26:57Z

src/QMCHamiltonians/SOECPotential.h

  ResourceHandle<SOECPotentialMultiWalkerResource> mw_res_handle_;

+  //flag to use fast evaluation
+  bool use_fast_evaluation_;


If we set its value via constructor, please add const.

src/QMCWaveFunctions/SPOSet.h

src/QMCWaveFunctions/tests/test_einset_spinor.cpp

src/QMCWaveFunctions/SpinorSet.h

src/QMCWaveFunctions/SPOSet.cpp

src/QMCHamiltonians/SOECPotential.cpp

ye-luo · 2024-04-10T03:33:19Z

Test this please

ye-luo · 2024-04-10T04:06:33Z

Test this please

ye-luo · 2024-04-11T23:42:08Z

Test this please

camelto2 added 8 commits February 15, 2024 14:19

add initial structure and failing test

14d6640

matrix setup and sizing

4b2c51f

finish SOECPotential, start component

56b11e1

add tests and implementation. Not passing yet

6c032d0

add implementation, passes unit test, and add input control

1c37096

Merge remote-tracking branch 'upstream/develop' into fast_soc

327de2a

small fix

1796101

Merge remote-tracking branch 'upstream/develop' into fast_soc

48d7c01

camelto2 requested review from rcclay and ye-luo and removed request for rcclay February 20, 2024 01:34

ye-luo reviewed Feb 20, 2024

View reviewed changes

fix failing tests

dce92d1

rcclay previously approved these changes Feb 20, 2024

View reviewed changes

ye-luo reviewed Feb 20, 2024

View reviewed changes

fix failures in mix builds

81fd1ed

camelto2 dismissed rcclay’s stale review via 81fd1ed February 20, 2024 20:32

change input name and add documentation

358c702

rcclay previously approved these changes Feb 21, 2024

View reviewed changes

ye-luo requested changes Feb 21, 2024

View reviewed changes

camelto2 added 3 commits April 4, 2024 10:30

address Ye comments

da77253

Merge remote-tracking branch 'upstream/develop' into fast_soc

6d7ca5c

Merge remote-tracking branch 'origin/fast_soc' into fast_soc

c77eab6

camelto2 requested a review from ye-luo April 4, 2024 16:37

ye-luo requested changes Apr 5, 2024

View reviewed changes

camelto2 added 5 commits April 8, 2024 11:01

Merge remote-tracking branch 'upstream/develop' into fast_soc

9bd850c

address comments

880e5dc

clang format

3c79f4c

remove header

f883078

remove header no longer needed

2747021

ye-luo added 2 commits April 9, 2024 22:42

Fix complex compilation.

e784553

Mixed precision fix.

8c2073f

camelto2 and others added 3 commits April 11, 2024 11:09

fix fast evaluation in constructor

428acd4

Merge remote-tracking branch 'origin/fast_soc' into fast_soc

f4f8d04

Add comment.

562ebab

ye-luo approved these changes Apr 11, 2024

View reviewed changes

ye-luo enabled auto-merge April 12, 2024 00:30

prckent approved these changes Apr 12, 2024

View reviewed changes

prckent disabled auto-merge April 12, 2024 02:23

prckent merged commit 1cd71b5 into QMCPACK:develop Apr 12, 2024

camelto2 deleted the fast_soc branch April 12, 2024 19:44

camelto2 mentioned this pull request Aug 9, 2024

Add mw_ version of exact spin integration for SOECPs #5119

Merged

Conversation

camelto2 commented Feb 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposed changes

What type(s) of changes does this code introduce?

Does this introduce a breaking change?

What systems has this change been tested on?

Checklist

Uh oh!

ye-luo commented Feb 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camelto2 commented Feb 20, 2024

Uh oh!

rcclay left a comment

Choose a reason for hiding this comment

Uh oh!

prckent commented Feb 20, 2024

Uh oh!

prckent commented Feb 20, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camelto2 commented Feb 20, 2024

Uh oh!

camelto2 commented Feb 20, 2024

Uh oh!

ye-luo commented Feb 21, 2024

Uh oh!

ye-luo left a comment

Choose a reason for hiding this comment

Uh oh!

rcclay commented Feb 27, 2024

Uh oh!

prckent commented Feb 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

prckent commented Mar 26, 2024

Uh oh!

camelto2 commented Mar 26, 2024

Uh oh!

camelto2 commented Apr 4, 2024

Uh oh!

PDoakORNL commented Apr 4, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

prckent commented Apr 5, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ye-luo commented Apr 10, 2024

Uh oh!

ye-luo commented Apr 10, 2024

Uh oh!

ye-luo commented Apr 11, 2024

camelto2 commented Feb 20, 2024 •

edited

Loading

ye-luo commented Feb 20, 2024 •

edited

Loading

prckent commented Feb 27, 2024 •

edited

Loading

PDoakORNL commented Apr 4, 2024 •

edited

Loading