KROME-X: a GPU-ready library for chemistry and microphysics in astrophysical simulations

Introduction

Chemistry plays a central role in many astrophysical environments: the formation of stars at low and high metallicity [106, 49], the interstellar medium [65, 146], starbursts and active galactic nuclei [91, 95], protoplanetary disks [127, 148], planetary atmospheres [18, 84] and the early Universe [24, 134, 43, 120]. It regulates gas cooling and thermodynamic evolution, and through it, the gravitational collapse [152, 107], the stability of disks [83], and their fragmentation (Li et al. 2003); at the same time, the abundances of atoms, ions, and molecules determine the line emission or absorption through which these objects are observed. An accurate and detailed model of gas thermodynamics and a meaningful comparison with observations, therefore, require chemistry in the numerical simulations themselves.

However, two difficulties stand in the way: i) chemical kinetics is expensive, as the rate equations are stiff, and a network of a few tens of species can dominate the cost of a hydrodynamical step, so that large simulations need both efficient solvers and reduced networks [52]; ii) building the system of ordinary differential equations (ODEs) of a given network, with its Jacobian and sparsity structure, is error-prone when the network has hundreds or thousands of reactions. Several codes addressed the second problem within their own frameworks [100, 77, 126, 75, 92, 144], but few were publicly available, and even fewer were designed to be embedded in a three-dimensional simulation (Richings et al. 2014b,a; Smith et al. 2017; Gong et al. 2023; Smith et al. 2023).

KROME [55] was originally written to fill this gap: for any network provided by the user, it generates the subroutines that solve the rate equations, together with the thermal processes tightly connected to the chemistry (cooling and heating by atoms, ions, molecules and dust, photochemistry, cosmic rays, dust evolution, grain-surface chemistry). The KROME pre-processor writes Fortran code that is called directly from the host simulation code, and integrates the rate equations with the sparse ‘backward differentiation formula’ (BDF) solver DLSODES [61], which was shown to be both accurate and efficient for astrochemical networks [12] Grassi et al. 2013). Since its release, KROME has been coupled to grid and particle codes and used to study different environments: i) the formation of the first stars and black holes (e.g., Bovino et al. 2014b,a; Van Borm et al. 2014; Katz et al. 2015; Bovino et al. 2016b; Riaz et al. 2018; Sharda et al. 2019; Díaz et al. 2024; Lenoble et al. 2024), ii) the chemistry and thermodynamics of the interstellar medium (ISM) of galaxies (e.g. [11, 19, 86, 85, 88, 131, 119]), iii) molecular clouds and star-forming filaments and cores ([14, 42, 87] (Seifried & Walch 2016)). Over time, the package has incorporated new modules, for instance, the tabulated dust cooling and H$_2$ formation on grains ([53]) or photo-bin-based X-ray secondary ionization and heating ([89]). Moreover, its networks have served to train machine-learning surrogates ([16, 68]) and have been adapted to the chemistry of dissipative dark matter ([26, 117]).

Since then, much of the microphysics of KROME models has been revised: the reaction databases (KIDA, UMIST; [144, 145, 94, 98]); the collisional data of the line coolants (LAMDA, BASECOL; [121, 141, 37, 38]), with quantum-mechanical fine-structure excitation of C, C$^{+}$, O, Si, and S by H and H$_2$ ([8, 82, 155 a] (Kłos et al. 2021)); the electron-impact data of CHIANTI ([29, 28, 40]); radiative and dielectronic recombination for whole isoelectronic sequences ([5, 4]); molecular photo-processes ([59]); the cosmological recombination with effective multi-level atoms ([122, 123, 151, 116, 1, 23]); and the optical properties of dust (Draine 2003a,b; Jones et al. 2017; Hensley & Draine 2023).

Chemistry libraries for hydrodynamical codes (e.g., [48, 111, 133]), gas-grain and photodissociation-region codes (e.g., [9, 74, 115, 41, 63, 22]) and disc models ([56]) are now widely used, and radiation-hydrodynamics codes transport the photons that drive photochemistry (e.g., [114] (Kannam et al. 2019)), so that the exchange of photons with the chemistry, and its conservation ([96]), is part of the problem. The coupling of non-equilibrium chemistry is a relevant topic, as demonstrated, for example, by Menon et al. [97], who integrate KROME into their radiation-hydrodynamics framework to model the interplay between radiation and chemical processes. Finally, most of the computing power of current supercomputers is in graphics processing units (GPUs): stiff kinetics has been ported to GPUs in combustion ([135, 103, 102]), supported by general solvers and portability layers ([20, 7, 44, 139]), but legacy KROME, a generated Fortran program with module-level state and DLSODES, is not suitable for those frameworks.

These developments motivated KROME-X, a rewrite of the package from scratch. It keeps what made KROME (“legacy” hereafter) useful (any network, the physics chosen at generation, the legacy networks, options, and tests) and changes how it is built. The network is described in a declarative language and checked before any code is written; the generated code is a compact C++ core linked to a fixed runtime library, integrated with the

BDF method of CVODE ([62, 44]), ported to C++, and automatically handles both analytic and finite-difference Jacobians; the same kernels run on CPUs and GPUs; and every physical process is validated against an independent implementation of its formulas and, where possible, against published data. The physics has been updated to reflect the data mentioned above and extended with new modules, including the co-evolution of the radiation field with chemistry and a consistent dust evolution.

The paper is organized as follows. Section 2 describes the restructuring, whereas section 3 presents the updated and new physics. Section 4 validates KROME-X against legacy KROME and compares their performance; Section 5 tests the new physical modules, and Section 6 compares the performance of CPUs and GPUs. We summarize our results in Section 7.

The restructuring

A code library

Legacy KROME consists of a Python pre-processor that reads a network file and a list of options, and fills a set of Fortran templates through pragmas: for every network it writes the system of ordinary differential equations (ODE system) to solve, the cooling, heating, dust and photochemistry modules, plus a user module, into a build directory. DLSODES files are copied along without any change. Every feature is an ad hoc branch of the templates, and the user extends the code by pasting Fortran code into it (custom ODEs, rate modifiers).

KROME-X separates what depends on the network from what does not (Figure 1). A code generator translates the network into C++: the species and their data, the per-cell parameters, the rate coefficients, the right-hand side, the analytic Jacobian (if used), the configured thermal processes, and the embedded tables. Unlike legacy KROME, KROME-X writes the table data in C++ files that are compiled together with the rest of the code. Everything else is a fixed, header-only runtime library: the integrator, the many-cell backends, the physics modules (photochemistry, dust, line coolants, shielding), the queries, and the interfaces. The generated project is self-contained and records the options used to create it, so it can be regenerated identically. To make the kernels compile as both CUDA and HIP device code, the code is written in C++17 in a C-like style (plain structures and arrays; no exceptions, virtual functions, or dynamic containers in the kernels); see § 2.6. HIP builds are compiled in the test suite but not explicitly run on AMD GPUs; the runs in this paper have been performed on NVIDIA GPUs only.

Vector (PDF)

The library is used through its C++ interface (the solver calls, for one cell or many, and custom queries: rates, fluxes, mean molecular weight, column densities, dust temperatures, etc.) or through a C interface, always built and vectorized over cells, on which optional Fortran and Python bindings give results bitwise identical to C++. A failed cell keeps its input and returns a named flag. The legacy tests are shipped as ports for comparison with legacy KROME. The option format and project setup have been modernized, and the migration table in the manual maps every legacy option to its replacement.

A revised network language

Networks are written in a declarative format (.knet) inspired by the one in PRIZMO [56, 125]: reactions are equations, rates are expressions in a small, backend-neutral language or typed constructors for the processes that need a runtime module (photo-processes, grain-surface processes, the line photodissociation of H$_2$), and network-level directives replace the raw Fortran of legacy KROME. The full description of the new format is described in Appendix A, with a simple comparison between the two formats reported in Fig. A.1.

The parser produces an intermediate representation with resolved symbols and typed rates, and checks the conservation of mass, charge, and number of nuclei in every reaction. In the case of reactions defined within a temperature range, legacy KROME set the rate to zero outside the range; KROME-X, instead, assumes by default an asymmetric rule, zero below the lowest bound and the value at the upper bound above it. This choice is motivated by the fact that extrapolating barrier reactions below their range can inflate them by orders of magnitude, whereas above the range, a constant continuation is safer. KROME-X has been coupled with a converter that automatically translates legacy networks to the new format on the fly: all 40 networks of legacy KROME convert, and every converted analytic rate was checked numerically against the legacy expression.

Generation and runtime options

As in legacy KROME, the physical processes to be included (cooling and heating, adiabatic index, dust, radiation, and the shielding models) are chosen when the code is generated. The configuration is a Python object, whose options, grouped by section (e.g. thermal.cooling), are checked when they are set, or an equivalent TOML file run from the command line; every generated project records it, so that it can be regenerated exactly. The chosen processes become compile-time constants in the generated code, so processes not configured are compiled out, and a configured process whose runtime module is missing results in a compilation error rather than a silent zero. At runtime, the only physics switch is the overall thermal evolution. The host code cannot turn individual processes on or off, which removes a whole class of inconsistencies between the generated code and its use. Everything else at run time is numerics (tolerances, Jacobian mode, linear solver, restarts, conservation corrections, backend) and the per-cell inputs: abundances, $T_{\mathrm{gas}}$, the time step, the parameters, and, if any, the photon fluxes and dust bins. The per-cell parameters replace the global setters of legacy KROME (krome_set_user_*): a network has exactly the parameters it uses, found by scanning the generated code and the configured processes, and those without a meaningful default (e.g., the free-fall time for compressional heating) must be set, or the solver refuses the call.

The core operation: code generation

After the first step that built the intermediate representation, the generator writes: the rate functions, split into those that depend only on $T_{\mathrm{gas}}$ and the parameters (evaluated once per call when $T_{\mathrm{gas}}$ is fixed) and the others; the right-hand side of the ODE system; and the Jacobian, differentiated symbolically from the rate expressions. The dependence of the rates on aggregate densities (e.g., $n_{\mathrm{tot}}$, the density of hydrogen nuclei $n_{\mathrm{H,tot}}$, $\rho$), which are all linear in the abundances ($q_k = \sum_j w_{kj}n_j$), is written as a low-rank term

