Opencl allocation

.basedpyright/baseline.json

@@ -8269,8 +8269,8 @@
                 "code": "reportUnknownArgumentType",
                 "range": {
                     "startColumn": 56,
-                    "endColumn": 17,
-                    "lineCount": 8
+                    "endColumn": 63,
+                    "lineCount": 1
                 }
             },
             {
@@ -8281,6 +8281,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

@@ -176,6 +176,109 @@ def supporting_temporary_names(
 
     return frozenset(result)
 
+
+def get_temporary_decl_blocks(
+        kernel: LoopKernel
+    ) -> tuple[dict[int, frozenset[str]], dict[int, frozenset[str]]]:
+    from loopy.kernel.data import AddressSpace
+    from loopy.schedule import CallKernel, EnterLoop
+
+    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
+    def get_temporaries_in_bounds(
+            linearization: Sequence[ScheduleItem],
+            lower_bound: int,
+            upper_bound: int
+        ) -> frozenset[str]:
+        temporaries: frozenset[str] = frozenset()
+        for sched_index in range(lower_bound, upper_bound+1):
+            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)
+                )
+        return temporaries & global_temporaries
+
+    block_boundaries = get_block_boundaries(kernel.linearization)
+
+    bounds: dict[int, frozenset[str]] = {}
+    sched_index = 0
+    while sched_index < len(kernel.linearization):
+        sched_item = kernel.linearization[sched_index]
+        if isinstance(sched_item, EnterLoop) or isinstance(sched_item, CallKernel):
+            if isinstance(sched_item, CallKernel):
+                block_end = block_boundaries[sched_index]
+                accessed_temporaries = (
+                    temporaries_written_in_subkernel(kernel, sched_item.kernel_name)
+                    | temporaries_read_in_subkernel(
+                        kernel, sched_item.kernel_name
+                    )
+                )
+            else:
+                block_end = block_boundaries[sched_index]
+                accessed_temporaries = get_temporaries_in_bounds(
+                    kernel.linearization, sched_index, block_end
+                )
+            bounds[sched_index] = accessed_temporaries
+            sched_index = block_end + 1
+        else:
+            sched_index += 1
+
+    def update_seen_storage_vars(
+            seen_sv: set[str],
+            new_temp_variables: frozenset[str]
+        ) -> frozenset[str]:
+        new_storage_variables: set[str] = set()
+        past_sv = frozenset(seen_sv)
+        for new_tv_name in new_temp_variables:
+            new_tv = kernel.temporary_variables[new_tv_name]
+            if new_tv.base_storage is None:
+                storage_var = new_tv_name
+            else:
+                storage_var = new_tv.base_storage
+            new_storage_variables.add(storage_var)
+            seen_sv.add(storage_var)
+        new_sv = frozenset(new_storage_variables)
+        return new_sv - past_sv
+    # forward pass for first accesses
+    first_accesses: dict[int, frozenset[str]] = {}
+    seen_storage_variables: set[str] = set()
+    for sched_index in range(0, len(kernel.linearization)):
+        if (sched_index not in bounds):
+            continue
+        sched_item = kernel.linearization[sched_index]
+        new_temporary_variables = bounds[sched_index]
+        new_storage_variables = update_seen_storage_vars(
+            seen_storage_variables, new_temporary_variables
+        )
+
+        if (len(new_storage_variables) > 0):
+            first_accesses[sched_index] = new_storage_variables
+
+    last_accesses: dict[int, frozenset[str]] = {}
+    seen_storage_variables.clear()
+    for sched_index in range(len(kernel.linearization)-1, -1, -1):
+        if (sched_index not in bounds):
+            continue
+        sched_item = kernel.linearization[sched_index]
+        new_temporary_variables = bounds[sched_index]
+        new_storage_variables = update_seen_storage_vars(
+            seen_storage_variables, new_temporary_variables
+        )
+
+        if (len(new_storage_variables) > 0):
+            last_accesses[sched_index] = new_storage_variables
+    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,21 @@ def get_temporary_decls(self, codegen_state: CodeGenerationState,
             schedule_index: int) -> ASTType:
         raise NotImplementedError
 
+    def get_temporary_var_declarator(self,
+            codegen_state: CodeGenerationState,
+            temp_var: TemporaryVariable) -> ASTType:
+        raise NotImplementedError()
+
+    def get_temporary_var_deallocator(self,
+            codegen_state: CodeGenerationState,
+            temp_var: TemporaryVariable) -> ASTType:
+        raise NotImplementedError()
+
+    def get_temporary_decl_at_index(
+            self, codegen_state: CodeGenerationState,
+            sched_index: int) -> tuple[ASTType | None, ASTType | None]:
+        raise NotImplementedError()
+
     def get_kernel_call(self, codegen_state: CodeGenerationState,
             subkernel_name: str,
             gsize: tuple[Expression, ...],
@@ -365,6 +381,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,
@@ -1107,6 +1108,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):
@@ -1240,6 +1247,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:
@@ -1272,6 +1280,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