Online Autotuning of Algebraic Multigrid via Contextual Bandits and Reinforcement Learning
Xudong He, Ruiyu Wang, Yilang Ding, Rui Peng Li, Yuanzhe Xi
Abstract
Algebraic multigrid (AMG) methods are widely used for solving large sparse linear systems, but their efficiency depends strongly on interacting setup and solve parameters. Effective parameter choices typically depend on the problem. We develop a phase-aware online autotuning method for hypre's BoomerAMG to reduce cumulative time-to-solution over sequences of related linear systems. A LinUCB-based contextual bandit selects mixed-type parameters for AMG hierarchy construction using a shared cost model over PDE descriptors and configuration encodings. A recursive LSTDQ controller uses residual and timing observations to estimate the AMG solve time within each attempt and adapt relaxation weights between multigrid cycles. Both learners update from ongoing solves without offline pretraining. Across grids of sizes $40^3$, $60^3$, and $80^3$, the joint autotuning method reduces total runtime, including recurring learner computation, by 36.1-43.3% for diffusion problems and 19.8-31.0% for diffusion-advection problems over sequences of 5000 linear systems, relative to the reference BoomerAMG default configuration. On a variable-flow diffusion-advection benchmark comprising 30,000 systems on a $60^3$ grid, 2 systems fail to converge with joint autotuning, compared with 1677 for the reference configuration, under their respective attempt budgets.