Precision Autotuning for Linear Solvers via Reinforcement Learning

TL;DR

The paper tackles the challenge of precision tuning for linear solvers by casting precision selection as a contextual bandit RL problem, using discretized state features derived from matrix properties. It introduces a general framework with per-step Q-tables and action-space reduction, and demonstrates its viability on GMRES-IR with four precision stages, balancing accuracy and computational cost. Key findings show that the RL approach can match FP64 accuracy while reducing cost on many dense cases, and that it can generalize to unseen data, though conservative behavior is observed for highly ill-conditioned sparse systems. Overall, the work advances mixed-precision numerical methods by providing a data-driven, generalizable precision autotuning framework with practical implications for scientific computing.

Abstract

We propose a reinforcement learning (RL) framework for adaptive precision tuning of linear solvers, and can be extended to general algorithms. The framework is formulated as a contextual bandit problem and solved using incremental action-value estimation with a discretized state space to select optimal precision configurations for computational steps, balancing precision and computational efficiency. To verify its effectiveness, we apply the framework to iterative refinement for solving linear systems $Ax = b$. In this application, our approach dynamically chooses precisions based on calculated features from the system. In detail, a Q-table maps discretized features (e.g., approximate condition number and matrix norm)to actions (chosen precision configurations for specific steps), optimized via an epsilon-greedy strategy to maximize a multi-objective reward balancing accuracy and computational cost. Empirical results demonstrate effective precision selection, reducing computational cost while maintaining accuracy comparable to double-precision baselines. The framework generalizes to diverse out-of-sample data and offers insight into utilizing RL precision selection for other numerical algorithms, advancing mixed-precision numerical methods in scientific computing. To the best of our knowledge, this is the first work on precision autotuning with RL and verified on unseen datasets.

Figures (10)

• Figure 1: Schematic of RL Framework: Phase I (top) depicts the iterative training loop where the agent updates the Q-table. Phase II (bottom) shows the inference workflow using the trained RL model.
• Figure 2: Comparison of RL and FP64 in terms of errors and iteration counts. Marker colors denote different matrix-size groups, while marker sizes represent the number of test samples in each group.
• Figure 3: RL($W_1$), $\tau=10^{-6}$
• Figure 4: RL($W_2$), $\tau=10^{-6}$
• Figure 5: RL($W_1$), $\tau=10^{-8}$