Array4: __cuda_array_interface__ v3

[ax3l]

Start implementing the __cuda_array_interface__ for zero-copy data exchange on Nvidia CUDA GPUs.

Optional: accessing an external __cuda_array_interface__ object in non-owning manner as AMReX Array4:
https://github.com/cupy/cupy/blob/a5b24f91d4d77fa03e6a4dd2ac954ff9a04e21f4/cupy/core/core.pyx#L2478-L2514

• mfab and mfab_device need to become functions, not fixtures. Otherwise they will be cached and outlive amrex.finalize(): AMReX Initialize/Finalize as Context Manager #81 MultiFab: Fix Fixture Lifetime #84
• Particle Iter & MFIter: Python does not destruct keys by default: add a GPU stream synchronize here?

  pyamrex/src/Base/MultiFab.cpp

  Lines 72 to 75 in 78bbbc7

  	if( !mfi.isValid() )
  	{
  	first_or_done = true;
  	throw py::stop_iteration();

  
  depends on MFIter::Finalize amrex#2983 and MFIter: Make Finalize Public amrex#2985
  if and for do not create a scope in Python (they do in C++):

In [1]: import numpy as np                                                                                                           

In [2]: x = np.array([1,2,3])                                                   

In [3]: for a in x: 
   ...:     print(a)                                                                
1
2
3

# a is still alive xD
In [4]: a                                                                       
Out[4]: 3

[ax3l]

⚠️ There is an nvcc host code generation bug that we fixed with Nvidia last night. Affects CUDA Toolkit 11.4-11.8 with pybind11 (pybind/pybind11#4193)
I will ship a work-around for pybind11 (pybind/pybind11#4220) before the next CUDA release, pls use an older NVCC (e.g. 11.3) in the meantime.

Update: patch that unlocks that broken compiler range in pybind/pybind11#4220
Add this to CMake:

cmake -S . -B build -DAMREX_GPU_BACKEND=CUDA -DpyAMReX_pybind11_repo=https://github.com/ax3l/pybind11.git -DpyAMReX_pybind11_branch=fix-nvcc-11.4-11.8

[ax3l]

@RemiLehe build logic from README.md is this:

• https://github.com/AMReX-Codes/pyamrex#developers
• https://github.com/AMReX-Codes/pyamrex#build

So concretely:

# Python packages if not already installed as described
python3 -m pip install -U pip setuptools wheel
python3 -m pip install -U cmake pytest
python3 -m pip install -U -r requirements.txt

# depending on what you try
python3 -m pip install cupy-cuda11x
python3 -m pip install numba
python3 -m pip install torch

# configure once (unless changing backend or versions heavily)
cmake -S . -B build -DAMReX_GPU_BACKEND=CUDA \
    -DpyAMReX_pybind11_repo=https://github.com/ax3l/pybind11.git \
    -DpyAMReX_pybind11_branch=fix-nvcc-11.4-11.8
# rinse & repeat: builds, packages & runs pip install
cmake --build build --target pip_install -j 8

and tests:

# Run all tests
python3 -m pytest tests/

# Run tests from a single file
python3 -m pytest tests/test_array4.py

# Run a single test (useful during debugging)
python3 -m pytest tests/test_array4.py::test_array4_cupy
python3 -m pytest tests/test_multifab.py::test_mfab_ops_cuda_cupy

# Run all tests, do not capture "print" output and be verbose
python3 -m pytest -s -vvvv tests/test_array4.py

and with nsight:

nsys profile -f true -t cuda,nvtx,osrt python3 -m pytest -s -vvv tests/test_multifab.py::test_mfab_ops_cuda_cupy

GUI:

nsight-sys

[ax3l]

Found a tiny bug, will rebase after #77 was merged. - Update: done.

Found another arena bug, will rebase after #78 was merged. - Update: done.

[ax3l]

First cupy progress. Gotta learn how to do in-place updates on arrays in kernels...

[ax3l]

With the new MFIter::Finalize, I can also see the cudaStreamSynchronize calls at the end of the iteration :)

[ax3l]

Wuup, wuup. First part done.
Larger tests and particles next :)