Accelerating grid transformation using CuPy

[hanaol]

Problem
Charge grid transformation becomes a performance bottleneck for large grids.

Proposed Solution
Leverage CuPy to offload memory-intensive operations—including FFT—to the GPU.

[Andrew-S-Rosen]

Closes #131.

[Andrew-S-Rosen]

Hi @jmmshn! @hanaol is a postdoc in my group and has been working on accelerating pyrho with GPU support. The tests here won't run in CI but can be run locally.

@hanaol can speak to the speedup factors, but they are quite significant, especially for large grid sizes!

[jmmshn]

@hanaol
So based on this plot

https://raw.githubusercontent.com/materialsproject/pyrho/e767795610011850d0f8330b11613dbdc018b51c/benchmarks/results/timing_3d.png

It's always better to use the cuda version when available right?

If that is the case I think it makes sense for the code to be more insistent on that default behavior.
You might use something like this to get xp instead of cp or np.

def get_namespace(arr):
    """Get array namespace (works with numpy, cupy, etc.)"""
    if hasattr(arr, '__array_namespace__'):
        return arr.__array_namespace__()
    return np

That would also allow you to avoid using the gpu flags since the behavior can be automatically determined by the type of the input array --- you might consider a warning or something if you ever see someone use the CPU version when the GPU is available. This also helps you avoid using a flag to determine function behavior which we should usually avoid especially if it is a "mode-switch" type of flag.

[Andrew-S-Rosen]

Perhaps change to pip install mp-pyrho[gpu], which will be valid at the next PyPI release.

[Andrew-S-Rosen]

Please add a docstring for all functions where you add the use_gpu kwarg. It looks like that should be get_transformed, _transform_data, and get_sc_interp since you got the other ones already.

[Andrew-S-Rosen]

Union[NDArray, "cp.ndarray"]

[hanaol]

@hanaol So based on this plot

https://raw.githubusercontent.com/materialsproject/pyrho/e767795610011850d0f8330b11613dbdc018b51c/benchmarks/results/timing_3d.png

It's always better to use the cuda version when available right?

If that is the case I think it makes sense for the code to be more insistent on that default behavior. You might use something like this to get xp instead of cp or np.

def get_namespace(arr):
    """Get array namespace (works with numpy, cupy, etc.)"""
    if hasattr(arr, '__array_namespace__'):
        return arr.__array_namespace__()
    return np

That would also allow you to avoid using the gpu flags since the behavior can be automatically determined by the type of the input array --- you might consider a warning or something if you ever see someone use the CPU version when the GPU is available. This also helps you avoid using a flag to determine function behavior which we should usually avoid especially if it is a "mode-switch" type of flag.

Hello @jmmshn, and thanks for your comments.

I agree that using the array get_namespace is a cleaner way to avoid the use_gpu flag. My understanding is that in the current implementation of the pgrid module (see here), we convert self.grid_data from NumPy to CuPy when CuPy is available, and then convert it back to NumPy after transformation (see here). With that setup, the utils module would define a get_namespace function, and everything downstream would operate based on the namespace of the input array.

Is that what you’re referring to?

[jmmshn]

@hanaol So based on this plot
https://raw.githubusercontent.com/materialsproject/pyrho/e767795610011850d0f8330b11613dbdc018b51c/benchmarks/results/timing_3d.png
It's always better to use the cuda version when available right?
If that is the case I think it makes sense for the code to be more insistent on that default behavior. You might use something like this to get xp instead of cp or np.

def get_namespace(arr):
    """Get array namespace (works with numpy, cupy, etc.)"""
    if hasattr(arr, '__array_namespace__'):
        return arr.__array_namespace__()
    return np

That would also allow you to avoid using the gpu flags since the behavior can be automatically determined by the type of the input array --- you might consider a warning or something if you ever see someone use the CPU version when the GPU is available. This also helps you avoid using a flag to determine function behavior which we should usually avoid especially if it is a "mode-switch" type of flag.

Hello @jmmshn, and thanks for your comments.

I agree that using the array get_namespace is a cleaner way to avoid the use_gpu flag. My understanding is that in the current implementation of the pgrid module (see here), we convert self.grid_data from NumPy to CuPy when CuPy is available, and then convert it back to NumPy after transformation (see here). With that setup, the utils module would define a get_namespace function, and everything downstream would operate based on the namespace of the input array.

Is that what you’re referring to?

Hi @hanaol,
Yes, I think that should work. But I think the requirements are:

• the functions in utils.py automatically converts the inputs if they come from somewhere else where np is defined differently.
• we are 100% that it's always better with CUDA no exception.

If we are sure about it think that'll be an easy change that make the code a bit cleaner.
Thanks!

[codecov-commenter]

Codecov Report

❌ Patch coverage is 14.51613% with 53 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.54%. Comparing base (b82e926) to head (e767795).
⚠️ Report is 56 commits behind head on main.

Files with missing lines	Patch %	Lines
src/pyrho/utils.py	14.51%	53 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #132       +/-   ##
===========================================
- Coverage   79.27%   68.54%   -10.73%     
===========================================
  Files           5        5               
  Lines         304      372       +68     
===========================================
+ Hits          241      255       +14     
- Misses         63      117       +54     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hanaol]

Hi @jmmshn, I’ve updated the utils module. Please take a look.

[jmmshn]

Thanks @hanaol, I can have a look by next Friday.
If I don't reply please ping me again!