experimental changes

.gitignore

@@ -1,3 +1,7 @@
+# nvidia
+*.ncu-rep
+
+
 ### Custom ####
 _build
 _local

tools/pyproject.toml

@@ -26,7 +26,8 @@ dependencies = [
   'icon4py-common',
   'tabulate>=0.8.9',
   'fprettify>=0.3.7',
-  'cffi>=1.5'
+  'cffi>=1.5',
+  'cupy-cuda12x'
 ]
 description = 'Tools and utilities for integrating icon4py code into the ICON model.'
 dynamic = ['version']

tools/src/icon4pytools/py2fgen/cli.py

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
-
+import os
 import pathlib
 
 import click
@@ -23,7 +23,7 @@
 )
 from icon4pytools.py2fgen.parsing import parse
 from icon4pytools.py2fgen.plugin import generate_and_compile_cffi_plugin
-from icon4pytools.py2fgen.utils import Backend
+from icon4pytools.py2fgen.utils import GT4PyBackend
 
 
 def parse_comma_separated_list(ctx, param, value):
@@ -52,9 +52,9 @@ def parse_comma_separated_list(ctx, param, value):
     help="Enable debug mode to print additional runtime information.",
 )
 @click.option(
-    "--gt4py-backend",
+    "--backend",
     "-g",
-    type=click.Choice([e.name for e in Backend], case_sensitive=False),
+    type=click.Choice([e.name for e in GT4PyBackend], case_sensitive=False),
     default="ROUNDTRIP",
     help="Set the gt4py backend to use.",
 )
@@ -64,16 +64,18 @@ def main(
     plugin_name: str,
     build_path: pathlib.Path,
     debug_mode: bool,
-    gt4py_backend: str,
+    backend: str,
 ) -> None:
     """Generate C and F90 wrappers and C library for embedding a Python module in C and Fortran."""
-    backend = Backend[gt4py_backend]
+    if backend == "GPU":
+        os.environ["GT4PY_GPU"] = "1"
+
     build_path.mkdir(exist_ok=True, parents=True)
 
     plugin = parse(module_import_path, functions, plugin_name)
 
     c_header = generate_c_header(plugin)
-    python_wrapper = generate_python_wrapper(plugin, backend.value, debug_mode)
+    python_wrapper = generate_python_wrapper(plugin, backend, debug_mode)
     f90_interface = generate_f90_interface(plugin)
 
     generate_and_compile_cffi_plugin(plugin.plugin_name, c_header, python_wrapper, build_path)

tools/src/icon4pytools/py2fgen/generate.py

@@ -45,15 +45,13 @@ def generate_c_header(plugin: CffiPlugin) -> str:
     return codegen.format_source("cpp", generated_code, style="LLVM")
 
 
-def generate_python_wrapper(
-    plugin: CffiPlugin, gt4py_backend: Optional[str], debug_mode: bool
-) -> str:
+def generate_python_wrapper(plugin: CffiPlugin, backend: Optional[str], debug_mode: bool) -> str:
     """
     Generate Python wrapper code.
 
     Args:
         plugin: The CffiPlugin instance containing information for code generation.
-        gt4py_backend: Optional gt4py backend specification.
+        backend: Optional gt4py backend specification.
         debug_mode: Flag indicating if debug mode is enabled.
 
     Returns:
@@ -65,7 +63,7 @@ def generate_python_wrapper(
         plugin_name=plugin.plugin_name,
         function=plugin.function,
         imports=plugin.imports,
-        gt4py_backend=gt4py_backend,
+        backend=backend,
         debug_mode=debug_mode,
     )
 

tools/src/icon4pytools/py2fgen/plugin.py

@@ -14,6 +14,7 @@
 from pathlib import Path
 
 import cffi
+import cupy as cp   # type: ignore
 import numpy as np
 from cffi import FFI
 from numpy.typing import NDArray
@@ -70,6 +71,46 @@ def unpack(ptr, *sizes: int) -> NDArray:
     return arr
 
 
+def unpack_gpu(ptr, *sizes: int) -> cp.ndarray:
+    if not sizes:
+        raise ValueError("Sizes must be provided to determine the array shape.")
+
+    length = np.prod(sizes)
+    c_type = ffi.getctype(ffi.typeof(ptr).item)
+
+    dtype_map = {
+        "int": cp.int32,
+        "double": cp.float64,
+    }
+    dtype = dtype_map.get(c_type, None)
+    if dtype is None:
+        raise ValueError(f"Unsupported C data type: {c_type}")
+
+    itemsize = ffi.sizeof(c_type)
+    total_size = length * itemsize
+
+    device_id = 0
+    with cp.cuda.Device(device_id):
+        ptr_val = int(ffi.cast("uintptr_t", ptr))
+
+        mem = cp.cuda.UnownedMemory(ptr_val, total_size, owner=ptr, device_id=device_id)
+        memptr = cp.cuda.MemoryPointer(mem, 0)
+
+        # invalid memory access?
+        # cp_arr = cp.ndarray(shape=sizes, dtype=dtype, memptr=memptr, order='F')
+
+        np_arr = np.frombuffer(ffi.buffer(ptr, length * ffi.sizeof(c_type)), dtype=dtype).reshape(  # type: ignore
+            sizes, order="F"
+        )
+        print(type(np_arr))
+
+        # Convert the NumPy array to a CuPy array
+        cp_arr = cp.array(np_arr, copy=False)
+        print(type(cp_arr))
+
+    return cp_arr
+
+
 def int_array_to_bool_array(int_array: NDArray) -> NDArray:
     """
     Converts a NumPy array of integers to a boolean array.

