Merge branch 'main' into move-to-versioningit

CONTRIBUTING.md

@@ -146,11 +146,11 @@ Check `pytest` documentation (`pytest --help`) for all the options to select and
 We recommended you to use `nox` for running the test suite in different environments. `nox` runs the package installation script in properly isolated environments to run tests in a reproducible way. A simple way to start with `nox` would be:
 
 ```bash
-# List all the available task environments
-nox list
+# List all available sessions
+nox --list
 
-# Run a specific task environment
-nox -e cartesian-py38-internal-cpu
+# Run a specific session
+nox -s "test_cartesian-3.10(internal, cpu)"
 ```
 
 Check `nox` documentation (`nox --help`) for the complete reference.

src/gt4py/cartesian/backend/cuda_backend.py

@@ -141,7 +141,7 @@ class CudaBackend(BaseGTBackend, CLIBackendMixin):
     GT_BACKEND_T = "gpu"
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
-        return self.make_extension(stencil_ir=self.builder.gtir, uses_cuda=True)
+        return self.make_extension(uses_cuda=True)
 
     def generate(self) -> Type[StencilObject]:
         self.check_options(self.builder.options)

src/gt4py/cartesian/backend/dace_backend.py

@@ -12,7 +12,6 @@
 import os
 import pathlib
 import re
-import textwrap
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
 
 import dace
@@ -432,20 +431,20 @@ def __call__(self, stencil_ir: gtir.Stencil) -> Dict[str, Dict[str, str]]:
 
 class DaCeComputationCodegen:
     template = as_mako(
-        """
-        auto ${name}(const std::array<gt::uint_t, 3>& domain) {
-            return [domain](${",".join(functor_args)}) {
-                const int __I = domain[0];
-                const int __J = domain[1];
-                const int __K = domain[2];
-                ${name}${state_suffix} dace_handle;
-                ${backend_specifics}
-                auto allocator = gt::sid::cached_allocator(&${allocator}<char[]>);
-                ${"\\n".join(tmp_allocs)}
-                __program_${name}(${",".join(["&dace_handle", *dace_args])});
-            };
-        }
-        """
+        """\
+auto ${name}(const std::array<gt::uint_t, 3>& domain) {
+    return [domain](${",".join(functor_args)}) {
+        const int __I = domain[0];
+        const int __J = domain[1];
+        const int __K = domain[2];
+        ${name}${state_suffix} dace_handle;
+        ${backend_specifics}
+        auto allocator = gt::sid::cached_allocator(&${allocator}<char[]>);
+        ${"\\n".join(tmp_allocs)}
+        __program_${name}(${",".join(["&dace_handle", *dace_args])});
+    };
+}
+"""
     )
 
     def generate_tmp_allocs(self, sdfg):
@@ -511,7 +510,7 @@ def _postprocess_dace_code(code_objects, is_gpu, builder):
                     lines = lines[0:i] + cuda_code.split("\n") + lines[i + 1 :]
                     break
 
-        def keep_line(line):
+        def keep_line(line: str) -> bool:
             line = line.strip()
             if line == '#include "../../include/hash.h"':
                 return False
@@ -521,11 +520,7 @@ def keep_line(line):
                 return False
             return True
 
-        lines = filter(keep_line, lines)
-        generated_code = "\n".join(lines)
-        if builder.options.format_source:
-            generated_code = codegen.format_source("cpp", generated_code, style="LLVM")
-        return generated_code
+        return "\n".join(filter(keep_line, lines))
 
     @classmethod
     def apply(cls, stencil_ir: gtir.Stencil, builder: StencilBuilder, sdfg: dace.SDFG):
@@ -563,17 +558,18 @@ def apply(cls, stencil_ir: gtir.Stencil, builder: StencilBuilder, sdfg: dace.SDF
             allocator="gt::cuda_util::cuda_malloc" if is_gpu else "std::make_unique",
             state_suffix=dace.Config.get("compiler.codegen_state_struct_suffix"),
         )
-        generated_code = textwrap.dedent(
-            f"""#include <gridtools/sid/sid_shift_origin.hpp>
-                             #include <gridtools/sid/allocator.hpp>
-                             #include <gridtools/stencil/cartesian.hpp>
-                             {"#include <gridtools/common/cuda_util.hpp>" if is_gpu else omp_header}
-                             namespace gt = gridtools;
-                             {computations}
-
-                             {interface}
-                             """
-        )
+        generated_code = f"""\
+#include <gridtools/sid/sid_shift_origin.hpp>
+#include <gridtools/sid/allocator.hpp>
+#include <gridtools/stencil/cartesian.hpp>
+{"#include <gridtools/common/cuda_util.hpp>" if is_gpu else omp_header}
+namespace gt = gridtools;
+
+{computations}
+
+{interface}
+"""
+
         if builder.options.format_source:
             generated_code = codegen.format_source("cpp", generated_code, style="LLVM")
 