where $J_{ij}$ is the Jacobian of the right-hand side $f_i = \mathrm{d}n_i/\mathrm{d}t$ (Eq. (2)), the first term is the derivative at fixed aggregate densities and the second their dependence on $n_j$ through the weights $w_{kj}$. This choice keeps the Jacobian exact and cheap, even for networks in which every surface rate depends on the total density (for instance, it reduces the number of evaluations of each rate in the Jacobian of the 655-species Semenov network from about 1300 to about 4). Functions that are not differentiated symbolically (tabulated or iterative rates) contribute through finite differences restricted to their known dependencies. Large tables (cooling and dust tables, cross sections) are embedded as constant arrays compiled once, so that the library reads no files at runtime, a requirement for the device backends.

To validate the code generator and prevent any silent errors, KROME-X includes a self-test suite that checks the generator on 31 networks: an independent Python evaluation of the intermediate representation reproduces the rates to $\leq5 \times10^{-10}$ (typically $10^{-13}$) and the right-hand sides exactly, and the analytic Jacobians agree with finite differences to their precision.

Numerical aspects

The right-hand side of the ODE is

where the first term is replicated for all the chemical species, with $n_i$ is the number density of the $i$-th species, $\nu_{ir}$ its net stoichiometric coefficient in reaction $r$, $k_r$ the rate coefficient of reaction $r$, a function of temperature and, for some reactions, other quantities, and $\mathcal{R}_r$ is the set of reactants of reaction $r$. $\zeta$ identifies the sign of the reaction, i.e., if the reaction forms or destroys the $i$-th species. In the second equation, $\gamma$ is the adiabatic index, $\Gamma$ and $\Lambda$ the heating and cooling functions, while $k_{\mathrm{B}}$ is the Boltzmann constant (see also Sect. 3 and App. D). These equations, together with the photon fluxes $\Phi_b$ for each photo-bin $b$, the dust temperatures $T_{\mathrm{d},s}$ and the grain sizes $a_{\mathrm{d},s}$, and derived from the state vector $(n_i, T_{\mathrm{gas}}, \ln\Phi_b, T_{\mathrm{d},s}, a_{\mathrm{d},s})$, are integrated with the variable-order BDF method of CVODE [62, 44] and a modified Newton iteration. The runtime does not link SUNDIALS explicitly: its integrator is a C++ port of the CVODE BDF method of SUNDIALS 7.9, one header that compiles for both host and device and keeps its work arrays with a stride, so that the same code runs on cells stored contiguously or interleaved (see § 2.6). On the benchmark problems of § 4.2, it reproduces CVODE bitwise, solver call by solver call (13 000 calls, stored as a test of the code). With $T_{\mathrm{gas}}$ fixed, an analytic Jacobian is used; when $T_{\mathrm{gas}}$ evolves, instead, the Jacobian is built by difference quotients by default (as legacy KROME). This choice is motivated by the fact that the thermal terms include tables and piecewise fits whose exact derivatives make the Newton iteration less robust. As an option, we also included a hybrid Jacobian with an analytic chemistry block and a finite-difference thermal row. With different quotients, the Newton iteration keeps the charge: the electron component of each correction follows from the charge balance, as it would with an exact Jacobian (the quotients of trace species leak charge at round-off, and in a collapsing cloud, where the electron fraction falls to $10^{-12}$, the leaked charge, compressed with the gas, can outgrow the electrons. Without the correction, 0.2-0.4 percent of cells in a collapse benchmark lost them on both the CPU and the GPU.

The linear systems are generally solved by dense LU factorization, either using the default CVODE implementation or LAPACK² [2 default above 64 equations]. In the case of very sparse networks, a sparse LU with static pivots is employed instead, as done by the DLSODES solver in legacy KROME. The generator orders species by minimum degree, computes the fill of the factors, and writes their structure as tables, which the integrator factorizes without pivoting (a zero pivot counts as a failed step), with $T_{\mathrm{gas}}$ and the other non-chemical equations eliminated as a dense border. The sparse LU in KROME-X is chosen automatically when it needs less than 10% of the floating-point operations of the dense LU. For instance, for the 655-species Semenov network, the factors fill 11% of the matrix, one factorization takes 73x fewer operations, and the whole run is 2.2x faster than with LAPACK and 16x faster than with the dense LU of CVODE, on one core. Unlike legacy KROME, where an error test or convergence failure simply re-entered the solver, where we restart the integration with a smaller step size. The default relative tolerances are tighter than legacy ($r_{\mathrm{tol}} = 10^{-8}$ against $10^{-4}$), whereas we employ the same absolute tolerance $a_{\mathrm{tol}} = 10^{-20}$.

GPU offloading

Although KROME and KROME-X can be used as stand-alone packages to evolve chemistry, their main advantage compared to other public libraries is in the coupling with hydrodynamic codes. In a simulation, the chemistry of many independent cells/particles, potentially distributed across a time-step hierarchy spanning orders of magnitude, must be solved at every time step. On CPU computers, calculations can be performed per cell, with cells evolved by a single CPU thread sequentially. On GPUs, such an approach is far from ideal due to their highly parallel nature and the additional cost (unless everything is already stored on the GPU) of transferring data to and from the GPU. For this reason, in KROME-X we incorporate two backends behind a single call, chosen at runtime based on the number of cells managed by each computing unit: a) per-cell independent solves, parallelized with OpenMP, the analog of legacy KROME; b) device mode, where the cells of a call are copied together to a Kokkos execution space [20, 139], and one kernel runs one thread per cell. Each thread integrates its cell with the same integrator as the CPU, over its own time step, with its own step sizes, order, Newton matrix, and error control, so that easy cells do not follow the hardest one. The integrator's arrays are stored either contiguously per cell or interleaved across cells (element $i$ of cell $c$ at $c+iN_{\mathrm{cell}}$, so that the threads of a warp read consecutive memory). Two kernels, one per layout, are chosen based on the number of cells in the call: contiguous is chosen if the cells fit within the device's resident threads, or interleaved above (see Appendix G). In the interleaved case, the right-hand side and the Jacobian operate on thread-local copies of the state and the cell's context, while in the contiguous one, they operate on the cell's own rows in global memory. The device buffers are kept between calls. A cell that fails on the device is resolved again on the host, so only cells that truly fail are reported as failures. A call larger than the device memory is split into chunks that fit, one after the other, in the same buffers (by default, within 80 percent of the free memory).

We also measured a batched integration of all the cells with SUNDIALS and MAGMA [7], in which every cell follows the steps of the stiffest one: faster on nearly identical cells, it is 2.5 times slower on $10^4$ cells that differ and more on larger batches (Appendix G). The code design was shaped by the need to make it device-ready. All generated and runtime kernels are free functions marked for host and device compilation, the per-cell inputs and the runtime modules (photochemistry cross sections, dust size bins) are passed as plain arrays and views rather than global state, and nothing in the right-hand side reads files, allocates memory or keeps history (an iterative dust temperature, for instance, is converged to round-off at every evaluation, so that the right-hand side is a smooth function of the state). The device integrates with the host's own right-hand side and Jacobian, the same functions compiled for the device, so that every configuration runs there.

Physics

KROME-X ports the physics of legacy KROME [55], and extends and updates it.³ Each process is checked in the test suite against an independent implementation of its formulas (typically to machine precision) and, where applicable, against published tables or code. The data comes from primary sources, which are stored with the code together with the scripts that convert them, so that every number can be traced and regenerated. A few legacy options whose physics have been superseded were removed (Appendix B), including the Galli & Palla [43] H$_2$ cooling, which was replaced by that of Glover & Abel [47], and the approximate adiabatic indices.

Radiation and photochemistry

The radiation field is the photon flux $\Phi_b$ (in cm$^{-2}$ s$^{-1}$ units) divided into $N_b$ energy bins, given per cell by the caller. The rate of a photo-process $r$ of species $i$ corresponds to the discretized integral $k_r = \sum_b \sigma_{r b}\,\bar{\sigma}_{r b}\,\Phi_b$, with $\bar{\sigma}_{r b}$ the cross section averaged over bin $b$ by quadrature and $s_r$ an optional shielding factor written in the network (e.g., the H$_2$ self-shielding of Draine & Bertoldi [34] or Wolcott-Green et al. [149], the default). Every spectrum (e.g., blackbody, Draine [32]-like interstellar field, power laws, the Haardt & Madau [57] background, tables) is integrated over each bin, $\Phi_b = (4\pi/h)\int_b J(E)\,\mathrm{d}\ln E$. In addition, analogously to legacy KROME, photoheating deposits the energy of each absorbed photon above the threshold. In legacy KROME the integral is computed by sampling both $J$ and $\sigma$ at the bin midpoint, which over-weights the bins containing an ionization threshold (with 10 bins, a 24% excess in the carbon photoionization rate of one of the legacy tests⁴), which is why an ad-hoc choice of the bins was required.

The radiation modeling in KROME-x has been modified to take into account a set of different scenarios:

• Co-evolved radiation: in legacy KROME the field was kept constant during a call, and a radiation-hydrodynamics code had to remove the absorbed photons afterward. When photochemistry is fast relative to the time-step, this operator splitting consumes many more photons than the field contains, likely over-ionizing species. In KROME-x, we introduce an option (radiation.evolve) that makes each bin (indexed $b$) part of the time-dependent variables, as $u_b = \ln\Phi_b$, integrating them alongside chemical species as

where $p_r$ is the number of reaction events per absorbed photon (1 for a photoionization; the dissociation probability per pumping, about 0.15, for the Solomon process).⁵ In logarithmic form, the absorption term does not depend on $\Phi_b$, preventing overshooting to negative values, and the coupling blocks of the Jacobian are analytic and smooth ($\partial\dot{n}/\partial u_b \propto k_r$, $\partial\dot{u}_b/\partial n = \mathrm{const}$). The field returned by the solver has lost exactly the photons that the chemistry used. § 5.1 shows the difference with an operator-split update. Without --evolve-rad, the same photon bookkeeping as in legacy KROME is available.

• Lyman–Werner rates and H$_2$ lines: rates may depend linearly on the mean intensity of the bins in the Lyman–Werner band. Such rates remove their photons only from those bins. The line photodissociation of H$_2$ (legacy kpd_H2) is a set of photo channels, one per level of the ground electronic state, with the Doppler broadening of the tabulated cross sections applied as an exact average over the line width (legacy widened each segment, increasing the integrated cross section). Legacy KROME attached the UV-pumping heat to a reaction chosen by a network tag (@next_Solomon, CHEM heating) or by its position (the first H$_2 \rightarrow$ H + H reaction, PHOTOAV heating). In the legacy KROME path, the latter is the direct continuum photodissociation or the cosmic-ray dissociation. KROME-x identifies the Solomon process from the rate itself, while the direct continuum heats the gas by the photon energy above threshold, as any other photo-reaction.

