Model.n_subjects: AOT-compile simulate, lock integer dtype to int32

benchmarks/bench_aca_baseline.py

@@ -54,7 +54,7 @@ def _build() -> tuple[object, object, object]:
         get_benchmark_params,
     )
 
-    model = create_benchmark_model()
+    model = create_benchmark_model(n_subjects=_N_SUBJECTS)
     _, model_params = get_benchmark_params(model=model)
     initial_conditions = get_benchmark_initial_conditions(
         model=model, n_subjects=_N_SUBJECTS, seed=0

pyproject.toml

@@ -98,7 +98,7 @@ tests-cuda13 = { features = [ "tests", "cuda13" ], solve-group = "cuda13" }
 tests-metal = { features = [ "tests", "metal" ], solve-group = "metal" }
 type-checking = { features = [ "type-checking", "tests" ], solve-group = "default" }
 [tool.pixi.feature.benchmarks.pypi-dependencies]
-aca-model = { git = "https://github.com/OpenSourceEconomics/aca-model.git", rev = "134286108b7445f3e17e8824bcdd1739a98b6089" }
+aca-model = { git = "https://github.com/OpenSourceEconomics/aca-model.git", rev = "f09b5e34102ff42f739b95be5a9d388795b734a1" }
 [tool.pixi.feature.cuda12]
 platforms = [ "linux-64" ]
 system-requirements = { cuda = "12" }

src/lcm/dtypes.py

@@ -0,0 +1,48 @@
+"""Boundary-cast helpers that pin user-supplied data to canonical pylcm dtypes.
+
+Used at every API boundary that accepts user data (params, initial
+conditions, regime-id arrays) — always called from Python, never inside
+JIT. Each helper validates that the value fits the target dtype and
+raises a clearly-named error if not.
+
+Casts further down the simulate stack (e.g. transition outputs landing
+in the state pool) use plain `.astype` and rely on the boundary cast
+above them having already pinned the canonical dtype.
+"""
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+_INT32_MIN = int(np.iinfo(np.int32).min)
+_INT32_MAX = int(np.iinfo(np.int32).max)
+
+
+def safe_to_int32(value: object, *, name: str) -> Array:
+    """Cast a scalar, sequence, or array to `jnp.int32`, checking int32 range.
+
+    Args:
+        value: A Python int, numpy/JAX integer scalar, or array-like of
+            integer values.
+        name: Qualified name of the leaf — surfaced in the error message
+            so the user can locate the offending input.
+
+    Returns:
+        A `jnp.int32` array (0-d if `value` was a scalar).
+
+    Raises:
+        ValueError: If any element of `value` is outside the int32 range
+            `[-2**31, 2**31 - 1]`. The message names the leaf via `name`.
+
+    """
+    np_value = np.asarray(value)
+    if np_value.size > 0:
+        lo = int(np_value.min())
+        hi = int(np_value.max())
+        if lo < _INT32_MIN or hi > _INT32_MAX:
+            msg = (
+                f"{name}: int32 overflow — value range [{lo}, {hi}] "
+                f"exceeds [{_INT32_MIN}, {_INT32_MAX}]."
+            )
+            raise ValueError(msg)
+    return jnp.asarray(np_value, dtype=jnp.int32)

src/lcm/grids/discrete.py

@@ -48,5 +48,13 @@ def batch_size(self) -> int:
         return self.__batch_size
 
     def to_jax(self) -> Int1D:
-        """Convert the grid to a Jax array."""
-        return jnp.array(self.codes)
+        """Convert the grid to a Jax array.
+
+        Discrete state/action codes are pinned to `int32` regardless of the
+        ambient `jax_enable_x64` setting. A single integer dtype across
+        transitions, V-array indexing, and action lookups keeps the JIT cache
+        unsplit and lets AOT-compiled programs ship one signature. `int32`
+        accommodates any realistic category count and matches the
+        `MISSING_CAT_CODE` sentinel.
+        """
+        return jnp.array(self.codes, dtype=jnp.int32)

src/lcm/model.py

@@ -1,7 +1,10 @@
 """Collection of classes that are used by the user to define the model and grids."""
 
 import dataclasses
+import logging
+import threading
 from collections.abc import Mapping
+from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 from types import MappingProxyType
 
@@ -30,6 +33,7 @@
 )
 from lcm.regime import Regime
 from lcm.regime_building.processing import InternalRegime
+from lcm.simulation.compile import compile_all_simulate_functions
 from lcm.simulation.initial_conditions import validate_initial_conditions
 from lcm.simulation.result import SimulationResult, get_simulation_output_dtypes
 from lcm.simulation.simulate import simulate
@@ -85,9 +89,56 @@ class Model:
     fixed_params: UserParams
     """Parameters fixed at model initialization."""
 
