LCAO GPU Optimization V1 reduce device pointer look up (OpenMP target map)

[anbenali]

Proposed changes

First set of optimizations allowing for a 2.2X speed up by reducing device pointer lookup via OpenMP.

The LCAO branch was about 3 to 4X slower than CPU. After investigation, It seemed like it was due to map(to:) not checking if data was already on device. So converted all calls to is_device_ptr which fixed the issue.

Tested on multiple systems, but more specifically "TOU ASN" molecule: O(3) S(1) C(8) N(5) H(17), with 106 electrons and 1185 basis functions.

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

Code Optimization

Does this introduce a breaking change?

• No

What systems has this change been tested on?

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
• No. This PR adds tests to cover any new code, or to catch a bug that is being fixed
• No. Documentation has been added (if appropriate)
  

[ye-luo]

Please undo changes to dev_ptr so this PR can focus on real optimization.

[prckent]

Can you please add some brief specifics on what the CPU/GPU comparison was? e.g. Processor, GPU, which molecule (basis, electron count), batch size?

[anbenali]

dev_ptr

I am not understanding this... What do you want me to do?

[ye-luo]

dev_ptr

I am not understanding this... What do you want me to do?

See my comment on the source code https://github.com/QMCPACK/qmcpack/pull/5342/files#r1972521897

[anbenali]

Can you please add some brief specifics on what the CPU/GPU comparison was? e.g. Processor, GPU, which molecule (basis, electron count), batch size?

Preparing file with summary.

[ye-luo]

I profiled runs with and without this change. This optimization pattern works cross platform (nvidia/intel checked)
old

auto* ptr = a.data(); // a is a dual space container.
#pragma target map(ptr)
{}

proposed

auto* dev_ptr = a.data(); // a is a dual space container.
#pragma target is_device_ptr(dev_ptr)
{}

the difference is that the old code needs a runtime table lookup to find out the dev_ptr and it requires mutex locking the table during lookup and it is a bottleneck. When threads doesn't do the look up frequently, the cost is negligible.

The reason of slow LCAO was its small kernels, high call counts nature.
https://github.com/QMCPACK/qmcpack/blob/develop/src/QMCWaveFunctions/LCAO/SoaLocalizedBasisSet.cpp#L251
The basis computation goes through atoms one by one. It should be changed eventually to group the computation by atomic species.

The proposed way does have a drawback. Device pointers are exposed in the host code. This may cased segfault if they are not managed correctly, for example de-referenced on the host by accident.

For the moment, I think we can take this PR that uses device ptr. Eventually once we change the code computing basis species by species. We can restore the code using the old code pattern.

[ye-luo]

Please only keep changes of switching to dev_ptr and run clang-format.

[anbenali]

Fantastic!!! Thanks Ye.
Will clean asap.

However, I am almost half through isolating per basiset. Will start in a new PR.

[ye-luo]

Fantastic!!! Thanks Ye. Will clean asap.

However, I am almost half through isolating per basiset. Will start in a new PR.

Please make small PRs. large ones are too painful to review.

[prckent]

Checking on the status of this, which I think should be merged ahead of #5357 . Looks like some quick tidying is all that remains? @anbenali What are your plans/preferences?

[anbenali]

Checking on the status of this, which I think should be merged ahead of #5357 . Looks like some quick tidying is all that remains? @anbenali What are your plans/preferences?

There is only one "serious" question to @ye-luo to see if we use

auto* SuperTwist_ptr = SuperTwist.data();

and transfer that to the GPU or not.

[ye-luo]

Checking on the status of this, which I think should be merged ahead of #5357 . Looks like some quick tidying is all that remains? @anbenali What are your plans/preferences?

There is only one "serious" question to @ye-luo to see if we use

auto* SuperTwist_ptr = SuperTwist.data();

and transfer that to the GPU or not.

OK. Here is what I found. In the current way using SuperTwist in the kernel, I believe the compiler does map(tofrom: this[:1]) for you. The are transfers both H2D and D2H. It works but not efficient. This can be made more efficient by keeping a resident copy on the GPU. We can deal with it later if it becomes a bottleneck.

[anbenali]

Here is the problem, this affects only pbc.. some other users might need it but obviously I won't be able to touch it then. I don't mind doing it on device right away.. was just not sure of the behavior

[ye-luo]

My consideration is to get this PR merged timely. So we can look at your other PR that depends on this one.

Here is the problem, this affects only pbc.. some other users might need it but obviously I won't be able to touch it then. I don't mind doing it on device right away.. was just not sure of the behavior

Right now, the price is being paid regardless of using PBC. If you are willing to address it right away, simply make a new PR.

[ye-luo]

Please do a review of your own code using the github web interface to prevent missing requests. There are two conversations being marked resolved without actually being addressed. All the rest LGTM.

[anbenali]

LGTM

[ye-luo]

Test this please