Don't merge: static args and hacks

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -384,6 +384,10 @@ def __init__(
         self._edge_params = edge_params
         self._cell_params = cell_params
 
+        import cupy as cp
+        # assert (m := cp.amin(self._grid.offset_providers["E2C"].ndarray)) >= 0, m
+        # assert (m := cp.amin(self._grid.offset_providers["E2C2E"].ndarray)) >= 0, m
+
         self.halo_exchange_wait = decomposition.create_halo_exchange_wait(
             self._exchange
         )  # wait on a communication handle
@@ -396,38 +400,80 @@ def __init__(
 
         self.mo_intp_rbf_rbf_vec_interpol_vertex = mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(
             self._backend
+        ).compile(
+            vertical_start=[0],
+            vertical_end=[self._grid.num_levels],
+            offset_provider_type=self._grid.offset_providers,
         )
         self.calculate_nabla2_and_smag_coefficients_for_vn = (
-            calculate_nabla2_and_smag_coefficients_for_vn.with_backend(self._backend)
+            calculate_nabla2_and_smag_coefficients_for_vn.with_backend(self._backend).compile(
+                vertical_start=[0],
+                vertical_end=[self._grid.num_levels],
+                offset_provider_type=self._grid.offset_providers,
+            )
         )
+
         self.calculate_diagnostic_quantities_for_turbulence = (
-            calculate_diagnostic_quantities_for_turbulence.with_backend(self._backend)
+            calculate_diagnostic_quantities_for_turbulence.with_backend(self._backend).compile(
+                vertical_start=[1],
+                vertical_end=[self._grid.num_levels],
+                offset_provider_type=self._grid.offset_providers,
+            )
+        )
+        self.apply_diffusion_to_vn = apply_diffusion_to_vn.with_backend(self._backend).compile(
+            limited_area=[self._grid.limited_area],
+            vertical_start=[0],
+            vertical_end=[self._grid.num_levels],
+            offset_provider_type=self._grid.offset_providers,
         )
-        self.apply_diffusion_to_vn = apply_diffusion_to_vn.with_backend(self._backend)
         self.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence = (
             apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_backend(
                 self._backend
+            ).compile(
+                type_shear=[int32(self.config.shear_type.value)],
+                nrdmax=[int32(self._vertical_grid.end_index_of_damping_layer + 1)],
+                vertical_start=[0],
+                vertical_end=[self._grid.num_levels],
+                offset_provider_type=self._grid.offset_providers,
             )
         )
         self.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools = (
             calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend(
                 self._backend
+            ).compile(
+                vertical_start=[(self._grid.num_levels - 2)],
+                vertical_end=[self._grid.num_levels],
+                offset_provider_type=self._grid.offset_providers,
             )
         )
-        self.calculate_nabla2_for_theta = calculate_nabla2_for_theta.with_backend(self._backend)
+        self.calculate_nabla2_for_theta = calculate_nabla2_for_theta.with_backend(
+            self._backend
+        ).compile(
+            vertical_start=[0],
+            vertical_end=[self._grid.num_levels],
+            offset_provider_type=self._grid.offset_providers,
+        )
         self.truly_horizontal_diffusion_nabla_of_theta_over_steep_points = (
             truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(self._backend)
-        )
-        self.update_theta_and_exner = update_theta_and_exner.with_backend(self._backend)
-        self.copy_field = copy_field.with_backend(self._backend)
-        self.scale_k = scale_k.with_backend(self._backend)
+        ).compile(
+            vertical_start=[0],
+            vertical_end=[self._grid.num_levels],
+            offset_provider_type=self._grid.offset_providers,
+        )
+        self.update_theta_and_exner = update_theta_and_exner.with_backend(self._backend).compile(
+            vertical_start=[0],
+            vertical_end=[self._grid.num_levels],
+            offset_provider_type={},
+        )
+        self.copy_field = copy_field.with_backend(self._backend).compile(offset_provider_type={})
+        self.scale_k = scale_k.with_backend(self._backend).compile(offset_provider_type={})
         self.setup_fields_for_initial_step = setup_fields_for_initial_step.with_backend(
             self._backend
-        )
+        ).compile(offset_provider_type={})
 
         self.init_diffusion_local_fields_for_regular_timestep = (
             init_diffusion_local_fields_for_regular_timestep.with_backend(self._backend)
-        )
+        ).compile(offset_provider_type={"Koff": dims.KDim})
 
         self._allocate_temporary_fields()
 
@@ -452,16 +498,17 @@ def __init__(
             offset_provider={"Koff": dims.KDim},
         )
 
