Add dynamic temporary allocation and deallocation

doc/tutorial.rst

@@ -1078,10 +1078,12 @@ Tagged prefetching
 Temporaries in global memory
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-:mod:`loopy` supports using temporaries with global storage duration. As with
-local and private temporaries, the runtime allocates storage for global
-temporaries when the kernel gets executed. The user must explicitly specify that
-a temporary is global. To specify that a temporary is global, use
+:mod:`loopy` supports using temporaries with global storage duration. Note that
+these temporaries will in fact be dynamically allocated before the first usage,
+and deallocated after the last. As with local and private temporaries, the
+runtime allocates storage for global temporaries when the kernel gets executed.
+The user must explicitly specify that a temporary is global. To specify that a
+temporary is global, use
 :func:`loopy.set_temporary_address_space`.
 
 Substitution rules
@@ -1276,15 +1278,17 @@ put those instructions into the schedule.
    ...
    ---------------------------------------------------------------------------
    LINEARIZATION:
-      0: CALL KERNEL rotate_v2
-      1:     tmp = arr[i_inner + i_outer*16]  {id=maketmp}
-      2:     tmp_save_slot[tmp_save_hw_dim_0_rotate_v2, tmp_save_hw_dim_1_rotate_v2] = tmp  {id=tmp.save}
-      3: RETURN FROM KERNEL rotate_v2
-      4: ... gbarrier
-      5: CALL KERNEL rotate_v2_0
-      6:     tmp = tmp_save_slot[tmp_reload_hw_dim_0_rotate_v2_0, tmp_reload_hw_dim_1_rotate_v2_0]  {id=tmp.reload}
-      7:     arr[(i_inner + i_outer*16 + 1) % n] = tmp  {id=rotate}
-      8: RETURN FROM KERNEL rotate_v2_0
+      0: ALLOCATE tmp_save_slot
+      1: CALL KERNEL rotate_v2
+      2:     tmp = arr[i_inner + i_outer*16]  {id=maketmp}
+      3:     tmp_save_slot[tmp_save_hw_dim_0_rotate_v2, tmp_save_hw_dim_1_rotate_v2] = tmp  {id=tmp.save}
+      4: RETURN FROM KERNEL rotate_v2
+      5: ... gbarrier
+      6: CALL KERNEL rotate_v2_0
+      7:     tmp = tmp_save_slot[tmp_reload_hw_dim_0_rotate_v2_0, tmp_reload_hw_dim_1_rotate_v2_0]  {id=tmp.reload}
+      8:     arr[(i_inner + i_outer*16 + 1) % n] = tmp  {id=rotate}
+      9: RETURN FROM KERNEL rotate_v2_0
+     10: DEALLOCATE tmp_save_slot
    ---------------------------------------------------------------------------
 
 Here's an overview of what :func:`loopy.save_and_reload_temporaries` actually
@@ -1294,7 +1298,8 @@ does in more detail:
    variables' live ranges cross a global barrier.
 
 2. For each temporary, :mod:`loopy` creates a storage slot for the temporary in
-   global memory (see :ref:`global_temporaries`).
+   global memory (see :ref:`global_temporaries`). Note that, as a global variable, the storage
+   slot will be dynamically allocated before its first usage and deallocated after its last.
 
 3. :mod:`loopy` saves the temporary into its global storage slot whenever it
    detects the temporary is live-out from a kernel, and reloads the temporary

loopy/check.py