• Secondary ionizations: X-rays are photons in high-energy bins. Legacy’s tabulated X-ray rates were based on Shull & van Steenberg [130] and were valid only for pure atomic H and He mixtures (Latif et al. [76]). In KROME-x, they are replaced by a consistent treatment of photochemistry and a model of fast photoelectrons based on results of Dalgarno et al. [25], valid for H, H$_2$, and He mixtures. The model has been expressed in the functional form of Shull & van Steenberg [130] as described in Quadri et al. [109], and also incorporates the ionization and heating by fast photoelectrons for heavier species, whose rates are rescaled by the ratio of their electron-impact ionization cross sections to that of H. Compared to legacy KROME, which assumed a single photoelectron energy of 4.46 keV, the new version accounts for the energy distribution of photoelectrons more consistently (see § 5.2). The secondary ionization options can also be included when $T_{\rm gas}$ is held fixed.

Dust

• Grain temperature: for binned dust, the temperature of each bin is the solution of its radiative and collisional energy balance, converged at every evaluation of the right-hand side (legacy KROME advanced it by one iteration per call, from the previous value, which in an optically thick test gave 134 K instead of 317 K). With dust.evolve_Td, the bin temperatures become state variables with a true energy equation

with the grain heat capacity $C_s$ of Draine & Li [36], and the terms on the RHS that are respectively the radiative heating, the radiative cooling, and the collisional energy exchange with the gas (Grassi et al. [53]). The latter term can be positive or negative depending on the temperature of the surrounding gas. The dust processes in each bin (accretion, sublimation, and heat exchange with the gas) use their own temperatures, while the grain-surface chemical rates use the average over the bins, weighted by grain area.

• Dust evolution: with dust.evolve the grain radii $a_{\mathrm{d},s}$ of each bin $s$ evolve; the state is the number of key-element atoms per bin, so that gas and dust together conserve each element to round-off. The processes are accretion of gas-phase atoms, with the thermal speed of the accreting atom (legacy used the proton’s, 3.5 and 5.3 times too fast for C and Si); thermal sputtering, as the mass flux of the Nozawa et al. [104] yields; and sublimation as the removal of atomic layers (Grassi et al. [53]),

independent of the size, with $E_0 = 7.2\,\mathrm{eV}$ for carbon (Lenzuni et al. [79]) and 4.66 eV for silicates (Nozawa et al. [104]), and an optional curvature (Kelvin) correction. Legacy KROME removed one atom per grain at a time, a radius change scaling as $a_{\mathrm{d},s}^{-2}$, and removed evaporated bins inside the right-hand side.

• Optical properties: the absorption and scattering efficiencies of the binned grains are computed with a custom Mie code (after Bohren & Huffman [10]) from the dielectric functions that Draine [33] currently recommends, which are not published as per-size tables. Astronomical silicates and graphite in the “1/3–2/3” approximation with the size-dependent free-electron term of Draine & Lee [35], whose parameters have been fit to Draine’s two published radii. The code reproduces Draine’s 1993 tables, used by legacy KROME, from his 1993 dielectric functions to $1.5\times10^{-4}$ (median, the

precision of the tables). Fig. E.1 compares the old and new efficiencies, where they differ mostly in the far-infrared for graphite (free electrons) and above 13.6 eV (X-ray edges).

– Photoelectric heating: besides the Bakes & Tielens [6] heating of legacy KROME, the heating of Weingartner & Draine [147] and their grain recombination cooling are available, with all 22 dust models and radiation fields of their tables, scaled with the dust-to-gas ratio of the cell. For Milky Way dust in the interstellar radiation field, it heats 1.4 times more than Bakes & Tielens [6] at $G_{0}T_{\mathrm{gas}}^{1/2}/n_{e}=10^{3}$ and 100 K.

– Grain-surface chemistry: surface rates are computed on the effective grain population or area-weighted over the runtime size bins (or, as in legacy, with species cloned per bin), with the Hasegawa et al. [58] frequency, thermal hopping and tunnelling, the competition between reaction and diffusion, cosmic-ray desorption, chemisorption, and reactive desorption with fixed efficiencies or those of Minissale et al. [99]. A dedicated comparison of a 655-species network with legacy KROME identified and corrected three causes of disagreement in the legacy implementation (a lost single-size grain population, binding energies matched case-sensitively, and the tunneling).

Line cooling and atomic processes

For metal-line cooling, we solve for the atom- and ion-level populations at each evaluation, as in the legacy code. Every level connected to the ground state is included. However, we added three extensions:

– optically thick lines: to account for optical thickness locally, with thermal.escape = "LVG" the radiative rates of each line are multiplied by the Sobolev escape probability $\beta=(1-e^{-\tau})/\tau$, with $\tau$ from the column of the species and the thermal and turbulent line width, and the level populations and escape probabilities iterated to self-consistency starting from the optically thin solution, as in RADEX [142], with under-relaxation of the populations.

– radiative pumping by the photon bins (now with the absorption in the level balance) and by the Cosmic Microwave Background (CMB) at the redshift of the cell, important at high redshift (the C$^{+}$ upper level has an occupation of 0.05 at $z=10$ and 0.6 at $z=30$ from the CMB alone, see Lupi et al. [88]);

– user coolant files in the same format used in legacy, converted from LAMDA or CHIANTI by the included converters, are added to the embedded coolants (without necessarily replacing them).

The coolant data were updated (Table E.1). The fine-structure excitation by H of C, Si, and S follows Yan & Babb [154], of C$^{+}$ and Si$^{+}$ Barinovs et al. [8], and that of O Lique et al. [82], in the corrected version distributed by LAMDA [153]. The excitation of C$^{+}$ and O by H$_{2}$, absent in legacy KROME, follows Kłos et al. (2021) and Lique et al. [82]. The electron excitation of O$^{+}$, Si$^{+}$, Si, and Fe$^{+}$ comes from CHIANTI [136, 137, 156, 39], and the radiative data from the NIST Atomic Spectra Database [73]. S I is a new coolant, and a 16-level model of Fe$^{+}$ is available as an optional file (the default uses only 5 levels). The molecular cooling tables of CO, H$_{2}$O, OH, and HCN [101, 105] are kept.

As for atomic processes, the radiative and dielectronic recombination rates follow Badnell et al. [5]; Badnell [4] for 370 ions. The recombination cooling is computed using the network’s own recombination reactions, and cosmological recombination is computed using RECFAST [122]. Further details on atomic chemistry, together with the modeling of reverse reactions, can be found in Appendix E.

Finally, all physical constants have been updated to values from the CODATA 2018 table, since the legacy version did not consistently use the number of significant digits.

Code validation and optimisation

Validation

As a validation check, we now compare KROME-X with the legacy KROME using the original version’s physics: the same rates, cooling and heating functions, and data, with only minor corrections made during the restructuring⁶. These runs were carried out on an Apple MacBook Pro equipped with M1, legacy KROME compiled with gfortran 12.1 at -O3, KROME-X with clang 21 at -O3. The outputs and the scripts that produced and compared them are archived with the code to ensure reproducibility. Two regressions run both codes on the same problem at the same, tight tolerance ($r_{\mathrm{tol}}=10^{-8}$, $a_{\mathrm{tol}}=10^{-20}$), to separate the equations from the solvers. Both evolve one cell of primordial gas with deuterium (the legacy network react_primordial) at constant density, from $n_{\mathrm{H}}=10^{4}\ \mathrm{cm}^{-3}$, $n_{\mathrm{H}}=n_{\mathrm{e}}=1\ \mathrm{cm}^{-3}$, $n_{\mathrm{H}}=775\ \mathrm{cm}^{-3}$ and $n_{\mathrm{D}}=0.253\ \mathrm{cm}^{-3}$, with outputs at $t=10^{6},10^{7},\ldots,10^{15}\ \mathrm{s}$, each reached by one call of the solver from the previous one. In the first, the chemistry alone at fixed temperatures of $10^{2}$, $10^{3}$, $5\times10^{3}$, $10^{4}$, and $3\times10^{4}$ K, the abundances of the two codes agree to $5$-$7\times10^{-8}$, and to $9\times10^{-6}$ for HD at late times. In the second, the same gas cools with atomic and H$_{2}$ cooling from $T_{\mathrm{gas}}=3\times10^{3}$, $10^{4}$ and $5\times10^{4}$ K. Since $T_{\mathrm{gas}}$ agrees to $4\times10^{-7}$ and the abundances to $\leq1.3\times10^{-5}$, the two codes therefore integrate the same equations to the precision of the solvers.

We then ran all the test problems available in legacy KROME that use ported features (21) with both codes, each at its own default tolerance. In Figure 2, we report the results for three representative tests: the one-zone collapse of primordial gas [55], labeled Collapse, of gas at metallicities between zero and 0.1 $Z_{\odot}$ [106], Collapse–Metal, and a dark cloud with gas-grain chemistry (655 species, DarkCloud). The curves of the two codes are indistinguishable: $T_{\mathrm{gas}}$ differs by $5\times10^{-4}$ (median) in the primordial collapse and by $4\times10^{-4}$ (median, 2% at most) along the low-metallicity tracks, and the main gas and ice species of the dark cloud by $\leq1\%$ over $10^{9}$ yr. The full comparison of the 21 tests is presented in Table B.1 in Appendix B. Most differences are small and come from the adopted tolerances (in collapseUV, for instance, legacy KROME itself changes by 2% between $r_{\mathrm{tol}}=10^{-4}$ and $10^{-8}$). The few larger ones come from defects of legacy KROME corrected here, most notably the sampling of the photon bins at their midpoint, which overestimated a photoionization rate by 24% with 10 bins (but no planned binning strategy), the energy of the photoelectrons in the secondary ionization, and the accretion and sublimation rates of dust.

Vector (PDF)

With the updated physics of Section 3, instead, the same test problems change where the updated data and rates apply (see Appendix F). In particular, they lower the temperature of collapsing low-metallicity gas by up to 43% at low density, raise the residual ionization after cosmological recombination by 25%, and the H$_2$ and HD formed in the early Universe by factors 1.7 and 2.5, and change the ionization balance and cooling of heavy elements by tens of percent. Elsewhere, they affect the results only through the physical constants, by $\lesssim10^{-5}$.

Computational performance compared to legacy KROME

We now compare KROME-X with legacy KROME in terms of computational efficiency. For the new version, we employ the updated physics described in Section 3, the custom port of the CVODE integrator, and the linear solver it chooses by default (i.e., sparse LU for large species, in this case, the 655-species network). Note that the physical processes employed for each test match exactly those of legacy KROME, i.e., we use the same options and the same exact definition of the adiabatic index (computed from the composition and the partition functions), while the new modules introduced in KROME-X are not employed. The updated physics alone marginally affects the integration time of KROME-X: with the same integrator (CVODE) and the cases of Table 1, we measure variations between $-2\%$ and $+8\%$ ($+8\%$ for Collapse, $+5\%$ for Collapse-Metals, $+3\%$ for DarkCloud and $-2\%$ to $+4\%$ for Photo), with at most 1% more steps. The updated physics makes the right-hand side slightly more expensive, not the problems stiffer. The comparison, therefore, does not favor KROME-X.

