Add driver bencher benchmarks

.cscs-ci/benchmark_bencher.yml

@@ -5,19 +5,29 @@ include:
 .bencher_feature_tests:
   extends: [.benchmark_nox_job]
   script:
-    - export PR_ID=$(echo "${CI_COMMIT_BRANCH}" | grep -o 'pr[0-9]*' | grep -o '[0-9]*')
-    - export FEATURE_BRANCH=$(curl -s https://api.github.com/repos/C2SM/icon4py/pulls/$PR_ID | jq -r '.head.ref')
-    - export GITHUB_ACTIONS=true
-    - export GITHUB_EVENT_NAME=pull_request
-    - export GITHUB_STEP_SUMMARY=$CI_PROJECT_DIR/step_summary.log
-    - export GITHUB_SHA=$CI_COMMIT_SHA
-    - export GITHUB_EVENT_PATH=$CI_PROJECT_DIR/event.json
-    - |
-      echo "{\"pull_request\": {\"head\": {\"repo\": {\"full_name\": \"C2SM/icon4py\"}}}, \"repository\": {\"full_name\": \"C2SM/icon4py\"}, \"number\": $PR_ID}" > $CI_PROJECT_DIR/event.json
+    - !reference [.bencher_github_event_setup, script]
     - !reference [.benchmark_nox_job, script]
   variables:
     NOX_SESSION: "__bencher_feature_branch_CI-${PYVERSION_SHORT}"
 
+.bencher_driver_mpi_feature_tests:
+  extends: [.test_runner_mpi, .test_template_aarch64, .benchmark_driver_mpi_nox_job]
+  stage: benchmark
+  script:
+    - !reference [.bencher_github_event_setup, script]
+    - !reference [.benchmark_driver_mpi_nox_job, script]
+  variables:
+    NOX_SESSION: "__bencher_driver_mpi_feature_branch_CI-${PYVERSION_SHORT}"
+
+.bencher_driver_feature_tests:
+  extends: [.test_runner_serial, .test_template_aarch64, .benchmark_driver_nox_job]
+  stage: benchmark
+  script:
+    - !reference [.bencher_github_event_setup, script]
+    - !reference [.benchmark_driver_nox_job, script]
+  variables:
+    NOX_SESSION: "__bencher_driver_feature_branch_CI-${PYVERSION_SHORT}"
+
 benchmark_bencher_stencils_feature_aarch64:
   extends: [.test_runner_serial, .test_template_aarch64, .bencher_feature_tests]
   variables:
@@ -28,3 +38,13 @@ benchmark_bencher_granules_feature_aarch64:
   extends: [.test_runner_serial, .test_template_aarch64, .bencher_feature_tests]
   variables:
     TEST_SELECTION: "${GRANULE_TEST_SELECTION}"
+
+benchmark_bencher_driver_mpi_feature_aarch64:
+  extends: [.bencher_driver_mpi_feature_tests]
+  variables:
+    SLURM_NTASKS: 4
+    SLURM_JOB_NUM_NODES: 1
+    ICON4PY_TEST_MPI_SUBCOMM_SIZE: $SLURM_NTASKS
+
+benchmark_bencher_driver_feature_aarch64:
+  extends: [.bencher_driver_feature_tests]

.cscs-ci/benchmark_bencher_baseline.yml

@@ -2,6 +2,18 @@ include:
   - local: '.cscs-ci/base.yml'
   - local: '.cscs-ci/benchmark_bencher_common.yml'
 
+.bencher_driver_mpi_baseline_tests:
+  extends: [.benchmark_driver_mpi_nox_job]
+  stage: benchmark
+  variables:
+    NOX_SESSION: "__bencher_driver_mpi_baseline_CI-${PYVERSION_SHORT}"
+
+.bencher_driver_baseline_tests:
+  extends: [.benchmark_driver_nox_job]
+  stage: benchmark
+  variables:
+    NOX_SESSION: "__bencher_driver_baseline_CI-${PYVERSION_SHORT}"
+
 benchmark_bencher_stencils_baseline_aarch64:
   extends: [.test_runner_serial, .test_template_aarch64, .benchmark_nox_job]
   variables:
@@ -14,3 +26,13 @@ benchmark_bencher_granules_baseline_aarch64:
   variables:
     NOX_SESSION: "__bencher_baseline_CI-${PYVERSION_SHORT}"
     TEST_SELECTION: "${GRANULE_TEST_SELECTION}"
+
+benchmark_bencher_driver_mpi_baseline_aarch64:
+  extends: [.test_runner_mpi, .test_template_aarch64, .bencher_driver_mpi_baseline_tests]
+  variables:
+    SLURM_NTASKS: 4
+    SLURM_JOB_NUM_NODES: 1
+    ICON4PY_TEST_MPI_SUBCOMM_SIZE: $SLURM_NTASKS
+
+benchmark_bencher_driver_baseline_aarch64:
+  extends: [.test_runner_serial, .test_template_aarch64, .bencher_driver_baseline_tests]

.cscs-ci/benchmark_bencher_common.yml

@@ -9,6 +9,13 @@
     OMP_PLACES: cores
     OMP_NUM_THREADS: 72
 
+.benchmark_driver_base_variables:
+  variables:
+    SLURM_TIMELIMIT: '02:00:00'
+    OMP_PROC_BIND: close
+    OMP_PLACES: cores
+    OMP_NUM_THREADS: 72
+
 .benchmark_nox_job:
   stage: benchmark
   script:
@@ -27,6 +34,52 @@
     SLURM_PARTITION: normal
     SLURM_TIMELIMIT: '01:30:00'
 
+.benchmark_driver_mpi_nox_job:
+  stage: benchmark
+  script:
+    - .cscs-ci/scripts/ci-mpi-wrapper.sh nox -s "${NOX_SESSION}" -- --backend=$BACKEND --grid=$GRID
+  # MPS is explicitly not enabled here. The job uses one GPU per task.
+  before_script:
+    - cd /icon4py
+    - source .cscs-ci/scripts/gt4py-cache.sh
+  extends: [.benchmark_driver_base_variables]
+  variables:
+    # allow-task-sharing lets NCCL open peer GPUs for GPU-GPU communication.
+    SLURM_GPUS_PER_NODE: 4
+    SLURM_CPUS_PER_TASK: 72
+    SLURM_GRES_FLAGS: allow-task-sharing
+  parallel:
+    matrix:
+      - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu]
+        GRID: [R02B06_GLOBAL]
+.benchmark_driver_nox_job:
+  stage: benchmark
+  script:
+    - nox -s "${NOX_SESSION}" -- --backend=$BACKEND --grid=$GRID
+  # MPS is explicitly not enabled here. The job uses one GPU per task.
+  before_script:
+    - cd /icon4py
+    - source .cscs-ci/scripts/gt4py-cache.sh
+  extends: [.benchmark_driver_base_variables]
+  variables:
+    SLURM_GPUS_PER_TASK: 1
+    SLURM_CPUS_PER_TASK: 72
+  parallel:
+    matrix:
+      - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu]
+        GRID: [R02B06_GLOBAL]
+.bencher_github_event_setup:
+  script:
+    - export PR_ID=$(echo "${CI_COMMIT_BRANCH}" | grep -o 'pr[0-9]*' | grep -o '[0-9]*')
+    - export FEATURE_BRANCH=$(curl -s https://api.github.com/repos/C2SM/icon4py/pulls/$PR_ID | jq -r '.head.ref')
+    - export GITHUB_ACTIONS=true
+    - export GITHUB_EVENT_NAME=pull_request
+    - export GITHUB_STEP_SUMMARY=$CI_PROJECT_DIR/step_summary.log
+    - export GITHUB_SHA=$CI_COMMIT_SHA
+    - export GITHUB_EVENT_PATH=$CI_PROJECT_DIR/event.json
+    - |
+      echo "{\"pull_request\": {\"head\": {\"repo\": {\"full_name\": \"C2SM/icon4py\"}}}, \"repository\": {\"full_name\": \"C2SM/icon4py\"}, \"number\": $PR_ID}" > $CI_PROJECT_DIR/event.json
+
 
 build_baseimage_aarch64:
   extends: [.build_baseimage_aarch64]

model/driver/src/icon4py/model/driver/driver.py

@@ -797,20 +797,16 @@ def initialize_driver(
     return icon4py_driver
 
 
-def run_driver(
-    *,
-    config: driver_config.ExperimentConfig,
-    grid_manager: gm.GridManager,
-    process_props: decomposition_defs.ProcessProperties,
-    backend: gtx.typing.Backend | None,
-) -> tuple[driver_states.DriverStates, Icon4pyDriver]:
-    icon4py_driver = initialize_driver(
-        config=config,
-        grid_manager=grid_manager,
-        process_props=process_props,
-        backend=backend,
-    )
-    allocator = model_backends.get_allocator(backend)
+def initialize_driver_states(
+    icon4py_driver: Icon4pyDriver,
+    allocator: gtx.typing.Allocator,
+) -> driver_states.DriverStates:
+    """Build and validate the initial driver states from a freshly initialized driver.
+
+    This wraps the creation of the prognostic/tracer/nonhydro/diagnostic states, the
+    application of the initial condition, the assembly of ``DriverStates``, and the
+    consistency check that precedes the time loop.
+    """
     prognostic_state_now = prognostics.initialize_prognostic_state(
         grid=icon4py_driver.grid,
         allocator=allocator,
@@ -859,5 +855,23 @@ def run_driver(
         granules=icon4py_driver.granules,
         states=ds,
     )
+    return ds
+
+
+def run_driver(
+    *,
+    config: driver_config.ExperimentConfig,
+    grid_manager: gm.GridManager,
+    process_props: decomposition_defs.ProcessProperties,
+    backend: gtx.typing.Backend | None,
+) -> tuple[driver_states.DriverStates, Icon4pyDriver]:
+    icon4py_driver = initialize_driver(
+        config=config,
+        grid_manager=grid_manager,
+        process_props=process_props,
+        backend=backend,
+    )
+    allocator = model_backends.get_allocator(backend)
+    ds = initialize_driver_states(icon4py_driver=icon4py_driver, allocator=allocator)
     icon4py_driver.time_integration(ds)
     return ds, icon4py_driver

model/driver/tests/driver/fixtures.py

@@ -5,8 +5,16 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from __future__ import annotations
+
+import gt4py.next.typing as gtx_typing
 import pytest
 
+from icon4py.model.common import model_backends, time
+from icon4py.model.common.decomposition import definitions as decomp_defs
+from icon4py.model.common.grid import grid_manager as gm
+from icon4py.model.driver import config as driver_config, driver_utils
+from icon4py.model.testing import datatest_utils as dt_utils, definitions as test_defs, grid_utils
 from icon4py.model.testing.fixtures import (
     backend,
     backend_like,
@@ -35,3 +43,88 @@
 @pytest.fixture
 def linit(timeloop_diffusion_linit_exit: bool) -> bool:
     return timeloop_diffusion_linit_exit
+
+
+BENCHMARK_EXPERIMENTS: list[test_defs.ExperimentDescription] = [test_defs.Experiments.JW]
+BENCHMARK_STEPS: int = 100
+BENCHMARK_ROUNDS: int = 5
+BENCHMARK_WARMUP_ROUNDS: int = 2
+
+
+def _resolve_grid(grid_option: str) -> test_defs.GridDescription:
+    name = grid_option.split(":", maxsplit=1)[0].strip()
+    try:
+        return getattr(test_defs.Grids, name)
+    except AttributeError:
+        raise pytest.UsageError(f"Unknown grid '{name}' in '--grid' option") from None
+
+
+def _make_config(
+    experiment: test_defs.ExperimentDescription,
+    grid: test_defs.GridDescription,
+    process_props: decomp_defs.ProcessProperties,
+) -> driver_config.ExperimentConfig:
+    dt_utils.download_experiment(experiment, process_props)
+    experiment_path = dt_utils.get_path_for_experiment(experiment, process_props)
+    config = driver_config.read_experiment_config_from_fortran(experiment_path)
+    return config.with_overrides(
+        driver={
+            "dtime": time.RelativeTime(seconds=50),
+            "enable_output": False,
+            "end_of_simulation": time.NumTimeSteps(BENCHMARK_STEPS),
+        }
+    )
+
+
+def _make_grid_manager(
+    config: driver_config.ExperimentConfig,
+    grid: test_defs.GridDescription,
+    process_props: decomp_defs.ProcessProperties,
+    backend: gtx_typing.Backend | None,
+) -> gm.GridManager:
+    allocator = model_backends.get_allocator(backend)
+    grid_file_path = grid_utils._download_grid_file(grid)
+    return driver_utils.create_grid_manager(
+        grid_file_path=grid_file_path,
+        vertical_grid_config=config.vertical_grid,
+        allocator=allocator,
+        process_props=process_props,
+    )
+
+
+@pytest.fixture(params=BENCHMARK_EXPERIMENTS, ids=lambda experiment: experiment.name)
+def driver_benchmark_experiment(request: pytest.FixtureRequest) -> test_defs.ExperimentDescription:
+    return request.param
+
+
+@pytest.fixture
+def driver_benchmark_grid(
+    request: pytest.FixtureRequest,
+    driver_benchmark_experiment: test_defs.ExperimentDescription,
+) -> test_defs.GridDescription:
+    grid_option = request.config.getoption("--grid")
+    return driver_benchmark_experiment.grid if grid_option is None else _resolve_grid(grid_option)
+
+
+@pytest.fixture
+def driver_benchmark_config(
+    driver_benchmark_experiment: test_defs.ExperimentDescription,
+    driver_benchmark_grid: test_defs.GridDescription,
+    process_props: decomp_defs.ProcessProperties,
+) -> driver_config.ExperimentConfig:
+    return _make_config(driver_benchmark_experiment, driver_benchmark_grid, process_props)
+
+
+@pytest.fixture
+def driver_benchmark_grid_manager(
+    driver_benchmark_config: driver_config.ExperimentConfig,
+    driver_benchmark_grid: test_defs.GridDescription,
+    process_props: decomp_defs.ProcessProperties,
+    backend: gtx_typing.Backend | None,
+) -> gm.GridManager:
+    return _make_grid_manager(
+        config=driver_benchmark_config,
+        grid=driver_benchmark_grid,
+        process_props=process_props,
+        backend=backend,
+    )