-        diffusion_utils._init_nabla2_factor_in_upper_damping_zone.with_backend(self._backend)(
+        diffusion_utils.init_nabla2_factor_in_upper_damping_zone.with_backend(self._backend)(
             physical_heights=self._vertical_grid.interface_physical_height,
+            diff_multfac_n2w=self.diff_multfac_n2w,
             end_index_of_damping_layer=self._vertical_grid.end_index_of_damping_layer,
             nshift=0,
             heights_nrd_shift=self._vertical_grid.interface_physical_height.ndarray[
                 self._vertical_grid.end_index_of_damping_layer + 1
             ].item(),
             heights_1=self._vertical_grid.interface_physical_height.ndarray[1].item(),
-            domain={dims.KDim: (1, self._vertical_grid.end_index_of_damping_layer + 1)},
-            out=self.diff_multfac_n2w,
+            vertical_start=gtx.int32(1),
+            vertical_end=gtx.int32(self._vertical_grid.end_index_of_damping_layer + 1),
             offset_provider={},
         )
 
@@ -635,15 +682,10 @@ def _do_diffusion_step(
             smag_offset:
 
         """
-        # dtime dependent: enh_smag_factor,
-        self.scale_k.with_connectivities(self.compile_time_connectivities)(
-            self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}
-        )
+        self.scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={})
 
         log.debug("rbf interpolation 1: start")
-        self.mo_intp_rbf_rbf_vec_interpol_vertex.with_connectivities(
-            self.compile_time_connectivities
-        )(
+        self.mo_intp_rbf_rbf_vec_interpol_vertex(
             p_e_in=prognostic_state.vn,
             ptr_coeff_1=self._interpolation_state.rbf_coeff_1,
             ptr_coeff_2=self._interpolation_state.rbf_coeff_2,
@@ -668,9 +710,7 @@ def _do_diffusion_step(
         log.debug("communication rbf extrapolation of vn - end")
 
         log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start")
-        self.calculate_nabla2_and_smag_coefficients_for_vn.with_connectivities(
-            self.compile_time_connectivities
-        )(
+        self.calculate_nabla2_and_smag_coefficients_for_vn(
             diff_multfac_smag=self.diff_multfac_smag,
             tangent_orientation=self._edge_params.tangent_orientation,
             inv_primal_edge_length=self._edge_params.inverse_primal_edge_lengths,
@@ -702,9 +742,7 @@ def _do_diffusion_step(
             log.debug(
                 "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start"
             )
-            self.calculate_diagnostic_quantities_for_turbulence.with_connectivities(
-                self.compile_time_connectivities
-            )(
+            self.calculate_diagnostic_quantities_for_turbulence(
                 kh_smag_ec=self.kh_smag_ec,
                 vn=prognostic_state.vn,
                 e_bln_c_s=self._interpolation_state.e_bln_c_s,
@@ -731,9 +769,7 @@ def _do_diffusion_step(
             log.debug("communication rbf extrapolation of z_nable2_e - end")
 
         log.debug("2nd rbf interpolation: start")
-        self.mo_intp_rbf_rbf_vec_interpol_vertex.with_connectivities(
-            self.compile_time_connectivities
-        )(
+        self.mo_intp_rbf_rbf_vec_interpol_vertex(
             p_e_in=self.z_nabla2_e,
             ptr_coeff_1=self._interpolation_state.rbf_coeff_1,
             ptr_coeff_2=self._interpolation_state.rbf_coeff_2,
@@ -758,7 +794,7 @@ def _do_diffusion_step(
         log.debug("communication rbf extrapolation of z_nable2_e - end")
 
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start")
-        self.apply_diffusion_to_vn.with_connectivities(self.compile_time_connectivities)(
+        self.apply_diffusion_to_vn(
             u_vert=self.u_vert,
             v_vert=self.v_vert,
             primal_normal_vert_v1=self._edge_params.primal_normal_vert[0],
@@ -790,13 +826,9 @@ def _do_diffusion_step(
             "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"
         )
         # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify?
-        self.copy_field.with_connectivities(self.compile_time_connectivities)(
-            prognostic_state.w, self.w_tmp, offset_provider={}
-        )
+        self.copy_field(prognostic_state.w, self.w_tmp, offset_provider={})
 
-        self.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_connectivities(
-            self.compile_time_connectivities
-        )(
+        self.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
             area=self._cell_params.area,
             geofac_n2s=self._interpolation_state.geofac_n2s,
             geofac_grg_x=self._interpolation_state.geofac_grg_x,
@@ -830,9 +862,7 @@ def _do_diffusion_step(
                 "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start"
             )
 
-            self.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_connectivities(
-                self.compile_time_connectivities
-            )(
+            self.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools(
                 theta_v=prognostic_state.theta_v,
                 theta_ref_mc=self._metric_state.theta_ref_mc,
                 thresh_tdiff=self.thresh_tdiff,
@@ -849,7 +879,7 @@ def _do_diffusion_step(
             )
 
             log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start")
-            self.calculate_nabla2_for_theta.with_connectivities(self.compile_time_connectivities)(
+            self.calculate_nabla2_for_theta(
                 kh_smag_e=self.kh_smag_e,
                 inv_dual_edge_length=self._edge_params.inverse_dual_edge_lengths,
                 theta_v=prognostic_state.theta_v,
@@ -866,9 +896,7 @@ def _do_diffusion_step(
                 "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start"
             )
             if self.config.apply_zdiffusion_t:
-                self.truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_connectivities(
-                    self.compile_time_connectivities
-                )(
+                self.truly_horizontal_diffusion_nabla_of_theta_over_steep_points(
                     mask=self._metric_state.mask_hdiff,
                     zd_vertoffset=self._metric_state.zd_vertoffset,
                     zd_diffcoef=self._metric_state.zd_diffcoef,
@@ -888,7 +916,7 @@ def _do_diffusion_step(
                     "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end"
                 )
             log.debug("running stencil 16 (update_theta_and_exner): start")
-            self.update_theta_and_exner.with_connectivities(self.compile_time_connectivities)(
+            self.update_theta_and_exner(
                 z_temp=self.z_temp,
                 area=self._cell_params.area,
                 theta_v=prognostic_state.theta_v,