Name	$N_{\mathrm{cells}}$	KROME	KROME-X	Ratio
Collapse	1	0.402 s	0.0986 s	4.1
Collapse-Metals	1	1.03 s	0.401 s	2.6
DarkCloud	1	28.6 s	2.12 s	13.5
Photo, 10 bins	30	14.9 ms	4.19 ms	3.6
Photo, 100 bins	30	15.4 ms	4.25 ms	3.6
Photo, 1000 bins	30	22.2 ms	5.04 ms	4.4
		Integrator steps		RHS calls
		KROME	KROME-X	KROME / KROME-X
Collapse		87 378	54 547	254 834 / 116 549
Collapse-Metals		85 791	71 573	334 939 / 224 225
DarkCloud		6 881	5 794	80 871 / 7 834
Photo, 100 bins		25 704	29 980	136 867 / 46 533

Table 1 and Figure 3 present the solver timings for four cases with matched tolerances, initial states, and output times. Collapse and Collapse-Metals correspond to the legacy tests collapse and collapseZ, DarkCloud corresponds to the cloudSurface, while Photo is a one-dimensional photochemistry problem of 30 cells with 10, 100, and 1000 photon bins (known as legacy rtchem). We report the total solver time per cell, i.e., the solver time for the whole case divided by the number of cells evolved during the test. Both codes make the same calls to the solver (i.e., the same output times). KROME-X is 2.6 to 13.5 times faster. The comparison shares equal tolerance and, as Appendix C shows, at equal accuracy. Against each code’s own converged solution, the two make errors of the same size at the same tolerance, and KROME-X is faster at every tolerance we tested, from $10^{-4}$ to $10^{-8}$. At its default $r_{\mathrm{tol}} = 10^{-8}$, twice as expensive as at $10^{-4}$ and three orders of magnitude more accurate, it is still faster than legacy KROME at the default $10^{-4}$ of the latter.

Vector (PDF)

The two integrators take a similar number of steps, with a smaller ratio than the timing ratio, suggesting cheaper right-hand-side evaluations and more efficient linear algebra. DLSODES builds its Jacobian by finite differences, while KROME-X uses an analytic form at fixed $T_{\mathrm{gas}}$ and, for evolving $T_{\mathrm{gas}}$, difference quotients as well, and re-evaluates the Jacobian 3.4–3.9$\times$ less often (thanks to the reuse heuristics of CVODE), resulting in 1.5 (Collapse-Metals) to 10 (DarkCloud) times lower number of evaluations of the right-hand side. Each evaluation is also cheaper, since at fixed $T_{\mathrm{gas}}$ the temperature-dependent rates and the bin sums of a static photon field are computed once per call, integer powers are expanded as products by the code generator, and the production build vectorizes the bin loops. As a consequence, the cost of the photochemistry hardly depends on the number of photon bins (+20% from 10 to 1000 bins, against +49% for legacy KROME). For DarkCloud, the cost is the factorization of its 655$\times$655 Newton matrix, which is performed about 1750 times during the run, whereas legacy KROME factorizes it as a sparse matrix. As a dense matrix, KROME-X takes 35 s with the default unblocked LU of the integrator, 1.2$\times$ legacy KROME, despite 10$\times$ fewer evaluations of the right-hand side. With the blocked LU of LAPACK, it takes instead 4.8 s, and with the sparse

LU, its default for this large network, 2.1 s, 13.5x less than legacy KROME. Rate tables, commonly employed in legacy KROME to accelerate the right-hand-side evaluations (option -useTabs), were implemented and measured in KROME-X as well, but were not retained as i) they saved only 7–10% of the solver time, and ii) they moved the final temperature of the low-metallicity collapse by 9% (also in legacy KROME).

The new physical modules

This section tests the new modules of Section 3, which the legacy test problems do not use.

Co-evolved photon bins

This test isolates how the photon budget is managed during a thermochemical step. When a cell holds fewer ionizing photons than atoms, the number of ionizations must be limited by the photons available, but a field kept fixed during a step $\Delta t$ (legacy KROME, or any radiation-hydrodynamics code that does not co-evolve the field) keeps ionizing at the initial rate and consumes photons it does not have, unless $\Delta t \ll t_{\mathrm{ion}}$. Since $t_{\mathrm{ion}} = 1/\Gamma$ is typically much shorter than a hydrodynamical time-step, in ionization fronts this condition is rarely met. We therefore use a cell of neutral hydrogen ($n_{\mathrm{H}} = 1\ \mathrm{cm}^{-3}$, $T_{\mathrm{gas}} = 10^{4}\ \mathrm{K}$) with a known photon budget (0.5 ionizing photons per H atom) and no source, where the exact answer is known in advance: an ionized fraction equal to the photons per atom, minus the recombinations. The photon spectrum is sampled in 20 logarithmic bins between 13.6 and 100 eV with the shape of a $10^{5}\ \mathrm{K}$ blackbody. We integrate it to 100 photoionization times $t_{\mathrm{ion}} = 1/\Gamma_{0}$ ($4.1 \times10^{7}\ \mathrm{s}$) in three ways (Fig. 4): (i) with the field co-evolved by the solver, (ii) with the field held fixed during steps of length $\Delta t$ and attenuated after each step by the absorption coefficient of the state at the start of the step (the operator split of a radiation-hydrodynamics code that does not co-evolve the field), and (iii) with a constant field (legacy KROME within a single call). The co-evolved field ionizes exactly the photons it holds. The final ionized fraction is 0.499 (the 0.5 photons per atom, minus the recombinations), and photons plus ions are conserved at all times. The operator splitting converges to this solution only for $\Delta t \ll t_{\mathrm{ion}}$, giving $x_{\mathrm{HII}} = 0.533$ for $\Delta t = 0.1\,t_{\mathrm{ion}}$, 0.81 for $\Delta t = t_{\mathrm{ion}}$ and 0.999 for $10\,t_{\mathrm{ion}}$, creating up to one ion per atom out of half a photon. The co-evolved field also hardens, as the bins near threshold, with the largest cross sections, are consumed first (Fig. 4c), which a single attenuation per step cannot follow within the step.

X-ray secondary ionization

We now consider the collapse of primordial gas irradiated by a UV field ($J_{21} = 0$ to 3000) and by X-rays ($J_{21,\mathrm{X}} = 0.1$, following a power law spectrum of index 1.5 between 2 and 10 keV, attenuated over the Jeans length), which allows us to test photochemistry of the X-ray bins and the secondary ionizations/heating (Fig. 5). The legacy tables of the X-ray rates are reproduced by the bins of KROME-X within 0.01 dex over $N_{\mathrm{H}} = 10^{16}$–$10^{25}\ \mathrm{cm}^{-2}$, so that the difference lies in the secondary ionizations. With the energy of each photoelectron, it is about half that of legacy KROME ($3.1 \times10^{-18}$ instead of $5.9 \times10^{-18}\ \mathrm{s}^{-1}$ per H at $x_{\mathrm{e}} = 10^{-3}$ in the optically thin limit), the heating efficiency differs by −11% to +16%, and fewer electrons form H$_2$: the tracks with $J_{21} = 0$ and 10 are 5%–15% warmer, and the gas with $J_{21} = 1000$ starts to cool by H$_2$ at $n \approx870\ \mathrm{cm}^{-3}$ instead of $\approx450\ \mathrm{cm}^{-3}$.

Dust evolution

The legacy dust test (DustEvol in Sect. 6) evolves gas-phase carbon and silicon with [93 MRN] distributions of carbon and silicate grains (10 bins per type, $5 \times10^{-7}$–$2.5 \times10^{-5}$ cm, dust-to-gas ratio $10^{-2}$) for $10^{3}$ yr, from the same initial state, at 100 fixed temperatures between 10 and $10^{4}$ K, with the dust temperature equal to $T_{\mathrm{gas}}$ (Fig. 6). Both codes show the same behavior: depletion by accretion at intermediate temperatures, and the return of the elements to the gas when the grains sublimate, with identical gas-phase values once all grains have sublimated. The differences come from the two corrections of Sect. 3.2, each visible in its own temperature range. At intermediate temperatures, the accretion rate scales with the thermal speed of the accreting atom, which legacy KROME took as the proton’s, $\sqrt{m/m_{\mathrm{p}}} = 3.5$ (C) and 5.3 (Si) times too large; since the depletion over $10^{3}$ yr is nearly exponential, this factor enters the exponent: at 327 K, the gas-phase carbon drops by a factor of $\sim10$ in KROME-X and $\sim5 \times10^{3}$ in legacy KROME ($1.9 \times10^{-22}$ and $3.7 \times10^{-25}\ \mathrm{g\ cm}^{-3}$). At high temperatures, the layer-removal sublimation and the binding energies of Sect. 3.2 return the elements to the gas at lower temperatures, by $\sim1300$ K for silicates and $\sim2600$ K for carbon, where legacy KROME still retains part of the dust. Gas plus dust conserve each element to round off in KROME-X.

Vector (PDF)

Vector (PDF)

Vector (PDF)

GPU versus CPU performance

Unlike the previous tests, for the comparison between CPUs and GPUs we employed a GPU node on Leonardo (CINECA, Italy), equipped with one Intel Xeon Platinum 8358 (32 cores, 2.6 GHz) and four NVIDIA A100 (SXM, 64 GB), of which only one was used. The CPU code was compiled using the Intel compiler icpx 2024.1 (-O3 -xHost -qopenmp -fp-model=precise, with mkl 2024.0 when lapack is enabled), whereas the device code was compiled with cuda 12.2 and Kokkos 4.7.1, both with the physics of § 3. The same runs on one core with gcc 12.2 (-O3 -march=native) show the effect of the compiler.

Setup

For this comparison, we consider three cases that together bracket the typical configurations of numerical simulations: (i) Collapse-Metals of Table 1 (28 equations), the thermal evolution with a metal-line network; (ii) Photo-RT, the Photo case with 10 bins co-evolved with the chemistry (44 equations); and (iii) DustEvol, the test of Sect. 5.3 with 10 silicate and 10 carbon bins evolving in size and temperature (55 equations). The large network is DarkCloud (655 equations, 74 calls, heterogeneous cells only).

