A hybrid GPU/CPU implementation of the average information restricted maximum likelihood algorithm for fitting multivariate linear mixed models
The multivariate linear mixed model combined with restricted maximum likelihood estimation of (co)variance components and solving the mixed model equations is widely used for the analysis of plant breeding data. Several algorithms for restricted maximum likelihood estimation have been proposed, with the average information algorithm being widely used for its computational efficiency, quadratic rate of convergence, and ability to handle large amounts of parameters. A key reason for the scalability of linear mixed models is the exploitation of sparse matrix techniques. Recently, decreasing costs of high-throughput data collection have led to high-dimensional sets of features characterizing genotypes and environments. A convenient way of using these features is through the use of kernel-based covariance structures. However, while convenient, these kernels typically have dense inverses meaning optimizations based on sparse matrix techniques are not always applicable. For these settings, the introduction of general-purpose computing on graphics processing units brought significant computational advantages for dense linear algebra operations. We present a hybrid GPU/CPU implementation of the average information restricted maximum likelihood algorithm. We eliminate sparse equations using supernodal Cholesky factorization and rely on CUDA GPU acceleration to solve the resulting reduced dense system. We show the significant computational advantages of this approach for fully dense problems, as well as problems with a mix of dense and sparse equations.