@@ -794,7 +790,7 @@ class DaceCPUBackend(BaseDaceBackend):
     options = BaseGTBackend.GT_BACKEND_OPTS
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
-        return self.make_extension(stencil_ir=self.builder.gtir, uses_cuda=False)
+        return self.make_extension(uses_cuda=False)
 
 
 @register
@@ -815,4 +811,4 @@ class DaceGPUBackend(BaseDaceBackend):
     options = {**BaseGTBackend.GT_BACKEND_OPTS, "device_sync": {"versioning": True, "type": bool}}
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
-        return self.make_extension(stencil_ir=self.builder.gtir, uses_cuda=True)
+        return self.make_extension(uses_cuda=True)

src/gt4py/cartesian/backend/gtc_common.py

@@ -236,19 +236,15 @@ def generate(self) -> Type[StencilObject]:
 
     def generate_computation(self) -> Dict[str, Union[str, Dict]]:
         dir_name = f"{self.builder.options.name}_src"
-        src_files = self.make_extension_sources(stencil_ir=self.builder.gtir)
+        src_files = self._make_extension_sources()
         return {dir_name: src_files["computation"]}
 
-    def generate_bindings(
-        self, language_name: str, *, stencil_ir: Optional[gtir.Stencil] = None
-    ) -> Dict[str, Union[str, Dict]]:
-        if not stencil_ir:
-            stencil_ir = self.builder.gtir
-        assert stencil_ir is not None
+    def generate_bindings(self, language_name: str) -> Dict[str, Union[str, Dict]]:
         if language_name != "python":
             return super().generate_bindings(language_name)
+
         dir_name = f"{self.builder.options.name}_src"
-        src_files = self.make_extension_sources(stencil_ir=stencil_ir)
+        src_files = self._make_extension_sources()
         return {dir_name: src_files["bindings"]}
 
     @abc.abstractmethod
@@ -260,32 +256,26 @@ def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
         """
         pass
 
-    def make_extension(
-        self, *, stencil_ir: Optional[gtir.Stencil] = None, uses_cuda: bool = False
-    ) -> Tuple[str, str]:
+    def make_extension(self, *, uses_cuda: bool = False) -> Tuple[str, str]:
         build_info = self.builder.options.build_info
         if build_info is not None:
             start_time = time.perf_counter()
 
-        if not stencil_ir:
-            stencil_ir = self.builder.gtir
-        assert stencil_ir is not None
-
         # Generate source
         gt_pyext_files: Dict[str, Any]
         gt_pyext_sources: Dict[str, Any]
-        if not self.builder.options._impl_opts.get("disable-code-generation", False):
-            gt_pyext_files = self.make_extension_sources(stencil_ir=stencil_ir)
-            gt_pyext_sources = {
-                **gt_pyext_files["computation"],
-                **gt_pyext_files["bindings"],
-            }
-        else:
+        if self.builder.options._impl_opts.get("disable-code-generation", False):
             # Pass NOTHING to the self.builder means try to reuse the source code files
             gt_pyext_files = {}
             gt_pyext_sources = {
                 key: gt_utils.NOTHING for key in self.PYEXT_GENERATOR_CLASS.TEMPLATE_FILES.keys()
             }
+        else:
+            gt_pyext_files = self._make_extension_sources()
+            gt_pyext_sources = {
+                **gt_pyext_files["computation"],
+                **gt_pyext_files["bindings"],
+            }
 
         if build_info is not None:
             next_time = time.perf_counter()
@@ -317,18 +307,19 @@ def make_extension(
 
         return result
 
-    def make_extension_sources(self, *, stencil_ir: gtir.Stencil) -> Dict[str, Dict[str, str]]:
+    def _make_extension_sources(self) -> Dict[str, Dict[str, str]]:
         """Generate the source for the stencil independently from use case."""
         if "computation_src" in self.builder.backend_data:
             return self.builder.backend_data["computation_src"]
+
         class_name = self.pyext_class_name if self.builder.stencil_id else self.builder.options.name
         module_name = (
             self.pyext_module_name
             if self.builder.stencil_id
             else f"{self.builder.options.name}_pyext"
         )
         gt_pyext_generator = self.PYEXT_GENERATOR_CLASS(class_name, module_name, self)
-        gt_pyext_sources = gt_pyext_generator(stencil_ir)
+        gt_pyext_sources = gt_pyext_generator(self.builder.gtir)
         final_ext = ".cu" if self.languages and self.languages["computation"] == "cuda" else ".cpp"
         comp_src = gt_pyext_sources["computation"]
         for key in [k for k in comp_src.keys() if k.endswith(".src")]:

src/gt4py/cartesian/backend/gtcpp_backend.py

@@ -129,7 +129,7 @@ class GTBaseBackend(BaseGTBackend, CLIBackendMixin):
     PYEXT_GENERATOR_CLASS = GTExtGenerator
 
     def _generate_extension(self, uses_cuda: bool) -> Tuple[str, str]:
-        return self.make_extension(stencil_ir=self.builder.gtir, uses_cuda=uses_cuda)
+        return self.make_extension(uses_cuda=uses_cuda)
 
     def generate(self) -> Type[StencilObject]:
         self.check_options(self.builder.options)