Each case first runs its test problem on the host and records the inputs of every solver call (abundances, $T_{\mathrm{gas}}$, time step, parameters, photon fluxes, dust temperatures, and grain sizes). Then, we consider two opposite cases, (i) an heterogeneous batch, in which the $N$ cells then sample this history evenly (cell $c$ takes call $\lfloor cK/N\rfloor$ of the $K$ recorded), so that one call of the many-cell solver contains the whole history of the problem at any $N$, from the easiest to the stiffest steps, with neighbouring cells at neighbouring stages, as the cells of a simulation; (ii) a uniform case, in which all cells take the same call. These two approaches allow us to isolate the cost of the divergence between cells. The time per cell is estimated as the shortest wall time of the call (divided by $N$) among five identical runs (three for DustEvol): on one CPU core ($N=5000$, which covers every recorded call, 100 for DustEvol), on the full node (the per-cell backend, OpenMP over the cells, one thread per core, $N=10^4$ and $10^5$) and on the GPU ($N=10^3$ to $10^6$, $10^5$ for DustEvol). The final states of the GPU and of the node are compared cell by cell.

Results

Our results are reported in Table 2 and Figure 7. To make a fair comparison, we report the results for one A100 GPU with its share of the node’s CPU cores, 8 cores. The time per cell of the CPU does not depend on $N$, as the cells are evolved in sequence by every thread separately; that of the GPU falls instead by 4–15 times from $10^{3}$ to $10^{4}$ cells, as the device fills, and the uniform batches saturate at $10^{5}$ cells while the heterogeneous ones still gain about 40 percent up to $10^{6}$. On heterogeneous cells, one A100 outperforms its eight cores across a few thousand cells: at $10^{4}$ cells, it has 2.0, 1.3, and 5.7 times their throughput for Collapse-Metals, Photo-RT, and DustEvol, and at the largest $N$, 10, 6.5, and 33 times. DarkCloud, the large network, is the exception: there, the A100 reaches only 0.61 of the throughput of its eight cores, at the largest batch it can hold (see below). The photochemistry gains the least, its time being in the dense factorization of 44 equations per thread (Table G.1), and DustEvol, the most expensive case per cell on the CPU, the most. The heterogeneous cells cost 1.7–2.7 times as much as the uniform ones on the GPU at $10^{5}$ cells, compared to 1.03–1.8 on the CPU, where only the per-cell work differs: the excess is the price of divergence within a warp. No cell failed.

case	cells	1 core	8 cores	A100, $10^{4}$	A100, $10^{5}$	A100, $10^{6}$	A100 / 8 cores $10^{4}$	max
Collapse-Metals	H	127	18.4	9.08	2.84	1.81	2.0	10.2
	U	123	18.0	3.54	1.64	1.68	5.1	11.0
Photo-RT	H	323	43.1	34.0	10.7	6.65	1.3	6.5
	U	222	30.3	11.2	5.46	5.42	2.7	5.6
DustEvol	H	$2.1 \times10^{5}$	$2.9 \times10^{4}$	4987	864	–	5.7	33.1
	U	$1.2 \times10^{5}$	$1.6 \times10^{4}$	642	323	–	24.9	49.5
DarkCloud	H	$3.8 \times10^{4}$	5956	$9720^{a}$	–	–	0.61	0.61

Vector (PDF)

The comparison depends on the CPU code compiler: on one core, the Intel build is about twice as fast as the GNU build on collapses and DustEvol, and up to 1.5 times slower on Photo-RT.

Networks of hundreds of species were measured with DarkCloud. With dense matrices (7 MB per cell), a GPU can hold too few cells to fill it (16 s per cell and call at 100 cells). With the sparse LU (§ 2.5) the A100 is 23–30 times faster, and its time per cell falls as $1/N$ up to the largest batch it holds (4 MB per cell, $\sim1.5 \times10^{4}$ cells): 13 ms per cell and call at 8000 cells and 9.7 ms at 12 000 (Table 2 and the last panel of Figure 7), against 6.0 ms for eight cores. The GPU memory, not the arithmetic, sets the limit: the time per cell is still falling at the largest batch (as $N^{-0.7}$), and about twice as many cells, $\sim2.5 \times10^{4}$, would be needed to match the eight cores. Networks of hundreds of species, such as the complete KIDA or the DarkCloud ones, are therefore still best solved on the CPU.

In terms of accuracy, instead, the GPU and the CPU agree to $2 \times10^{-3}$–$4 \times10^{-3}$ (the largest relative difference of the final states over the heterogeneous batches, which include the chaotic calls of Collapse-Metals; DarkCloud $2 \times10^{-3}$) and to $10^{-11}$–$1.4 \times10^{-3}$ on the uniform ones. For DustEvol, the dust temperatures and grain radii of the bins that still hold dust agree to $1.3 \times10^{-3}$ and $6 \times10^{-3}$.

Conclusions

In this work, we have introduced and validated KROME-x, a full rewrite of the KROME package for chemistry and microphysics in astrophysical simulations. Our main results can be summarised as follows:

• A real library A code generator turns a declarative, validated network description into a compact C++ core with an analytic/numerical Jacobian, linked to a fixed runtime that integrates it with a C++ port of the BDF method of CVODE, identical to it bit for bit, which needs no external library. The physics is fixed when the code is generated. At runtime, the host code selects only the numerics and the per-cell inputs. A C interface, with optional Fortran and Python bindings, serves the host codes, and all interfaces give bitwise identical results.

– A GPU-ready tool. The kernels and the integrator are written so that they compile as device code. On a kokkos device, each thread integrates one cell, with its own time step, order, and error control, and every configuration of the host runs there, including co-evolved photon bins and evolving dust. One A100 outperforms eight CPU cores on its share of a node on the Leonardo machine (Italy), increasing throughput from a few thousand cells per call to 5-30 times that at $10^{5}$–$10^{6}$ cells (§ 6). The per-thread kernel wins on networks of tens of species; networks of hundreds of species currently remain faster on the CPU, even with sparse LU on the GPU, due to the GPU’s memory limit.

– Accurate and faster. With the physics of legacy KROME, KROME-X reproduces it to the precision of the solvers in multiple test problems, with the largest differences traced to slightly discrepant numerical choices and small inconsistencies in legacy KROME. With the updated physics, it is 2.6-13.5 times faster on the benchmark cases at equal tolerance and accuracy.

– New physics. The photon field can be co-evolved with the chemistry, which conserves photons exactly; the secondary ionizations follow the energy of each photoelectron ([109 appendix B]); dust grains can grow, be sputtered, and sublimate with element conservation, and their temperature can follow a true energy equation. Line coolants can be optically thick and pumped by the photon bins and the CMB.

– Updated data. The line coolants are rebuilt from current collision and radiative data, while recombination follows [5, 4], the cosmological recombination RECFAST, the dust optical properties [33], and the constants CODATA 2018. On standard problems, these updates change the final results by up to a factor of 2.5, especially at low or zero metallicity.

Every process is validated in an automated test suite against independent implementations and published references, and every data set is stored with its source and the script used to convert it. Future work includes the CPU and GPU performance on production problems, an updated UV background and dust tables, and a tool to reduce networks for a given application.

Author contribution

The motivation for this work arose from discussions among the authors at the Como Lake School for Advanced Studies, "A decade of KROME," which the authors organized in July 2026. AL took care of most of the restructuring of KROME, the tests, and the paper draft. All authors equally contributed to the editing and revision of the manuscript.

Disclaimer about the use of Artificial Intelligence in the present work

Given the extent of the restructuring of KROME, we worked with the AI coding assistant Claude Code (Anthropic, Claude Opus models) under the authors’ direction in an iterative, supervised process. The authors of the original code (TG and SB) are actively involved in the process. All the authors defined the design, physics, data, reference implementations, and tests; the assistant wrote most of the source code, the scripts for the data tables, benchmarks, and figures. Every result was checked by the authors against independent references (legacy KROME, the original data sources, the literature, or independent implementations). The GPU and cluster runs were performed by the authors, who take full responsibility for the code, the results, and the text.

Data availability

KROME-X is available at https://github.com/alupi24/krome-x. The data and scripts behind the figures in this paper, as well as the comparison with legacy KROME the timing and benchmarking of the physics modules, are available upon request.

The network language

Networks are written in a declarative format (.knet), see Fig. A.1 for an example. Reactions are equations with case-sensitive species, e- for the electron, isotopes, isomers, and state labels. Rates are expressed in a small, backend-neutral language (with variables, tables, piecewise temperature ranges, and a lazy conditional), or typed constructors for processes that need a runtime module. In particular, photo() declares a photo-process (with the cross section from the database, a file, or an expression, an optional shielding factor, and the number of reaction events per absorbed photon) and surface() a grain-surface process; the line photodissociation of H$_2$ and rates proportional to the Lyman–Werner intensity are also supported. Reverse reactions are generated by detailed balance (<->), and two network-level directives replace the raw FORTRAN of legacy KROME: a global rate scaling and the freezing of species under a condition on the parameters.

Converting networks and coolant data. The distribution includes converters to the .knet format of the KIDA networks (Wakelam et al. [144, 145]) and of the UMIST ratefiles (McElroy et al. [94]; Millar et al. [98]). They implement the rate formulae of each database (modified Arrhenius, the ion-dipole capture formulae of KIDA, cosmic-ray ionization and cosmic-ray-induced photoreactions scaled by the ionization rate, and photoreactions attenuated by the visual extinction), and write rates fitted over several temperature intervals as piecewise rates. Pseudo-species (cosmic rays, photons) are removed from the reactions and enter through their rates, and duplicated entries are kept once. Parameters inconsistent with their formula are flagged. Two more converters write the coolant files of Sect. 3.3 from the LAMDA (Schöier et al. [121]) and CHIANTI (Dere et al. [29]; Del Zanna et al. [28]) databases.

Documenting a network. A separate tool, which replaces the legacy docmake, documents a network from its .knet file without generating or compiling it. It evaluates the rate coefficients from the network expressions and plots them as functions of the temperature and of the parameters, including the behavior outside the declared temperature ranges, and reports the jumps of piecewise rates at their boundaries. For photo-processes, it shows the cross section used by the code and the unshielded rate in the Draine [32] field. It also lists candidate missing reactions and branches (reactant sets the network does not use, and products with the same elements and charge it does not include), ranked by the reaction enthalpy from the Burcat & Rusciс [17] database, draws the network graph of each element, and writes the reaction table in LaTeX.

The legacy test problems