+    n_subjects: int | None = None
+    """Expected simulate batch size; enables AOT compile of simulate functions.
+
+    Dispatch by call shape:
+
+    - `None`: purely lazy behaviour, no AOT.
+    - First `simulate(...)` with `actual_n == n_subjects`: AOT-compiles all
+      simulate functions for that batch shape in parallel and caches them.
+    - Subsequent `simulate(...)` with the same matching size: reuses the
+      cached compiled programs.
+    - `simulate(...)` with a mismatching size: warns once per size and falls
+      back to the runtime-traced path.
+
+    Param-shape contract: the cache is keyed only on `n_subjects`. The shapes
+    and dtypes of `internal_params` leaves at the first matching call become
+    part of the AOT signature; subsequent calls must keep them stable. MSM-
+    style estimation (varying values, fixed shapes) is the target use case;
+    construct a fresh `Model` whenever a param array's shape or dtype changes.
+    """
+
     _params_template: ParamsTemplate
     """Template for the model parameters."""
 
+    _simulate_compile_cache: dict[int, MappingProxyType[RegimeName, InternalRegime]]
+    """AOT-compiled `internal_regimes` per matching `n_subjects`."""
+
+    _warned_n_subjects: set[int]
+    """Mismatching `actual_n_subjects` already warned about (one warning each)."""
+
+    _simulate_compile_future: (
+        Future[MappingProxyType[RegimeName, InternalRegime]] | None
+    )
+    """Pending background AOT compile started by `solve(...)`, or `None`.
+
+    `solve(...)` kicks off `compile_all_simulate_functions` in a single
+    background thread so XLA compilation overlaps with the GPU-bound
+    backward induction. `simulate(...)` awaits the future before
+    dispatching the AOT-compiled program. Cleared after the result lands
+    in `_simulate_compile_cache`.
+    """
+
+    _simulate_compile_lock: threading.Lock
+    """Serialises mutations of `_simulate_compile_cache`, `_warned_n_subjects`,
+    and `_simulate_compile_future`.
+
+    The check-then-set on each container is held under this lock. The
+    consequent `log.warning` call sits outside the lock so concurrent
+    simulate() calls don't serialise on logging I/O.
+    """
+
     def __init__(
         self,
         *,
@@ -100,6 +151,7 @@ def __init__(
         derived_categoricals: Mapping[FunctionName, DiscreteGrid] = MappingProxyType(
             {}
         ),
+        n_subjects: int | None = None,
     ) -> None:
         """Initialize the Model.
 
@@ -115,17 +167,27 @@ def __init__(
                 not in states/actions. Broadcast to all regimes (merged with
                 each regime's own `derived_categoricals`). Raises if a regime
                 already has a conflicting entry.
+            n_subjects: Expected simulate batch size; if set, the first matching
+                `simulate(...)` call AOT-compiles all simulate functions for
+                batch shape `n_subjects` in parallel. `None` keeps the purely
+                lazy behaviour.
 
         """
         self.description = description
         self.ages = ages
         self.n_periods = ages.n_periods
         self.fixed_params = ensure_containers_are_immutable(fixed_params)
+        self.n_subjects = n_subjects
+        self._simulate_compile_cache = {}
+        self._warned_n_subjects = set()
+        self._simulate_compile_future = None
+        self._simulate_compile_lock = threading.Lock()
 
         validate_model_inputs(
             n_periods=self.n_periods,
             regimes=regimes,
             regime_id_class=regime_id_class,
+            n_subjects=n_subjects,
         )
         self.regime_names_to_ids = MappingProxyType(
             dict(
@@ -149,6 +211,31 @@ def __init__(
             regime_names_to_ids=self.regime_names_to_ids,
         )
 
+    def __getstate__(self) -> dict[str, object]:
+        """Return a copy of `__dict__` with per-process AOT compile state removed.
+
+        Drops `_simulate_compile_lock` (a `threading.Lock`, not pickleable),
+        `_simulate_compile_cache` (compiled XLA programs that can't survive
+        a process boundary), `_warned_n_subjects` (its companion set), and
+        `_simulate_compile_future` (a `Future` tied to the originating thread
+        pool).
+        `__setstate__` restores all three to their fresh state.
+        """
+        state = self.__dict__.copy()
+        state.pop("_simulate_compile_lock", None)
+        state.pop("_simulate_compile_cache", None)
+        state.pop("_warned_n_subjects", None)
+        state.pop("_simulate_compile_future", None)
+        return state
+
+    def __setstate__(self, state: dict[str, object]) -> None:
+        """Restore AOT compile state to a fresh empty cache."""
+        self.__dict__.update(state)
+        self._simulate_compile_cache = {}
+        self._warned_n_subjects = set()
+        self._simulate_compile_future = None
+        self._simulate_compile_lock = threading.Lock()
+
     def get_params_template(self) -> UserFacingParamsTemplate:
         """Get a human-readable params template.
 
@@ -210,6 +297,11 @@ def solve(
             internal_params=internal_params,
             ages=self.ages,
         )
+        self._maybe_start_simulate_compile_async(
+            internal_params=internal_params,
+            max_compilation_workers=max_compilation_workers,
+            logger=get_logger(log_level=log_level),
+        )
         try:
             period_to_regime_to_V_arr = solve(
                 internal_params=internal_params,
@@ -240,6 +332,101 @@ def solve(
             )
         return period_to_regime_to_V_arr
 
+    def _maybe_start_simulate_compile_async(
+        self,
+        *,
+        internal_params: InternalParams,
+        max_compilation_workers: int | None,
+        logger: logging.Logger,
+    ) -> None:
+        """Spawn `compile_all_simulate_functions` in a background thread.
+
+        Called from `solve(...)` so the simulate-side XLA compilation runs in
+        parallel with the GPU-bound backward induction. No-op when
+        `n_subjects is None`, when the cache for this size is already
+        populated, or when a compile is already in flight.
+        """
+        if self.n_subjects is None:
+            return
+        with self._simulate_compile_lock:
+            if self.n_subjects in self._simulate_compile_cache:
+                return
+            if self._simulate_compile_future is not None:
+                return
+            executor = ThreadPoolExecutor(
+                max_workers=1, thread_name_prefix="lcm-simulate-compile"
+            )
+            self._simulate_compile_future = executor.submit(
+                compile_all_simulate_functions,
+                internal_regimes=self.internal_regimes,
+                internal_params=internal_params,
+                ages=self.ages,
+                n_subjects=self.n_subjects,
+                max_compilation_workers=max_compilation_workers,
+                logger=logger,
+            )
+            executor.shutdown(wait=False)
+
+    def _resolve_simulate_internal_regimes(
+        self,
+        *,
+        actual_n_subjects: int,
+        internal_params: InternalParams,
+        log: logging.Logger,
+        max_compilation_workers: int | None,
+    ) -> MappingProxyType[RegimeName, InternalRegime]:
+        """Return internal_regimes to use for simulate; AOT cache when matching.
+
+        Three dispatch cases:
+
+        - `n_subjects is None`: return the original `internal_regimes`
+          (purely lazy path).
+        - `actual_n_subjects != n_subjects`: return the original
+          `internal_regimes` and log a warning the first time each
+          mismatching size is seen.
+        - `actual_n_subjects == n_subjects`: return the cached AOT-compiled
+          regimes. If `solve(...)` started a background compile, await it
+          here; otherwise compile synchronously.
+        """
+        if self.n_subjects is None:
+            return self.internal_regimes
+        if actual_n_subjects != self.n_subjects:
+            with self._simulate_compile_lock:
+                already_warned = actual_n_subjects in self._warned_n_subjects
+                if not already_warned:
+                    self._warned_n_subjects.add(actual_n_subjects)
+            if not already_warned:
+                log.warning(
+                    "simulate called with n_subjects=%d but model declared "
+                    "n_subjects=%d; falling back to runtime compile.",
+                    actual_n_subjects,
+                    self.n_subjects,
+                )
+            return self.internal_regimes
+        with self._simulate_compile_lock:
+            if self.n_subjects in self._simulate_compile_cache:
+                return self._simulate_compile_cache[self.n_subjects]
+            future = self._simulate_compile_future
+        if future is not None:
+            compiled = future.result()
+            with self._simulate_compile_lock:
+                self._simulate_compile_cache[self.n_subjects] = compiled
+                self._simulate_compile_future = None
+            return compiled
+        with self._simulate_compile_lock:
+            if self.n_subjects not in self._simulate_compile_cache:
+                self._simulate_compile_cache[self.n_subjects] = (
+                    compile_all_simulate_functions(
+                        internal_regimes=self.internal_regimes,
+                        internal_params=internal_params,
+                        ages=self.ages,
+                        n_subjects=self.n_subjects,
+                        max_compilation_workers=max_compilation_workers,
+                        logger=log,
+                    )
+                )
+            return self._simulate_compile_cache[self.n_subjects]
+
     def simulate(
         self,
         *,
@@ -341,10 +528,17 @@ def simulate(
                     )
                     exc.add_note(f"Snapshot saved to {snap_dir}")
                 raise
+        actual_n_subjects = len(next(iter(initial_conditions.values())))
+        simulate_internal_regimes = self._resolve_simulate_internal_regimes(
+            actual_n_subjects=actual_n_subjects,
+            internal_params=internal_params,
+            log=log,
+            max_compilation_workers=max_compilation_workers,
+        )
         result = simulate(
             internal_params=internal_params,
             initial_conditions=initial_conditions,
-            internal_regimes=self.internal_regimes,
+            internal_regimes=simulate_internal_regimes,
             regime_names_to_ids=self.regime_names_to_ids,
             logger=log,
             period_to_regime_to_V_arr=period_to_regime_to_V_arr,