feat[next-dace]: Support external memory for transient arrays

src/gt4py/next/program_processors/runners/dace/transformations/__init__.py

@@ -17,6 +17,7 @@
     GT4PyAutoOptHook,
     GT4PyAutoOptHookFun,
     GT4PyAutoOptHookStage,
+    TransientMemoryMode,
     gt_auto_optimize,
 )
 from .concat_where_mapper import (
@@ -84,7 +85,7 @@
     gt_propagate_strides_from_access_node,
     gt_propagate_strides_of,
 )
-from .utils import gt_make_transients_persistent
+from .utils import gt_configure_transient_lifetime
 
 
 __all__ = [
@@ -119,20 +120,21 @@
     "SingleStateGlobalSelfCopyElimination",
     "SplitAccessNode",
     "SplitConsumerMemlet",
+    "TransientMemoryMode",
     "VerticalMapFusionCallback",
     "VerticalMapSplitCallback",
     "constants",
     "gt_apply_concat_where_replacement_on_sdfg",
     "gt_auto_optimize",
     "gt_change_strides",
     "gt_check_if_concat_where_node_is_replaceable",
+    "gt_configure_transient_lifetime",
     "gt_create_local_double_buffering",
     "gt_eliminate_dead_dataflow",
     "gt_gpu_transform_non_standard_memlet",
     "gt_gpu_transformation",
     "gt_horizontal_map_split_fusion",
     "gt_inline_nested_sdfg",
-    "gt_make_transients_persistent",
     "gt_map_strides_to_dst_nested_sdfg",
     "gt_map_strides_to_src_nested_sdfg",
     "gt_multi_state_global_self_copy_elimination",

src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py

@@ -111,11 +111,42 @@ class GT4PyAutoOptHook(enum.Enum):
 ]
 
 