Table B.1 lists the 21 legacy test problems run in both codes for Section 4 (KROME-X with the physics of legacy KROME, at its default $r_{\mathrm{tol}} = 10^{-8}$; legacy KROME at its default $10^{-4}$; both at $a_{\mathrm{tol}} = 10^{-20}$, and both at the tolerance a test sets: $r_{\mathrm{tol}} = 10^{-6}$ for CloudSurface, $a_{\mathrm{tol}} = 10^{-40}$ for collapseZ and collapseCO), with the relative difference in $T_{\mathrm{gas}}$ and the cause of the differences $\gtrsim10^{-2}$. Each test’s documentation in the code quantifies them, mostly by patching the legacy build to isolate each cause.

Legacy options not carried over. Of the 107 options of the legacy pre-processor, the following have no counterpart in KROME-X; the migration table of the manual gives the replacement of every other one:

• superseded physics: the Galli & Palla [43] H$_2$ cooling (-cooling=H2GP98), replaced by that of Glover & Abel [47]; the approximations of the adiabatic index that freeze a degree of freedom or leave a species out (-gamma=FULL, VIB, ROT, REDUCED), replaced by the exact computation from the composition and the partition functions (the constant and the Popovas & Jørgensen 2016 options remain); and the switch-off of the cooling below a temperature (-cooling=quenсh), replaced by cooling floors;

• a cooling function not ported: the photoionized-gas cooling of Gnedin & Hollon [50] (-cooling=GH);

• user FORTRAN in the generated code (-customODE, -useODEConstant, -useCustomCoe, and the @ODModifier blocks of the network files), whose uses are covered by the directives of the network language (species frozen under a condition, a scale factor on all the rates, rates written in the network);

• rate tables (-useTabs; Section 4.2);

• tools outside the solver: the automatic networks from the reaction databases (-useAutoNetwork, -listAutomatics, -listSWRI), user cross-section kernels (-execKernelFunction), and the network reducer (-reducer), that will be replaced by a separate tool that writes and validates a reduced network.

The chemistry is always evolved (-useChemoToggle is not ported: the thermal evolution is the only physics switch at run time).

Accuracy against cost

Figure C.1 compares legacy KROME and KROME-X at a grid of relative tolerances on Collapse and DarkCloud, on one core of the Apple MacBook Pro M1, with the current codes: KROME-X with the physics of Section 3, its integrator and, for DarkCloud, the sparse LU; legacy KROME regenerated at each tolerance (the codes and options of Table 1, three runs per point). The error of a run is the largest relative difference of its outputs. $(T_{\mathrm{gas}}$, $x_{\mathrm{H}_2}$ and $x_{\mathrm{e}}$ for Collapse; gas-phase and ice CO and H$_2$O and the electrons, above $10^{-10}$, for DarkCloud), median over the outputs, against a converged run of the same code: KROME-X at $r_{\mathrm{tol}}=10^{-12}$, legacy KROME at $10^{-8}$, since it does not integrate Collapse at tighter tolerances (at $10^{-9}$ and $10^{-10}$ it stops at its limit of 100 restarts). Against its own $10^{-8}$ run, the errors of KROME-X at $10^{-4}$–$10^{-6}$ are the same within 30 per cent.

At the same tolerance, the two codes make errors of the same size, the medians within a factor of 2 and the worst outputs of KROME-X smaller or similar (Collapse at $10^{-4}$: $5.7 \times10^{-3}$ against $9.3 \times10^{-2}$), so that the speed-ups of Table 1 hold at equal accuracy: KROME-X is 2.0–3.5 times faster on Collapse and 12–18 times on DarkCloud at every tolerance from $10^{-4}$ to $10^{-8}$. Its cost grows slowly as the tolerance tightens: from $10^{-4}$ to its default $10^{-8}$, by factors of 2.3 on Collapse and 2.0 on DarkCloud, for errors that are 5000 and 3000 times smaller, respectively. With the defaults of the two codes, KROME-X ($10^{-8}$) is still 1.5 and 9 times faster than legacy KROME ($10^{-4}$).

test	$\Delta T_{\mathrm{gas}}/T_{\mathrm{gas}}$	cause
hello, slowmanifold, reverse	–	(agree to $\le5 \times10^{-3}$)
customCooling	$4 \times10^{-9}$ / $3 \times10^{-8}$	–
cooling_molecules	$8 \times10^{-5}$ / $1.4 \times10^{-3}$	–
collapse	$5 \times10^{-4}$ / $4 \times10^{-3}$	legacy tolerance
collapseZ	$4 \times10^{-4}$ / $2 \times10^{-2}$	legacy tolerance
collapseUV	$2 \times10^{-4}$ / $2 \times10^{-2}$	legacy tolerance
collapseUV_Xrays	$7 \times10^{-2}$ / 0.77	photoelectron energy (legacy: a fixed 4.46 keV)
collapseCO	$9 \times10^{-3}$ / 1	legacy’s dust-temperature table beyond its grid
earlyUniverse	$1 \times10^{-8}$ / $4 \times10^{-2}$	two Hubble constants in legacy
shock1D, shock1Dcool, euler1D	$\le5 \times10^{-4}$ / $6 \times10^{-3}$	legacy tolerance
cloudSurface	–	(species $\le1$ per cent)
rtchem	–	photon-bin integration (fronts within 2 percent)
lamda	$5 \times10^{-9}$	legacy’s $hc/k_{\mathrm{B}}$ in the level energies
auto, map	– / 0.57	legacy’s midpoint sampling of the photon bins
chianti	$3 \times10^{-7}$	legacy’s rate table; lines with $\lambda\le0$
dust	–	accretion speed and sublimation (Section 3.2)

Vector (PDF)

Cooling and heating processes

Table D.1 lists the cooling and heating processes of KROME-X. They are chosen when the code is generated, by name (the names of legacy KROME), and only the selected ones are compiled into the core; the temperature equation then evolves with their sum, and each term is also available separately at run time. Processes whose species are absent from the network vanish automatically. Most are ported from legacy KROME; the photoheating of the photon bins, the binned dust and the metal lines are described in Sects. 3.1–3.3, and the reaction enthalpies in Appendix E.

Atomic processes, constants and data of the updated physics

We describe the changes to atomic processes, constants, and data here.

– Recombination. The rates of radiative recombination [4] and dielectronic recombination [5] of 370 ions are now available automatically from the reaction database (legacy KROME used, for several abundant ions, fits without dielectronic recombination), with the radiative fits of Dr. Verner’s atomic-data compilation for the others; the collisional ionization follows Voronov [143]. At $10^4$ K the new recombination rates are 1.6 (C$^+$), 2.7 (Si$^+$) and 2.9 (S$^+$) times the legacy ones.

– Recombination cooling. The recombination term of the atomic cooling is computed from the network’s own recombination reactions,

the mean kinetic energy carried away by the captured electrons, exact for any capture rate $k_X(T)$ on a Maxwellian distribution. Legacy KROME used the H and He fits of Cen [21], independently of the recombination rates of the network: for the primordial network, the new term is 0.85 (H$^+$) times Cen’s at $10^4$ K, and 0.36 at $10^6$ K.

– Cosmological recombination. The early-Universe network uses the RECFAST 1.5.2 model [122, 123, 151, 116] for H, D and He I, written as recombination and CMB photoionization rate coefficients, so that it becomes part of the network. The rates reproduce the derivatives of the RECFAST code to $3.4 \times10^{-4}$, and the integrated ionization history follows it within 2.5 percent from $z=2500$ to 10.

– Constants and defaults. All physical constants all come from one table, CODATA 2018 [138] and the IAU units, used by the generated code, the runtime, and the converters (legacy KROME mixed several older sets, and had a typo in the $h c/k_{\mathrm{B}}$ factor of the partition functions, which lowered the rotational and vibrational temperatures by 0.7 percent). The H$_2$ self-shielding defaults to Wolcott-Green et al. [149], and the reaction enthalpies of the chemical heating are taken at $T_{\mathrm{gas}}$.

name	process	reference
Cooling
ATOMIC	H, He, He$^{+}$ collisional excitation and ionization; recombination cooling from the network’s own recombination rates, $n_{\mathrm{e}}n_{\mathrm{X}}k_{\mathrm{B}}T\left(\frac{3}{2}k+T\,\mathrm{d}k/\mathrm{d}T\right)$	Cen [21]
H2	H$_2$ rovibrational lines: collisions with H, H$^{+}$, H$_2$, He, e$^{-}$, and the LTE limit; optional optically thick suppression	Glover & Abel [47]; Hollenbach & McKee [65]; Galli & Palla [43]; Glover [46]; Ripamonti & Abel [113]
HD	HD rovibrational lines	Lipovka et al. [81]
CIE	H$_2$ collision-induced emission, with its opacity	Ripamonti & Abel [113]
CONT	continuum emission of primordial gas, $4\sigma T^{4}\kappa_{\mathrm{P}}\rho$, escaping over a Jeans length
COMPTON	inverse Compton scattering of the CMB	Cen [21]
FF	free–free emission of the atomic ions
EXPANSION	adiabatic cosmological expansion	Galli & Palla [43]
CHEM	collisional dissociation of H$_2$ (4.48 eV per event)
Z	fine-structure lines of C, C$^{+}$, O, O$^{+}$, Si, Si$^{+}$, Fe, Fe$^{+}$, S from the level populations (Sect. 3.3, Table E.1); optional escape probability (LVG), radiative and CMB pumping	Maio et al. [90]; van der Tak et al. [142]
Z_CIE	metal cooling and heating in ionization equilibrium, with or without the UV background, scaled by the metallicity	Shen et al. [129]; Haardt & Madau [57]
Z_EXTENDED	Z below $10^{4}$ K, Z_CIE above
CO, OH, H2O, HCN	molecular lines, tabulated in temperature, density, and column density	Neufeld & Kaufman [101]; Omukai et al. [105]
DUST	gas–grain collisional exchange with one effective grain population	Hollenbach & McKee [65]
DUSTTAB	the same, with the dust temperature tabulated for a given radiation field	Grassi et al. [53]
DUSTBINS	gas–grain exchange of the dust size bins, each at its own temperature (Sect. 3.2)	Draine [33]
DH	reaction enthalpies at $T_{\mathrm{gas}}$ (cooling and heating)	Burcat & Ruscic [17]
Heating
COMPRESS	compression in a collapse, $n k_{\mathrm{B}}T/t_{\mathrm{ff}}$
VISCOUS	viscous dissipation in a disc, $\frac{9}{4}\nu\rho\Omega^{2}$
PHOTO	the photoelectron energy of every photoionization and photodissociation of the photon bins (Sect. 3.1); X-ray secondary electrons (Sect. 5.2)	Shull & van Steenberg [130]; Dalgarno et al. [25]; Quadri et al. [109]
CR	cosmic-ray ionization, 4.3 eV per ionization of H and He, a density-dependent value for H$_2$	Glassgold et al. [45]
CHEM	H$_2$ formation in the gas phase and on grains, and UV pumping, the thermalized fraction $1/(1+n_{\mathrm{cr}}/n)$	Hollenbach & McKee [65]
PHOTOAV	H$_2$ photodissociation and UV pumping in a field $G_{0}$ attenuated by $A_{\mathrm{V}}$
PHOTODUST	grain photoelectric heating; PHOTODUSTNET net of grain recombination cooling	Bakes & Tielens [6]; Wolfire et al. [150]; Weingartner & Draine [147]
NUCLEAR	nuclear energy release of the reactions

