Accelerated linear algebra compiler for computationally efficient numerical models: Success and potential area of...

Area of Science:

• Computer Science
• Computational Science
• Numerical Analysis

Background:

• Machine learning advancements are driven by high-performance computing tools like the Accelerated Linear Algebra (XLA) compiler.
• XLA optimizes array operations, fusing them to reduce memory usage and generate platform-specific high-performance code.
• Numerical models, similar to machine learning models, rely heavily on array operations, suggesting potential performance gains from XLA.

Purpose of the Study:

• To evaluate the efficiency of the XLA compiler for numerical models.
• To compare XLA's performance against optimal implementations on CPU and GPU platforms.
• To identify factors influencing XLA's performance with different numerical models and computing architectures.

Main Methods:

• Rigorous performance testing of XLA on CPU and GPU shared-memory platforms.
• Comparison of XLA-compiled models with optimally implemented counterparts.
• Development of simple equations to model numerical model computing speed based on measurable parameters.
• Analysis of XLA's operation optimization (fusion) and compilation efficiency across various workloads and platforms.

Main Results:

• XLA successfully fuses low-level array operations (slice, concatenation, arithmetic) but struggles with high-level operations (convolution, roll).
• High efficiency (>80%) achieved for large problems on certain CPUs and simple GPU models; acceptable efficiency (10%-80%) for medium-size problems, impacted by Python overhead.
• Suboptimal performance observed on some CPU platforms and for high-dimensional GPU models due to inefficient use of specialized hardware instructions.
• Performance degradation on GPUs occurs when XLA attempts to leverage instructions for processing multiple floating-point numbers per thread, which are underutilized in complex models.

Conclusions:

• XLA offers significant performance improvements for numerical models, particularly for large problems on specific hardware configurations.
• The efficiency of XLA is dependent on the model's complexity, the type of operations used, and the target computing platform.
• Recommendations include adding compilation control flags to mitigate performance issues in non-optimal scenarios, especially for complex GPU models.