@@ -1373,8 +1373,10 @@ def _check_for_unused_hw_axes_in_kernel_chunk(
             sched_index: int | None = None
         ) -> int:
     from loopy.schedule import (
+        AllocTemp,
         Barrier,
         CallKernel,
+        DeallocTemp,
         EnterLoop,
         LeaveLoop,
         ReturnFromKernel,
@@ -1469,7 +1471,8 @@ def _check_for_unused_hw_axes_in_kernel_chunk(
                         "Calling loopy.add_inames_for_unused_hw_axes(...) "
                         "might help.")
 
-        elif isinstance(sched_item, (Barrier, EnterLoop, LeaveLoop)):
+        elif isinstance(sched_item, (Barrier, EnterLoop, LeaveLoop,
+                                     AllocTemp, DeallocTemp)):
             i += 1
             continue
 

loopy/codegen/control.py

@@ -31,8 +31,10 @@
 from loopy.codegen.result import CodeGenerationResult, merge_codegen_results, wrap_in_if
 from loopy.diagnostic import LoopyError
 from loopy.schedule import (
+    AllocTemp,
     Barrier,
     CallKernel,
+    DeallocTemp,
     EnterLoop,
     LeaveLoop,
     RunInstruction,
@@ -184,6 +186,12 @@
                 "instruction %s" % insn.id,
                 lambda inner_cgs: generate_instruction_code(inner_cgs, insn))
 
+    elif isinstance(sched_item, AllocTemp):
+        return codegen_state.ast_builder.emit_alloc_temp(codegen_state, sched_item.temp)
+    elif isinstance(sched_item, DeallocTemp):
+        return codegen_state.ast_builder.emit_dealloc_temp(
+            codegen_state, sched_item.temp)
+
     else:
         raise RuntimeError("unexpected schedule item type: %s"
                 % type(sched_item))
@@ -197,7 +205,7 @@
 
     result = None
     for _, sched_item in generate_sub_sched_items(kernel.linearization, sched_index):
-        if isinstance(sched_item, Barrier):
+        if isinstance(sched_item, (Barrier, AllocTemp, DeallocTemp)):
             my_preds = frozenset()
         elif isinstance(sched_item, RunInstruction):
             my_preds = kernel.id_to_insn[sched_item.insn_id].predicates
@@ -278,7 +286,7 @@
             assert i <= codegen_state.schedule_index_end, \
                     "schedule block extends beyond schedule_index_end"
 
-        elif isinstance(sched_item, (Barrier, RunInstruction)):
+        elif isinstance(sched_item, (Barrier, RunInstruction, AllocTemp, DeallocTemp)):
             i += 1
         else:
             raise RuntimeError("unexpected schedule item type: %s"

loopy/schedule/__init__.py

@@ -25,6 +25,7 @@
 
 import logging
 import sys
+from collections import defaultdict
 from dataclasses import dataclass, replace
 from typing import (
     TYPE_CHECKING,
@@ -39,13 +40,13 @@
 from pytools.persistent_dict import WriteOncePersistentDict
 
 from loopy.diagnostic import LoopyError, ScheduleDebugInputError, warn_with_kernel
+from loopy.kernel.data import AddressSpace
 from loopy.tools import LoopyKeyBuilder, caches
 from loopy.typing import not_none as not_none
 from loopy.version import DATA_MODEL_VERSION
 
 
 if TYPE_CHECKING:
-    from collections import defaultdict
     from collections.abc import (
         Collection,
         Hashable,
@@ -113,6 +114,16 @@
     insn_id: str
 
 
+@dataclass(frozen=True)
+class AllocTemp(ScheduleItem):
+    temp: str
+
+
+@dataclass(frozen=True)
+class DeallocTemp(ScheduleItem):
+    temp: str
+
+
 @dataclass(frozen=True)
 class CallKernel(BeginBlockItem):
     kernel_name: str
@@ -513,6 +524,10 @@
         elif isinstance(sched_item, Barrier):
             lines.append(indent + "... %sbarrier" %
                          sched_item.synchronization_kind[0])
+        elif isinstance(sched_item, AllocTemp):
+            lines.append(indent + "ALLOCATE %s" % sched_item.temp)
+        elif isinstance(sched_item, DeallocTemp):
+            lines.append(indent + "DEALLOCATE %s" % sched_item.temp)
         else:
             raise AssertionError()
 
@@ -2158,7 +2173,9 @@
                 dep_tracker.add_source(sched_item.insn_id)
                 i += 1
 
-            elif isinstance(sched_item, (CallKernel, ReturnFromKernel)):
+            # TODO: check if this is correct
+            elif isinstance(sched_item, (CallKernel, ReturnFromKernel,
+                                         AllocTemp, DeallocTemp)):
                 result.append(sched_item)
                 i += 1
 
@@ -2188,7 +2205,8 @@
             i = new_i
 
         elif isinstance(sched_item,
-                (Barrier, RunInstruction, CallKernel, ReturnFromKernel)):
+                (Barrier, RunInstruction, CallKernel,
+                 ReturnFromKernel, AllocTemp, DeallocTemp)):
             result.append(sched_item)
             i += 1
 
@@ -2274,9 +2292,102 @@
         # Device mapper only gets run once.
         new_kernel = map_schedule_onto_host_or_device(new_kernel)
 
+    new_kernel = insert_temporary_alloc_dealloc(new_kernel)
     return new_kernel
 
 
+def insert_temporary_alloc_dealloc(kernel: LoopKernel):
+    schedule = kernel.linearization
+    assert schedule is not None
+    inserted_allocs = set()
+    inserted_deallocs = set()
+
+    kernel_entry_idx = -1
+    loops_entry_idxs = []
+
+    insert_alloc_idxs = defaultdict(frozenset)
+    for idx, sched_item in enumerate(schedule):
+        if isinstance(sched_item, EnterLoop):
+            loops_entry_idxs.append(idx)
+        elif isinstance(sched_item, LeaveLoop):
+            loops_entry_idxs.pop()
+        elif isinstance(sched_item, CallKernel):
+            kernel_entry_idx = idx
+        elif isinstance(sched_item, ReturnFromKernel):
+            kernel_entry_idx = -1
+        elif isinstance(sched_item, RunInstruction):
+            insn = kernel.id_to_insn[sched_item.insn_id]
+            for var_name in set(insn.assignee_var_names()):
+                if var_name not in kernel.temporary_variables:
+                    continue
+                tv = kernel.temporary_variables[var_name]
+                if tv.base_storage:
+                    var_name = tv.base_storage
+                    tv = kernel.temporary_variables[var_name]
+                if var_name in inserted_allocs:
+                    continue
+                if tv.address_space == AddressSpace.GLOBAL:
+                    inserted_allocs.add(var_name)
+                    target_idx = idx
+                    if (kernel_entry_idx != -1):
+                        # lift the alloc out of ALL loops
+                        target_idx = min(idx, kernel_entry_idx, *loops_entry_idxs)
+                    else:
+                        target_idx = min(idx, *loops_entry_idxs)
+                    insert_alloc_idxs[target_idx] |= {var_name}
+
+    kernel_entry_idx = -1
+    loops_entry_idxs = []
+    insert_dealloc_idxs = defaultdict(frozenset)
+    for idx, sched_item in enumerate(reversed(schedule)):
+        idx = len(schedule) - idx - 1
+        if isinstance(sched_item, LeaveLoop):
+            loops_entry_idxs.append(idx)
+        elif isinstance(sched_item, EnterLoop):
+            loops_entry_idxs.pop()
+        elif isinstance(sched_item, ReturnFromKernel):
+            kernel_entry_idx = idx
+        elif isinstance(sched_item, CallKernel):
+            kernel_entry_idx = -1
+        elif isinstance(sched_item, RunInstruction):
+            insn = kernel.id_to_insn[sched_item.insn_id]
+            for var_name in set(insn.read_dependency_names()):
+                if var_name not in kernel.temporary_variables:
+                    continue
+                tv = kernel.temporary_variables[var_name]
+                if tv.base_storage:
+                    var_name = tv.base_storage
+                    tv = kernel.temporary_variables[var_name]
+                if var_name in inserted_deallocs:
+                    continue
+                if tv.address_space != AddressSpace.GLOBAL:
+                    continue
+                if var_name in inserted_allocs:
+                    inserted_deallocs.add(var_name)
+                    target_idx = idx
+                    if (kernel_entry_idx != -1):
+                        # lift the dealloc out of ALL loops
+                        target_idx = max(idx, kernel_entry_idx, *loops_entry_idxs)
+                    else:
+                        target_idx = max(idx, *loops_entry_idxs)
+                    insert_dealloc_idxs[target_idx] |= {var_name}
+
+    insert_dealloc_idxs[len(schedule)-1] |= (inserted_allocs - inserted_deallocs)
+
+    new_schedule = []
+    for idx, sched_item in enumerate(schedule):
+        if idx in insert_alloc_idxs:
+            for var_name in insert_alloc_idxs[idx]:
+                new_schedule.append(AllocTemp(var_name))
+        new_schedule.append(sched_item)
+
+        if idx in insert_dealloc_idxs:
+            for var_name in insert_dealloc_idxs[idx]:
+                new_schedule.append(DeallocTemp(var_name))
+
+    return kernel.copy(linearization=new_schedule)
+
+
 def _generate_loop_schedules_inner(
         kernel: LoopKernel,
         callables_table: CallablesTable,

loopy/target/__init__.py

@@ -323,6 +323,11 @@ def emit_comment(self, s):
     def emit_noop_with_comment(self, s):
         raise NotImplementedError()
 
+    def emit_alloc_temp(self, codegen_state, var_name):
+        raise NotImplementedError()
+
+    def emit_dealloc_temp(self, codegen_state, var_name):
+        raise NotImplementedError()
     # }}}
 
     def process_ast(self, node: ASTType):

loopy/target/c/__init__.py

@@ -67,6 +67,7 @@
     CallInstruction,
     VarAtomicity,
 )
+from loopy.schedule import CallKernel
 from loopy.symbolic import IdentityMapper
 from loopy.target import ASTBuilderBase, DummyHostASTBuilder, TargetBase
 from loopy.tools import remove_common_indentation
@@ -83,7 +84,6 @@
     from loopy.codegen.result import CodeGenerationResult
     from loopy.kernel import LoopKernel
     from loopy.kernel.instruction import MultiAssignmentBase
-    from loopy.schedule import CallKernel
     from loopy.target.execution import ExecutorBase
     from loopy.translation_unit import (
         CallableId,
@@ -1009,10 +1009,13 @@ def get_function_declaration(self,
             ) -> tuple[Sequence[tuple[str, str]], Generable]:
         kernel = codegen_state.kernel
 
-        assert codegen_state.kernel.linearization is not None
+        assert kernel.linearization is not None
+        while not isinstance(kernel.linearization[schedule_index], CallKernel):
+            schedule_index += 1
+            assert schedule_index < len(kernel.linearization)
         subkernel_name = cast(
                         "CallKernel",
-                        codegen_state.kernel.linearization[schedule_index]
+                        kernel.linearization[schedule_index]
                         ).kernel_name
 
         from cgen import FunctionDeclaration, Value
@@ -1105,6 +1108,11 @@ def get_temporary_decls(self, codegen_state, schedule_index):
             temporaries_read_in_subkernel,
             temporaries_written_in_subkernel,
         )
+
+        while not isinstance(kernel.linearization[schedule_index], CallKernel):
+            schedule_index += 1
+            assert schedule_index < len(kernel.linearization)
+
         subkernel_name = kernel.linearization[schedule_index].kernel_name
         sub_knl_temps = (
                 temporaries_read_in_subkernel(kernel, subkernel_name)
@@ -1144,6 +1152,16 @@ def get_temporary_decls(self, codegen_state, schedule_index):
 
         return result
 
+    @override
+    def emit_alloc_temp(self, codegen_state, var_name):
+        from cgen import Line
+        return Line()
+
+    @override
+    def emit_dealloc_temp(self, codegen_state, var_name):
+        from cgen import Line
+        return Line()
+
     @property
     @override
     def ast_block_class(self):