• Other chemistry. Reverse reactions are generated by detailed balance from the NASA polynomials of the Burcat database [17], $k_{\mathrm{rev}}=k_{\mathrm{fwd}}/K_{\mathrm{c}}$, $K_{\mathrm{c}}=e^{-\Delta G/RT}(\rho_{0}/k_{\mathrm{B}}T)^{\Delta n}$ (legacy applied the pressure factor twice, and its table assigned to H$_2$CO and HCO the data of their isomers HCOH and COH); electron screening of nuclear reactions [67] is applied in the ODEs; and kinetic nucleation with detailed-balance destruction is supported.

The effect of the updated physics on the legacy tests

We ran the 21 legacy test problems with the physics of Section 3 and compared them with the runs of Section 4, which include the physics of legacy KROME. Where none of the updates apply, the tests move only through the physical constants, by $\lesssim10^{-5}$ (e.g. the dark cloud with grain-surface chemistry, $4\times10^{-7}$ in the median). The tests that change are discussed below; Table F.1 summarizes them.

• Metal line cooling. Fig. F.1 compares the cooling per coolant particle of the legacy and the new line-cooling data in atomic, molecular, and ionized gas. In atomic gas ($n_{\mathrm{H}}=30\ \mathrm{cm}^{-3}$, $x_{\mathrm{e}}=10^{-4}$) the new rates change the cooling of C$^{+}$ by $-10$ to $+21$ per cent and of C by up to $+22$ per cent, and increase that of O by 1.4–2.4 below 1000 K and that of Si by up to a factor of 5 at 3000 K; the cooling of Si$^{+}$ changes from 0.6 (30 K) to 2 times (8000 K) the legacy value, and that of O$^{+}$ is 0.22 times the legacy value, mostly because of the corrected ground-level degeneracy. The largest changes are in molecular gas: without H$_2$ collisions in the legacy data, C$^{+}$ and O cooled 12–22 and 70–230 times less at 20–300 K ($n_{\mathrm{H}_{2}}=50\ \mathrm{cm}^{-3}$). In ionized gas, the new electron rates lower the Fe$^{+}$ cooling by factors 2–3.5. The 16-level Fe$^{+}$ adds the $a^{4}\mathrm{F}$ and $a^{4}\mathrm{D}$ terms, which matter above $\sim2000$ K (26 per cent more cooling at 4000 K).

Vector (PDF)

Vector (PDF)

	levels, radiative data	collisions
C	LAMDA, 3 levels	H: Yan & Babb [155]; e, H$^{+}$, He, H$_2$: LAMDA
C$^{+}$	LAMDA, 2	H$_2$: Kłos et al. (2021); H: Barinovs et al. [8]
O	LAMDA, 3	H, He, H$_2$: Lique et al. [82]; e, H$^{+}$: LAMDA
O$^{+}$	CHIANTI, 5; NIST	e: Tayal [136]
Si	CHIANTI, 3; NIST	H: Yan & Babb [154]; e: Dufresne et al. [39]
Si$^{+}$	CHIANTI, 2; NIST	e: Tayal [137]; H: Barinovs et al. [8]
S	NIST, 3	H: Yan & Babb [154]
Fe	legacy, 5	legacy
Fe$^{+}$	CHIANTI, 5 (16); NIST	e: Zhang & Pradhan [156]; H: legacy

test	$\Delta T_{\mathrm{gas}}/T_{\mathrm{gas}}$	cause
collapse	$8 \times10^{-6}$ / $7 \times10^{-3}$	exact adiabatic index
collapseZ	$1.9 \times10^{-3}$ / $0.43$	metal line cooling (H$_2$ collisions)
collapseUV	$5 \times10^{-6}$ / $3.5 \times10^{-3}$	adiabatic index, recombination cooling
collapseUV_Xrays	$1 \times10^{-5}$ / $1.1 \times10^{-2}$	recombination cooling
earlyUniverse	$0$ / $0.32$	RECFAST
map	$1.9 \times10^{-3}$ / $0.58$	recombination cooling (3 of 2500 models change branch)
chianti	fixed $T_{\mathrm{gas}}$	dielectronic recombination: cooling $\times0.37$–$1.39$
auto	fixed $T_{\mathrm{gas}}$	C$^{+}$ recombination: C fraction $\times1.2$–$1.7$
shock1Dcool	$0$ / $5.5 \times10^{-4}$	enthalpies at $T_{\mathrm{gas}}$
others	$\lesssim10^{-5}$	physical constants

The collapse of low-metallicity gas (Fig. F.2) shows the consequence: the additional cooling between 100 and 300 K, mostly by C$^{+}$ and O excited by H$_2$, lowers the temperature along the low-density part of the tracks, by up to 43 percent at $Z = 10^{-2}\,Z_{\odot}$ (from 151 to 86 K at $n = 5\ \mathrm{cm}^{-3}$), 26 per cent at $0.1\,Z_{\odot}$, 23 per cent at $10^{-3}\,Z_{\odot}$ and 14 per cent at $10^{-4}\,Z_{\odot}$. At high density, the tracks converge again, the remaining change (0.6 percent) coming from the exact adiabatic index.

Vector (PDF)

• Cosmological recombination. Fig. F.3 shows the chemistry of the early Universe with the recombination rates of legacy krome (from Sasaki & Takahara [118]) and with RECFAST. The residual ionization at freeze-out is $x_{\mathrm{e}} = 1.62 \times10^{-4}$ instead of $1.29 \times10^{-4}$ (RECFAST itself gives $1.63 \times10^{-4}$), and the larger electron fraction forms more H$_2$ ($1.85 \times10^{-6}$ instead of $1.10 \times10^{-6}$ at $z = 10$) and HD ($1.5 \times10^{-9}$ instead of $5.9 \times10^{-10}$); the gas, coupled longer to the CMB, is warmer at $z = 10$ (1.33 K instead of 0.90 K).

• Recombination. The dielectronic recombination changes the ionization balance of heavy elements (Fig. F.4): the cooling of carbon in collisional and photoionization equilibrium changes by factors 0.37–1.39 between $10^{4}$ and $10^{6}$ K, and the neutral carbon fraction of a gas photoionized by a $J_{21}$ power law increases from 0.153 to 0.225 at $J_{21} = 0.1$ (and from 0.022 to 0.037 at $J_{21} = 1$), because C$^{+}$ recombines 1.6 times faster at $1.2 \times10^{4}$ K. In the same test, the larger difference between legacy krome and krome-x comes from the photon bins, integrated exactly in krome-x (§ 3.1).

GPU offloading: the device kernel and the alternatives

This appendix collects the measurements behind the device backend of § 2.6: where its time goes, and the alternatives that were implemented and measured against it.

The device kernel

Table G.1 shows, for four cases, which parts of the code the threads of the device kernel spend their time in. The arithmetic of the chemistry is cheap on the GPU: the rates, the right-hand

	Collapse	Coll.-Metals (U/H)	Photo
integrator (BDF), total	77	80/85	93
dense LU	20	41/48	67
Newton matrix $I - \gamma J$	19	17/22	19
rest of the step	38	22/15	7
rates and right-hand side	6	5/4	6
thermal terms	10	11/8	< 1
rest of the kernel	7	3/2	1

case	$N = 10^{3}$ cells cont.	$N = 10^{3}$ cells int.	$N = 10^{4}$ cells cont.	$N = 10^{4}$ cells int.	$N = 10^{5}$ cells cont.	$N = 10^{5}$ cells int.
Collapse U	2.48	2.69	0.570	0.519	0.605	0.351
Collapse H	17.2	19.1	1.84	1.94	0.759	0.555
Collapse-Metals U	13.2	14.1	2.67	2.70	2.55	1.29
Collapse-Metals H	68.2	114	9.18	13.9	3.45	2.91
Photo U	29.7	40.3	8.18	9.26	7.94	3.88
Photo H	82.2	143	16.9	22.9	7.11	5.13

side, and the thermal terms take 6–17 percent of the time, while the factorization and the assembly of the Newton matrix take 40–86 percent, rising with the size of the network (10, 28, and 34 equations for Collapse, Collapse-Metals, and Photo). On one CPU core, the picture is the opposite: three-quarters of the collapse time goes into the right-hand side. Each thread factorizes its own dense matrix in global memory, and the kernel is latency-bound (occupancy of 6 percent with 255 registers per thread at $10^{4}$ cells).

• Layout and launch. An A100 holds $108 \times256 = 27\,648$ resident threads of the integrator. Below this number of cells, the contiguous layout wins (1.1–1.8 times faster at $10^{3}$ cells, up to 1.5 at $10^{4}$), each thread’s working set staying in the L1 cache; above it, the interleaved layout (64 registers per thread, blocks of 128) keeps all the cells resident and coalesces the loads (3.3 times fewer memory sectors), and is 1.2–2.1 times faster at $10^{5}$ cells (Table G.2). The default selects the layout based on the number of cells relative to the device’s resident threads.

• Heterogeneous cells. Each thread follows its own cell, so the cost of heterogeneity is the divergence of the threads of a warp: a Collapse-Metals batch of heterogeneous cells costs 3.4 times a uniform one at $10^{4}$ cells and 2.3 times at $10^{5}$ (Table G.2). We tested the possibility of using a queue from which idle threads take the next unsolved cell, but it gave no measurable gain, since it could not remove the divergence; sorting the cells by expected work would, but needs a re-ordering in the host code around every call.

• A team of threads per Newton matrix. Since the factorization dominates, we tried factorizing each Newton matrix

@page 17/19 starts_mid=1 ends_mid=1

Vector (PDF)

mostly by C$^{+}$ and O excited by H$_2$, lowers the temperature along the low-density part of the tracks, by up to 43 percent at $Z = 10^{-2} Z_{\odot}$ (from 151 to 86 K at $n = 5\ \mathrm{cm}^{-3}$), 26 per cent at $0.1 Z_{\odot}$, 23 per cent at $10^{-3} Z_{\odot}$ and 14 per cent at $10^{-4} Z_{\odot}$. At high density, the tracks converge again, the remaining change (0.6 percent) coming from the exact adiabatic index.