tools/src/icon4pytools/py2fgen/template.py

@@ -11,7 +11,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 import inspect
-from typing import Any, Optional, Sequence
+from typing import Any, Sequence
 
 from gt4py.eve import Node, datamodels
 from gt4py.eve.codegen import JinjaTemplate as as_jinja, TemplatedGenerator
@@ -22,8 +22,8 @@
     BUILTIN_TO_CPP_TYPE,
     BUILTIN_TO_ISO_C_TYPE,
 )
-from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack
-from icon4pytools.py2fgen.utils import flatten_and_get_unique_elts
+from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack, unpack_gpu
+from icon4pytools.py2fgen.utils import GT4PyBackend, flatten_and_get_unique_elts
 
 
 CFFI_DECORATOR = "@ffi.def_extern()"
@@ -70,11 +70,16 @@ class CffiPlugin(Node):
 
 
 class PythonWrapper(CffiPlugin):
-    gt4py_backend: Optional[str]
+    backend: str
     debug_mode: bool
     cffi_decorator: str = CFFI_DECORATOR
     cffi_unpack: str = inspect.getsource(unpack)
+    cffi_unpack_gpu: str = inspect.getsource(unpack_gpu)
     int_to_bool: str = inspect.getsource(int_array_to_bool_array)
+    gt4py_backend: str = datamodels.field(init=False)
+
+    def __post_init__(self, *args: Any, **kwargs: Any) -> None:
+        self.gt4py_backend = GT4PyBackend[self.backend].value
 
 
 def build_array_size_args() -> dict[str, str]:
@@ -184,6 +189,7 @@ class PythonWrapperGenerator(TemplatedGenerator):
 # necessary imports for generated code to work
 from {{ plugin_name }} import ffi
 import numpy as np
+import cupy as cp
 from numpy.typing import NDArray
 from gt4py.next.ffront.fbuiltins import int32
 from gt4py.next.iterator.embedded import np_as_located_field
@@ -206,7 +212,6 @@ class PythonWrapperGenerator(TemplatedGenerator):
                     format=log_format,
                     datefmt='%Y-%m-%d %H:%M:%S')
 
-
 # We need a grid to pass offset providers
 grid = SimpleGrid()
 
@@ -216,8 +221,18 @@ class PythonWrapperGenerator(TemplatedGenerator):
 
 {{ cffi_unpack }}
 
+{{ cffi_unpack_gpu }}
+
 {{ int_to_bool }}
 
+def list_available_gpus():
+    num_gpus = cp.cuda.runtime.getDeviceCount()
+    logging.debug('Total GPUs available: %d' % num_gpus)
+
+    for i in range(num_gpus):
+        device = cp.cuda.Device(i)
+        logging.debug(device)
+
 {% for func in _this_node.function %}
 
 {{ cffi_decorator }}
@@ -234,14 +249,16 @@ def {{ func.name }}_wrapper(
         logging.info("Python Execution Context Start")
         {% endif %}
 
+        list_available_gpus()
+
         # Unpack pointers into Ndarrays
         {% for arg in func.args %}
         {% if arg.is_array %}
         {%- if _this_node.debug_mode %}
         msg = '{{ arg.name }} before unpacking: %s' % str({{ arg.name}})
         logging.debug(msg)
         {% endif %}
-        {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }})
+        {{ arg.name }} = unpack{%- if _this_node.backend == 'GPU' -%}_gpu{%- endif -%}({{ arg.name }}, {{ ", ".join(arg.size_args) }})
 
         {%- if arg.d_type.name == "BOOL" %}
         {{ arg.name }} = int_array_to_bool_array({{ arg.name }})

tools/src/icon4pytools/py2fgen/utils.py

@@ -17,7 +17,7 @@
 from gt4py.next.type_system.type_specifications import FieldType, ScalarKind, ScalarType, TypeSpec
 
 