+class TransientMemoryMode(str, enum.Enum):
+    """
+    Policy selecting the lifetime/allocation strategy of transient arrays.
+
+    Supported strategies are:
+    - `SCOPED`: Transients are allocated and deallocated in the scope of the SDFG
+        where they are defined, being it the top-level SDFG or a nested one.
+        This is the default strategy.
+    - `PERSISTENT`: Transients are allocated the first time the SDFG is called and
+        retained through the entire life of the compiled SDFG, and freed only once
+        it goes out of scope.
+    - `POOL`: Transients are allocated in a memory pool, associated to the GPU
+        default stream. These allocations are managed by an asynchronous allocator,
+        since all memory is allocated and freed in stream order.
+    - `EXTERNAL`: Transients are allocated and deallocated by an external allocator.
+        This strategy allows to reuse a workspace memory for multiple SDFGs, relying
+        on sequential execution of the programs on the default stream.
+    Note:
+        The `EXTERNAL` strategy requires that the `external_memory_allocator`
+        attribute of the dace backend workflow is set to a callable that takes
+        `(required_nbytes, device_type)` and returns the allocated memory, in the
+        form of an array object that can handled by `dace.dtypes.array_interface_ptr()`.
+        The callable is expected to raise an exception if the allocation fails.
+    """
+
+    SCOPED = "SCOPED"
+    PERSISTENT = "PERSISTENT"
+    POOL = "POOL"
+    EXTERNAL = "EXTERNAL"
+
+
 def gt_auto_optimize(
     sdfg: dace.SDFG,
     gpu: bool,
     unit_strides_kind: Optional[gtx_common.DimensionKind] = None,
-    make_persistent: bool = False,
+    transient_memory_mode: TransientMemoryMode = TransientMemoryMode.POOL,
     gpu_block_size: Optional[Sequence[int | str] | str] = (32, 8, 1),
     gpu_block_size_1d: Optional[Sequence[int | str] | str] = (64, 1, 1),
     gpu_block_size_2d: Optional[Sequence[int | str] | str] = None,
@@ -130,7 +161,6 @@ def gt_auto_optimize(
     reuse_transients: bool = False,
     gpu_launch_bounds: Optional[int | str] = None,
     gpu_launch_factor: Optional[int] = None,
-    gpu_memory_pool: bool = True,
     constant_symbols: Optional[dict[str, Any]] = None,
     assume_pointwise: bool = True,
     optimization_hooks: Optional[dict[GT4PyAutoOptHook, GT4PyAutoOptHookFun]] = None,
@@ -174,9 +204,7 @@ def gt_auto_optimize(
         gpu: Optimize for GPU or CPU.
         unit_strides_kind: All dimensions of this kind are considered to have unit
             strides, see `gt_set_iteration_order()` for more.
-        make_persistent: Turn all transients to persistent lifetime, thus they are
-            allocated over the whole lifetime of the program, even if the kernel exits.
-            Thus the SDFG can not be called by different threads.
+        transient_memory_mode: Lifetime for transient arrays.
         gpu_block_size: This is used as default thread block size for GPU Maps. See
             also the `gpu_block_size_*d` arguments
         gpu_block_size_{1, 2, 3}d: Allows to specify the GPU thread block size for
@@ -194,7 +222,6 @@ def gt_auto_optimize(
         gpu_launch_bounds: Use this value as `__launch_bounds__` for _all_ GPU Maps.
         gpu_launch_factor: Use the number of threads times this value as `__launch_bounds__`
             for _all_ GPU Maps.
-        gpu_memory_pool: Enable CUDA memory pool in gpu codegen.
         constant_symbols: Symbols listed in this `dict` will be replaced by the
             respective value inside the SDFG. This might increase performance.
         assume_pointwise: Assume that the SDFG has no risk for race condition in
@@ -400,13 +427,12 @@ def gt_auto_optimize(
         sdfg = _gt_auto_post_processing(
             sdfg=sdfg,
             gpu=gpu,
-            make_persistent=make_persistent,
+            transient_memory_mode=transient_memory_mode,
             # TODO(phimuell): In general `TransientReuse` is a good idea, but the
             #   current implementation also reuses transients scalars inside Map
             #   scopes, which I do not like. Thus we should fix the transformation
             #   to avoid that.
             reuse_transients=reuse_transients,
-            gpu_memory_pool=gpu_memory_pool,
             validate_all=validate_all,
         )
 
@@ -918,9 +944,8 @@ def _gt_auto_configure_maps_and_strides(
 def _gt_auto_post_processing(
     sdfg: dace.SDFG,
     gpu: bool,
-    make_persistent: bool,
+    transient_memory_mode: TransientMemoryMode,
     reuse_transients: bool,
-    gpu_memory_pool: bool,
     validate_all: bool,
 ) -> dace.SDFG:
     """Perform post processing on the SDFG.
@@ -939,27 +964,34 @@ def _gt_auto_post_processing(
     # TODO(phimuell): Fix the bug, it uses the tile value and not the stack array value.
     dace_aoptimize.move_small_arrays_to_stack(sdfg)
 
-    if make_persistent and gpu_memory_pool:
-        raise ValueError("Cannot set both 'make_persistent' and 'gpu_memory_pool'.")
-
-    if make_persistent:
-        device = dace.DeviceType.GPU if gpu else dace.DeviceType.CPU
-        gtx_transformations.gt_make_transients_persistent(sdfg=sdfg, device=device)
-
-        if device == dace.DeviceType.GPU:
-            # NOTE: For unknown reasons the counterpart of the
-            #   `gt_make_transients_persistent()` function in DaCe, resets the
-            #   `wcr_nonatomic` property of every memlet, i.e. makes it atomic.
-            #   However, it does this only for edges on the top level and on GPU.
-            #   For compatibility with DaCe (and until we found out why) the GT4Py
-            #   auto optimizer will emulate this behaviour.
-            for state in sdfg.states():
-                assert isinstance(state, dace.SDFGState)
-                for edge in state.edges():
-                    edge.data.wcr_nonatomic = False
-
-    if gpu and gpu_memory_pool:
-        gtx_transformations.gpu_utils.gt_gpu_apply_mempool(sdfg)
+    match transient_memory_mode:
+        case TransientMemoryMode.PERSISTENT:
+            gtx_transformations.gt_configure_transient_lifetime(
+                sdfg=sdfg, lifetime=dace.AllocationLifetime.Persistent
+            )
+            if gpu:
+                # NOTE: For unknown reasons the counterpart of the
+                #   `gt_make_transients_persistent()` function in DaCe, resets the
+                #   `wcr_nonatomic` property of every memlet, i.e. makes it atomic.
+                #   However, it does this only for edges on the top level and on GPU.
+                #   For compatibility with DaCe (and until we found out why) the GT4Py
+                #   auto optimizer will emulate this behaviour.
+                for state in sdfg.states():
+                    assert isinstance(state, dace.SDFGState)
+                    for edge in state.edges():
+                        edge.data.wcr_nonatomic = False
+
+        case TransientMemoryMode.EXTERNAL:
+            gtx_transformations.gt_configure_transient_lifetime(
+                sdfg=sdfg, lifetime=dace.AllocationLifetime.External
+            )
+
+        case TransientMemoryMode.POOL:
+            if gpu:
+                gtx_transformations.gpu_utils.gt_gpu_apply_mempool(sdfg)
+
+        case TransientMemoryMode.SCOPED:
+            pass
 
     if validate_all:
         sdfg.validate()

src/gt4py/next/program_processors/runners/dace/transformations/utils.py

@@ -10,48 +10,42 @@
 
 from __future__ import annotations
 
-from typing import Optional, Sequence, TypeVar, Union
+from typing import Optional, Sequence, Union
 
 import dace
 from dace import data as dace_data, subsets as dace_sbs, symbolic as dace_sym
 from dace.libraries import standard as dace_stdlib
 from dace.sdfg import graph as dace_graph, nodes as dace_nodes
-from dace.transformation import pass_pipeline as dace_ppl
 from dace.transformation.passes import analysis as dace_analysis
 from ordered_set import OrderedSet
 
 from gt4py.next.program_processors.runners.dace import library_nodes as gtx_lib
 
 
-_PassT = TypeVar("_PassT", bound=dace_ppl.Pass)
-
-
-def gt_make_transients_persistent(
+def gt_configure_transient_lifetime(
     sdfg: dace.SDFG,
-    device: dace.DeviceType,
+    lifetime: dace.AllocationLifetime,
 ) -> dict[int, set[str]]:
     """
-    Changes the lifetime of certain transients to `Persistent`.
+    Configure transient lifetime for eligible data nodes in the given SDFG and all nested SDFGs.
 
-    A persistent lifetime means that the transient is allocated only the very first
-    time the SDFG is executed and only deallocated if the underlying `CompiledSDFG`
-    object goes out of scope. The main advantage is, that memory must not be
-    allocated every time the SDFG is run. The downside is that the SDFG can not be
-    called by different threads.
+    Eligible data nodes are transient arrays or scalars excluding:
+    - data nodes with storage type `Register` (relevant for scalars)
+    - data nodes with lifetime `External` (already externally managed)
+    - data nodes used inside a scope (e.g., maps)
+    - data nodes whose size is not fully determined by the SDFG's free symbols (dynamic allocations)
 
     Args:
         sdfg: The SDFG to process.
-        device: The device type.
+        lifetime: The desired lifetime to set for eligible transient data nodes.
 
     Returns:
-        A `dict` mapping SDFG IDs to a set of transient arrays that
-        were made persistent.
-
-    Note:
-        This function is based on a similar function in DaCe. However, the DaCe
-        function does, for unknown reasons, also reset the `wcr_nonatomic` property,
-        but only for GPU.
+        A dictionary mapping SDFG configuration IDs to the data node names whose
+        lifetimes were modified.
     """
+    if lifetime not in {dace.AllocationLifetime.Persistent, dace.AllocationLifetime.External}:
+        raise ValueError(f"Unsupported transient lifetime '{lifetime}'.")
+
     result: dict[int, set[str]] = {}
     for nsdfg in sdfg.all_sdfgs_recursive():
         fsyms: set[str] = nsdfg.free_symbols
@@ -70,6 +64,7 @@ def gt_make_transients_persistent(
 
                 desc = dnode.desc(nsdfg)
                 if not desc.transient or type(desc) not in {dace.data.Array, dace.data.Scalar}:
+                    # TODO(phimuell): Find out why scalars are processed.
                     not_modify_lifetime.add(dnode.data)
                     continue
                 if desc.storage == dace.StorageType.Register:
@@ -81,10 +76,8 @@ def gt_make_transients_persistent(
                     continue
 
                 # If the data is referenced inside a scope, such as a map, it might be possible
-                #  that it is only used inside that scope. If we would make it persistent, then
-                #  it would essentially be allocated outside and be shared among the different
-                #  map iterations. So we can not make it persistent.
-                #  The downside is, that we might have to perform dynamic allocation.
+                #  that it is only used inside that scope. If we would make it global-lifetime,
+                #  it would effectively be shared among map iterations, which is unsafe.
                 if scope_dict[dnode] is not None:
                     not_modify_lifetime.add(dnode.data)
                     continue
@@ -100,13 +93,24 @@ def gt_make_transients_persistent(
                 except AttributeError:  # total_size is an integer / has no free symbols
                     pass
 
-                # Make it persistent.
                 modify_lifetime.add(dnode.data)
 
-        # Now setting the lifetime.
         result[nsdfg.cfg_id] = modify_lifetime - not_modify_lifetime
         for aname in result[nsdfg.cfg_id]:
-            nsdfg.arrays[aname].lifetime = dace.AllocationLifetime.Persistent
+            adesc = nsdfg.arrays[aname]
+            adesc.lifetime = lifetime
+            if adesc.storage == dace.StorageType.Default and isinstance(adesc, dace.data.Array):
+                # GPU transformation have already changed the storage for GPU arrays.
+                # NOTE: If we do not change the storage during lowering / transformations,
+                #  it will be done in code generation when calling `sdfg.compile()`.
+                #  This side effect is a potential issue, because we do not store
+                #  the program handle, see `sdfg.compile(return_program_handle=False)`
+                #  in `compilation.py`; therefore, we do not have the modified SDFG.
+                #  The original SDFG is deserialized, at call time, and the storage
+                #  type is reset to `Default`. This is a problem for arrays with
+                #  external storage, because `CompiledSDFG.set_workspace()` will
+                #  try to load a symbol from the library for the wrong storage type.
+                adesc.storage = dace.StorageType.CPU_Heap
 
     return result
 

src/gt4py/next/program_processors/runners/dace/workflow/backend.py

@@ -9,13 +9,15 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Callable
 from typing import Any, Final
 
 import factory
 
 import gt4py.next.custom_layout_allocators as next_allocators
 from gt4py._core import definitions as core_defs
 from gt4py.next import backend, common, config
+from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations
 from gt4py.next.program_processors.runners.dace.workflow.factory import DaCeWorkflowFactory
 
 
@@ -35,6 +37,7 @@ class Meta:
     class Params:
         name_device = "cpu"
         name_postfix = ""
+        external_memory_allocator = None
         gpu = factory.Trait(
             allocator=next_allocators.StandardGPUFieldBufferAllocator(),
             device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA,
@@ -46,6 +49,7 @@ class Params:
             cached_translation=True,
             device_type=factory.SelfAttribute("..device_type"),
             auto_optimize=factory.SelfAttribute("..auto_optimize"),
+            external_memory_allocator=factory.SelfAttribute("..external_memory_allocator"),
         )
         auto_optimize = factory.Trait(name_postfix="_opt")
 
@@ -60,6 +64,7 @@ def make_dace_backend(
     auto_optimize: bool = True,
     async_sdfg_call: bool = True,
     optimization_args: dict[str, Any] | None = None,
+    external_memory_allocator: Callable[[int, core_defs.DeviceType], Any] | None = None,
     unstructured_horizontal_has_unit_stride: bool = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE,
     use_metrics: bool = True,
     use_zero_origin: bool = False,
@@ -74,6 +79,9 @@ def make_dace_backend(
             of GPU kernel execution with the Python driver code.
         optimization_args: A `dict` containing configuration parameters for
             the SDFG auto-optimize pipeline, see `gt_auto_optimize()`.
+        external_memory_allocator: Callable taking `(required_nbytes, storage_type)`
+            used later for external-memory workspace allocation. Threaded through
+            the backend workflow for now.
         unstructured_horizontal_has_unit_stride: When the memory layout has unit stride
             in the horizontal dimension, replace the field stride symbol with '1'.
         use_metrics: Add SDFG instrumentation to collect the metric for stencil
@@ -110,11 +118,26 @@ def make_dace_backend(
         else None
     }
 
+    if external_memory_allocator is not None:
+        expected_mode = gtx_transformations.TransientMemoryMode.EXTERNAL
+        if "transient_memory_mode" in optimization_args:
+            if (
+                transient_memory_mode := optimization_args["transient_memory_mode"]
+            ) is not expected_mode:
+                warnings.warn(
+                    f"External memory allocator provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{expected_mode}'.",
+                    stacklevel=2,
+                )
+        else:
+            optimization_args["transient_memory_mode"] = expected_mode
+
     return DaCeBackendFactory(  # type: ignore[return-value] # factory-boy typing not precise enough
         gpu=gpu,
         auto_optimize=auto_optimize,
+        external_memory_allocator=external_memory_allocator,
         otf_workflow__bare_translation__async_sdfg_call=(async_sdfg_call if gpu else False),
         otf_workflow__bare_translation__auto_optimize_args=optimization_args,
+        otf_workflow__compilation__external_memory_allocator=external_memory_allocator,
         otf_workflow__bare_translation__unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride,
         otf_workflow__bare_translation__use_metrics=use_metrics,
         otf_workflow__bare_translation__disable_field_origin_on_program_arguments=use_zero_origin,