• Cosmological recombination. Fig. F.3 shows the chemistry of the early Universe with the recombination rates of legacy krome (from Sasaki & Takahara [118]) and with RECFAST. The residual ionization at freeze-out is $x_{\mathrm{e}} = 1.62 \times10^{-4}$ instead of $1.29 \times10^{-4}$ (RECFAST itself gives $1.63 \times10^{-4}$), and the larger electron fraction forms more H$_2$ ($1.85 \times10^{-6}$ instead of $1.10 \times10^{-6}$ at $z = 10$) and HD ($1.5 \times10^{-9}$ instead of $5.9 \times10^{-10}$); the gas, coupled longer to the CMB, is warmer at $z = 10$ (1.33 K instead of 0.90 K).

• Recombination. The dielectronic recombination changes the ionization balance of heavy elements (Fig. F.4): the cooling of carbon in collisional and photoionization equilibrium changes by factors 0.37–1.39 between $10^{4}$ and $10^{6}$ K, and the neutral carbon fraction of a gas photoionized by a $J_{21}$ power law increases from 0.153 to 0.225 at $J_{21} = 0.1$ (and from 0.022 to 0.037 at $J_{21} = 1$), because C$^{+}$ recombines 1.6 times faster at $1.2 \times10^{4}$ K. In the same test, the larger difference between legacy krome and krome-x comes from the photon bins, integrated exactly in krome-x (§ 3.1).

GPU offloading: the device kernel and the alternatives

This appendix collects the measurements behind the device backend of § 2.6: where its time goes, and the alternatives that were implemented and measured against it.

The device kernel

Table G.1 shows, for four cases, which parts of the code the threads of the device kernel spend their time in. The arithmetic of the chemistry is cheap on the GPU: the rates, the right-hand

	Collapse	Coll.-Metals (U/H)	Photo
integrator (BDF), total	77	80/85	93
dense LU	20	41/48	67
Newton matrix $I - \gamma J$	19	17/22	19
rest of the step	38	22/15	7
rates and right-hand side	6	5/4	6
thermal terms	10	11/8	< 1
rest of the kernel	7	3/2	1

case	$N = 10^{3}$ cells cont.	$N = 10^{3}$ cells int.	$N = 10^{4}$ cells cont.	$N = 10^{4}$ cells int.	$N = 10^{5}$ cells cont.	$N = 10^{5}$ cells int.
Collapse U	2.48	2.69	0.570	0.519	0.605	0.351
Collapse H	17.2	19.1	1.84	1.94	0.759	0.555
Collapse-Metals U	13.2	14.1	2.67	2.70	2.55	1.29
Collapse-Metals H	68.2	114	9.18	13.9	3.45	2.91
Photo U	29.7	40.3	8.18	9.26	7.94	3.88
Photo H	82.2	143	16.9	22.9	7.11	5.13

side, and the thermal terms take 6–17 percent of the time, while the factorization and the assembly of the Newton matrix take 40–86 percent, rising with the size of the network (10, 28, and 34 equations for Collapse, Collapse-Metals, and Photo). On one CPU core, the picture is the opposite: three-quarters of the collapse time goes into the right-hand side. Each thread factorizes its own dense matrix in global memory, and the kernel is latency-bound (occupancy of 6 percent with 255 registers per thread at $10^{4}$ cells).

• Layout and launch. An A100 holds $108 \times256 = 27\,648$ resident threads of the integrator. Below this number of cells, the contiguous layout wins (1.1–1.8 times faster at $10^{3}$ cells, up to 1.5 at $10^{4}$), each thread’s working set staying in the L1 cache; above it, the interleaved layout (64 registers per thread, blocks of 128) keeps all the cells resident and coalesces the loads (3.3 times fewer memory sectors), and is 1.2–2.1 times faster at $10^{5}$ cells (Table G.2). The default selects the layout based on the number of cells relative to the device’s resident threads.

• Heterogeneous cells. Each thread follows its own cell, so the cost of heterogeneity is the divergence of the threads of a warp: a Collapse-Metals batch of heterogeneous cells costs 3.4 times a uniform one at $10^{4}$ cells and 2.3 times at $10^{5}$ (Table G.2). We tested the possibility of using a queue from which idle threads take the next unsolved cell, but it gave no measurable gain, since it could not remove the divergence; sorting the cells by expected work would, but needs a re-ordering in the host code around every call.

• A team of threads per Newton matrix. Since the factorization dominates, we tried factorizing each Newton matrix

Vector (PDF)

Vector (PDF)

by a team of threads in shared memory, keeping one cell per thread for the rest of the integration: the integrator pauses whenever it needs a new matrix, and the warp builds and factorizes the matrices of all its pending cells together. The factorization alone becomes 3–30 times faster for 28 and 34 equations (Table G.3), but the whole integration is slower in most cases: the pauses, the barriers, and the resumable state of the integrator, kept in memory instead of registers, cost about what the faster factorization saves. On DarkCloud, whose $655 \times655$ matrices the team factorizes in global memory, it was 1.2 times faster than the dense per-thread kernel at 100 cells and 0.79 times at 500, both far slower than the per-thread sparse LU. The team kernel is not part of the released code.

– Batched integration. The alternative to one integration per cell is one integration of all the cells of a call as a single system, as provided by sundials on GPUs ([7]): cvode on the concatenated state, with one step size and order for all the cells and a block-diagonal Newton matrix factorized by a batched solver. We measured it with the sundials cuda vectors and fused kernels and two batched linear

case (eq.)	$N$	thread	team	integration
Collapse (10)	$10^3$	0.060	0.024	0.48
	$10^4$	0.0060	0.0067	0.66
	$10^5$	0.0039	0.0053	0.58
Collapse-Metals (28)	$10^3$	1.15	0.084	0.63
	$10^4$	0.199	0.039	0.80
	$10^5$	0.095	0.030	0.76
Photo (34)	$10^3$	3.39	0.119	0.93
	$10^4$	0.743	0.073	1.26
	$10^5$	0.260	0.060	1.13

solvers, the dense LU of magma 2.8 and the sparse QR of cusparse, on one A100, through the whole history of each case (Table G.4). In the uniform batches, the cells differ by a

case	$N$	per cell	MAGMA	cuSPARSE	8 cores
Collapse	$10^3$	0.0167	0.0189	0.0326	0.0203
	$10^4$	0.00347	0.00437	0.00811	0.0201
	$10^5$	0.00238	0.00351	0.00703	–
Collapse-Metals	$10^3$	0.102	0.0657	–	0.0730
	$10^4$	0.0200	0.0178	–	0.0729
	$10^5$	0.0108	0.0161	–	–
C.-Metals, H	$10^3$	0.185	0.211	–	0.0805
	$10^4$	0.0327	0.0825	–	0.0799
	$10^5$	0.0148	$\ge0.089$	–	–
Photo	$10^3$	0.00286	0.00106	0.00313	0.000875
	$10^4$	0.000706	0.000417	0.00101	0.000755
	$10^5$	0.000453	0.000385	0.000910	–
DarkCloud, 3 calls	$10^2$	63.2	0.128	5.27	0.103
	$10^3$	8.48	0.0320	1.99	0.0952
DarkCloud	$10^3$	–	3.71	–	1.46

30 percent spread in $T_{\mathrm{gas}}$ or by the attenuation of their field; in the heterogeneous batch of Collapse-Metals the densities span a factor of $10^3$ and the metallicities $10^{-3}$–$10^{-1} Z_\odot$. On uniform batches, the batched integration is competitive: MAGMA is faster than the per-cell integrator on Photo (1.2–2.7 times) and on Collapse-Metals up to $10^4$ cells (1.6 and 1.1 times), slower on Collapse. cuSPARSE is slower than MAGMA everywhere. The shared step is what breaks it on cells that differ: every cell takes the steps of the stiffest one (each still meets its own tolerance, as the error test takes the largest of the per-cell norms). On the heterogeneous batch, it takes 2.5–4 times the steps of the per-cell integrations, and the per-cell integrator is 1.1, 2.5, and at least 6 times faster at $10^3$, $10^4$ and $10^5$ cells. On DarkCloud, the batched LU of MAGMA is far faster than a thread’s dense factorization, but the shared step costs 13 times the steps over the whole history (3.7 s per cell at $10^3$ cells, 2.5 times the eight cores). A looser error norm does not help: one norm over all the cells (the default of SUNDIALS) makes the heterogeneous batch only 1.05 times faster at $10^4$ cells, and leaves 759 of the $10^4$ cells more than $10^{-2}$ from a reference at $r_{\mathrm{tol}}=10^{-8}$, against 21 with the per-cell maximum, which over-resolves the easy cells (the per-cell integrations at $r_{\mathrm{tol}}=10^{-4}$ leave about a quarter of these chaotic histories that are far from the reference).

Linear solvers. Table G.5 gives the solver time of CVODE with different linear solvers on one core of the Apple MacBook Pro M1, for Collapse (10 equations) and DarkCloud (655 species). Of the general solvers, the dense LU, from LAPACK above $\sim$64 equations, is the fastest in every case: the Jacobian of Collapse is 68 percent full, and for DarkCloud (5.7 per cent) the factors of KLU [27] fill 26 per cent of the matrix, with pivoting and a symbolic analysis at every setup. The Krylov solvers of GINKGO [3] and SUNDIALS suffer from systems that are both stiff and badly scaled: a preconditioner good enough is close to the full LU, and inexact solves degrade the accuracy. The sparse LU of KROME-X (Section 2.5) avoids the pivoting and the analysis, its structure being fixed when the code is generated; for DarkCloud, its factors fill 11 percent of the matrix, one factorization takes $2.6 \times10^6$ operations instead of $1.9 \times10^8$, and the run is 2.2 times faster than with LAPACK. On the networks of tens of species of Section 6, whose factors are 45–87 percent full, it does not pay, on the CPU or on the GPU.

linear solver	Collapse	DarkCloud
dense LU (SUNDIALS)	0.13	35.4
dense LU (LAPACK, Accelerate)	0.14	4.75
sparse LU (KLU)	0.20	17.4
sparse LU, generated order	–	2.17
BiCGSTAB + ILU(0) (Ginkgo)	2.6	14.0
GMRES + block Jacobi (Ginkgo)	$>120$	$135^a$
GMRES + exact sparse LU (Ginkgo)	$>120$	$>600$
GMRES, matrix-free, no preconditioner	$0.21^b$	$>300$
GMRES, matrix-free + Jacobi	$0.48^c$	$>300$