Opencl allocation

.basedpyright/baseline.json

@@ -8261,8 +8261,8 @@
                 "code": "reportUnknownArgumentType",
                 "range": {
                     "startColumn": 56,
-                    "endColumn": 17,
-                    "lineCount": 8
+                    "endColumn": 63,
+                    "lineCount": 1
                 }
             },
             {
@@ -8273,6 +8273,14 @@
                     "lineCount": 1
                 }
             },
+            {
+                "code": "reportUnknownArgumentType",
+                "range": {
+                    "startColumn": 52,
+                    "endColumn": 59,
+                    "lineCount": 1
+                }
+            },
             {
                 "code": "reportUnknownMemberType",
                 "range": {

loopy/codegen/control.py

@@ -85,14 +85,23 @@ def generate_code_for_sched_index(
             glob_grid, loc_grid = kernel.get_grid_sizes_for_insn_ids_as_exprs(
                     get_insn_ids_for_block_at(kernel.linearization, sched_index),
                     codegen_state.callables_table)
-            return merge_codegen_results(codegen_state, [
-                codegen_result,
 
+            prefixes, suffixes = (
+                codegen_state.ast_builder.get_temporary_decl_at_index(
+                    codegen_state, sched_index
+                )
+            )
+            results = [
+                prefixes,
+                codegen_result,
                 codegen_state.ast_builder.get_kernel_call(
                     codegen_state,
                     sched_item.kernel_name,
-                    glob_grid, loc_grid)
-                ])
+                    glob_grid, loc_grid),
+                suffixes
+            ]
+            results = [r for r in results if r is not None]
+            return merge_codegen_results(codegen_state, results)
         else:
             # do not generate host code for non-entrypoint kernels
             return codegen_result
@@ -136,7 +145,14 @@ def generate_code_for_sched_index(
                     "for '%s', tagged '%s'"
                     % (sched_item.iname, ", ".join(str(tag) for tag in tags)))
 
-        return func(codegen_state, sched_index)
+        prefixes, suffixes = (
+            codegen_state.ast_builder.get_temporary_decl_at_index(
+                codegen_state, sched_index
+            )
+        )
+        results = [prefixes, func(codegen_state, sched_index), suffixes]
+        results = [r for r in results if r is not None]
+        return merge_codegen_results(codegen_state, results)
 
     elif isinstance(sched_item, Barrier):
         # {{{ emit barrier code

loopy/schedule/tools.py

@@ -177,6 +177,137 @@ def supporting_temporary_names(
 
     return frozenset(result)
 
+
+def _get_temporaries_accessed_in_schedule(
+            kernel: LoopKernel,
+            sched_idx_lower_bound: int,
+            sched_idx_upper_bound: int
+        ) -> frozenset[str]:
+    from loopy.schedule import CallKernel, EnterLoop, LeaveLoop
+
+    linearization = kernel.linearization
+    assert linearization is not None
+
+    temporaries: frozenset[str] = frozenset()
+    for sched_index in range(sched_idx_lower_bound, sched_idx_upper_bound):
+        sched_item = linearization[sched_index]
+        if isinstance(sched_item, CallKernel):
+            temporaries = (
+                temporaries_written_in_subkernel(kernel, sched_item.kernel_name)
+                | temporaries_read_in_subkernel(
+                    kernel, sched_item.kernel_name
+                )
+                | (temporaries)
+            )
+        elif isinstance(sched_item, (EnterLoop, LeaveLoop)):
+            # ignore further outside-kernel loops
+            pass
+
+        else:
+            raise NotImplementedError("kernel with non-CallKernel outermost")
+
+    return temporaries
+
+
+def _map_to_base_storage(kernel: LoopKernel, tv_names: Set[str]) -> Set[str]:
+    result: set[str] = set()
+    for tv_name in tv_names:
+        while True:
+            tv = kernel.temporary_variables[tv_name]
+            if tv.base_storage is not None:
+                tv_name = tv.base_storage
+            else:
+                break
+
+        result.add(tv_name)
+
+    return result
+
+
+@memoize_on_first_arg
+def get_sched_index_to_first_and_last_used(
+        kernel: LoopKernel
+    ) -> tuple[Mapping[int, Set[str]], Mapping[int, Set[str]]]:
+    """
+    Returns the tuple (first_used, last_used), where first_used is
+    a dict such that first_used[sched_index] is the set of all global temporary
+    variable names first used at sched_index.
+
+    Likewise, last_used[sched_index] is the set of all global temporary
+    variable names last used at sched_index.
+    """
+    from loopy.kernel.data import AddressSpace
+    from loopy.schedule import CallKernel, EnterLoop, Barrier
+
+    assert kernel.linearization is not None
+
+    global_temporaries = frozenset(
+        tv.name for tv in kernel.temporary_variables.values()
+        if tv.address_space == AddressSpace.GLOBAL
+    )
+
+    # Collapse into blocks
+    block_boundaries = get_block_boundaries(kernel.linearization)
+
+    tvs_accessed_at: dict[int, frozenset[str]] = {}
+    sched_index = 0
+    while sched_index < len(kernel.linearization):
+        sched_item = kernel.linearization[sched_index]
+        if isinstance(sched_item, CallKernel):
+            block_end = block_boundaries[sched_index]
+            tvs_accessed_at[sched_index] = (
+                temporaries_written_in_subkernel(kernel, sched_item.kernel_name)
+                | temporaries_read_in_subkernel(
+                    kernel, sched_item.kernel_name
+                )
+            ) & global_temporaries
+
+            sched_index = block_end + 1
+
+        elif isinstance(sched_item, EnterLoop):
+            block_end = block_boundaries[sched_index]
+            tvs_accessed_at[sched_index] = _get_temporaries_accessed_in_schedule(
+                kernel, sched_index, block_end+1
+            ) & global_temporaries
+
+            sched_index = block_end + 1
+
+        elif isinstance(sched_item, Barrier):
+            sched_index += 1
+        else:
+            raise ValueError(
+                    f"unexpected schedule item at outermost level: {type(sched_item)}")
+
+    storage_vars_accessed_at = {
+        sched_index: _map_to_base_storage(kernel, accessed)
+        for sched_index, accessed in tvs_accessed_at.items()
+    }
+    del tvs_accessed_at
+
+    # forward pass for first_accesses
+    first_accesses: dict[int, Set[str]] = {}
+    seen_storage_vars: set[str] = set()
+    for sched_index in range(0, len(kernel.linearization)):
+        accessed = storage_vars_accessed_at.get(sched_index, set())
+        new_storage_vars = accessed - seen_storage_vars
+        seen_storage_vars.update(accessed)
+
+        if new_storage_vars:
+            first_accesses[sched_index] = new_storage_vars
+
+    # backward pass for last_accesses
+    last_accesses: dict[int, Set[str]] = {}
+    seen_storage_vars = set()
+    for sched_index in range(len(kernel.linearization)-1, -1, -1):
+        accessed = storage_vars_accessed_at.get(sched_index, set())
+        new_storage_vars = accessed - seen_storage_vars
+        seen_storage_vars.update(accessed)
+
+        if new_storage_vars:
+            last_accesses[sched_index] = new_storage_vars
+
+    return (first_accesses, last_accesses)
+
 # }}}
 
 

loopy/target/__init__.py

@@ -59,6 +59,7 @@
     from loopy.codegen import CodeGenerationState, PreambleInfo
     from loopy.codegen.result import CodeGenerationResult
     from loopy.kernel import LoopKernel
+    from loopy.kernel.data import TemporaryVariable
     from loopy.target.c import DTypeRegistry
     from loopy.target.execution import ExecutorBase
     from loopy.translation_unit import CallableId, CallablesTable, TranslationUnit
@@ -251,6 +252,27 @@ def get_temporary_decls(self, codegen_state: CodeGenerationState,
             schedule_index: int) -> ASTType:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_temporary_var_declarator(self,
+                codegen_state: CodeGenerationState,
+                temp_var: TemporaryVariable
+            ) -> ASTType | None:
+        ...
+
+    @abstractmethod
+    def get_temporary_var_deallocator(self,
+                codegen_state: CodeGenerationState,
+                temp_var: TemporaryVariable
+            ) -> ASTType | None:
+        ...
+
+    @abstractmethod
+    def get_temporary_decl_at_index(
+                self, codegen_state: CodeGenerationState,
+                sched_index: int
+            ) -> tuple[ASTType | None, ASTType | None]:
+        ...
+
     def get_kernel_call(self, codegen_state: CodeGenerationState,
             subkernel_name: str,
             gsize: tuple[Expression, ...],
@@ -365,6 +387,27 @@ def get_expression_to_code_mapper(self, codegen_state):
     def get_kernel_call(self, codegen_state, name, gsize, lsize):
         return None
 
+    @override
+    def get_temporary_var_declarator(
+        self, codegen_state: CodeGenerationState,
+        temp_var: TemporaryVariable
+    ) -> None:
+        return None
+
+    @override
+    def get_temporary_var_deallocator(
+        self, codegen_state: CodeGenerationState,
+        temp_var: TemporaryVariable
+    ) -> None:
+        return None
+
+    @override
+    def get_temporary_decl_at_index(
+        self, codegen_state: CodeGenerationState,
+        sched_index: int
+    ) -> tuple[None, None]:
+        return (None, None)
+
     @property
     def ast_block_class(self):
         return _DummyASTBlock

loopy/target/c/__init__.py

@@ -37,6 +37,7 @@
 from cgen import (
     Block,
     Collection,
+    Comment,
     Const,
     Declarator,
     Generable,
@@ -1109,6 +1110,12 @@ def get_temporary_decls(self, codegen_state, schedule_index):
 
         return result
 
+    @override
+    def get_temporary_decl_at_index(
+            self, codegen_state: CodeGenerationState, sched_index: int
+            ) -> tuple[Generable | None, Generable | None]:
+        return (None, None)
+
     @property
     @override
     def ast_block_class(self):
@@ -1242,6 +1249,7 @@ def arg_to_cgen_declarator(
             raise ValueError(f"unexpected type of argument '{passed_name}': "
                     f"'{type(var_descr)}'")
 
+    @override
     def get_temporary_var_declarator(self,
             codegen_state: CodeGenerationState,
             temp_var: TemporaryVariable) -> Declarator:
@@ -1274,6 +1282,12 @@ def get_temporary_var_declarator(self,
         return self.wrap_decl_for_address_space(temp_var_decl,
                 temp_var.address_space)
 
+    @override
+    def get_temporary_var_deallocator(self,
+                codegen_state: CodeGenerationState,
+                temp_var: TemporaryVariable
+            ) -> Generable:
+        return Comment("Dynamic freeing of temp vars not supported")
     # }}}
 
     @override