-class Backend(Enum):
+class GT4PyBackend(Enum):
     CPU = "run_gtfn"
     GPU = "run_gtfn_gpu"
     ROUNDTRIP = "run_roundtrip"

tools/src/icon4pytools/py2fgen/wrappers/simple.py

@@ -85,3 +85,17 @@ def square_error(
     result: Field[[CellDim, KDim], float64],
 ):
     raise Exception("Exception foo occurred")
+
+
+def identity(
+    inp: Field[[CellDim, KDim], float64],
+):
+    print("inp BEFORE running gtfn gpu stencil on Python side")
+    print(inp.ndarray)
+
+    from gt4py.next.program_processors.runners.gtfn import run_gtfn_gpu
+
+    square.with_backend(run_gtfn_gpu)(inp, inp, offset_provider={})
+
+    print("inp AFTER running gtfn gpu stencil on Python side")
+    print(inp.ndarray)

tools/tests/py2fgen/fortran_samples/test_gpu.f90

@@ -0,0 +1,42 @@
+program call_identity_cffi_plugin
+   use openacc
+   use, intrinsic :: iso_c_binding
+   use identity_plugin
+   implicit none
+   integer(c_int) :: cdim, kdim, i, rc
+   real(c_double), dimension(:, :), allocatable :: input
+
+   ! array dimensions
+   cdim = 5
+   kdim = 5
+
+   ! allocate arrays (allocate in column-major order)
+   allocate(input(cdim, kdim))
+
+   ! initialise arrays
+   input = 5.0d0
+
+   ! print array shapes and values before computation
+   print *, "Fortran Arrays before calling Python:"
+   print *, "input = ", input
+   print *
+
+   !$acc enter data create(input)
+   ! call cffi code
+   call identity(input, rc)
+
+   if (rc /= 0) then
+       print *, "Python failed with exit code = ", rc
+       call exit(1)
+   end if
+
+   !$acc update host(input)
+
+   print *, "input after calling Python = ", input
+   print *
+
+   !$acc exit data delete(input)
+
+   deallocate (input)
+
+end program call_identity_cffi_plugin

tools/tests/py2fgen/test_cli.py

@@ -42,15 +42,18 @@ def run_test_case(
     backend: str,
     samples_path: Path,
     fortran_driver: str,
+    compiler: str = "gfortran",
     extra_compiler_flags: tuple[str, ...] = (),
     expected_error_code: int = 0,
 ):
     with cli.isolated_filesystem():
-        result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend, "-d"])
+        result = cli.invoke(main, [module, function, plugin_name, "--backend", backend, "-d"])
         assert result.exit_code == 0, "CLI execution failed"
 
         try:
-            compile_fortran_code(plugin_name, samples_path, fortran_driver, extra_compiler_flags)
+            compile_fortran_code(
+                plugin_name, samples_path, fortran_driver, compiler, extra_compiler_flags
+            )
         except subprocess.CalledProcessError as e:
             pytest.fail(f"Compilation failed: {e}")
 
@@ -65,11 +68,15 @@ def run_test_case(
 
 
 def compile_fortran_code(
-    plugin_name: str, samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...]
+    plugin_name: str,
+    samples_path: Path,
+    fortran_driver: str,
+    compiler: str,
+    extra_compiler_flags: tuple[str, ...],
 ):
     subprocess.run(["gfortran", "-c", f"{plugin_name}.f90", "."], check=True)
     command = [
-        "gfortran",
+        f"{compiler}",
         "-cpp",
         "-I.",
         "-Wl,-rpath=.",
@@ -104,7 +111,7 @@ def run_fortran_executable(plugin_name: str):
         ("ROUNDTRIP", ""),
     ],
 )
-def test_py2fgen_compilation_and_execution_square(
+def test_py2fgen_compilation_and_execution_square_cpu(
     cli_runner, backend, samples_path, wrapper_module, extra_flags
 ):
     run_test_case(
@@ -119,6 +126,26 @@ def test_py2fgen_compilation_and_execution_square(
     )
 
 
+@pytest.mark.parametrize(
+    "backend, extra_flags",
+    [("GPU", ("-acc",))],
+)
+def test_py2fgen_compilation_and_execution_square_gpu(
+    cli_runner, backend, samples_path, wrapper_module, extra_flags
+):
+    run_test_case(
+        cli_runner,
+        wrapper_module,
+        "identity",
+        "identity_plugin",
+        backend,
+        samples_path,
+        "test_gpu",
+        "/opt/nvidia/hpc_sdk/Linux_x86_64/2024/compilers/bin/nvfortran",
+        extra_flags,
+    )
+
+
 def test_py2fgen_python_error_propagation_to_fortran(cli_runner, samples_path, wrapper_module):
     """Tests that Exceptions triggered in Python propagate an error code (1) up to Fortran."""
     run_test_case(