QiMeng-TensorOp: Automatically Generating High-Performance Tensor Operators with Hardware Primitives

TL;DR

QiMeng-TensorOp is a framework that enables large language models to automatically generate high-performance tensor operators tailored to diverse hardware architectures, significantly improving efficiency and reducing development costs.

Contribution

This work introduces a novel auto-generation framework that leverages LLMs to produce optimized tensor operators with hardware primitives across various architectures.

Findings

Achieves up to 1291x performance improvement over vanilla LLMs.

Reaches 251% of OpenBLAS performance on RISC-V CPUs.

Reduces development costs by 200x compared to human experts.

Abstract

Computation-intensive tensor operators constitute over 90\% of the computations in Large Language Models (LLMs) and Deep Neural Networks.Automatically and efficiently generating high-performance tensor operators with hardware primitives is crucial for diverse and ever-evolving hardware architectures like RISC-V, ARM, and GPUs, as manually optimized implementation takes at least months and lacks portability.LLMs excel at generating high-level language codes, but they struggle to fully comprehend hardware characteristics and produce high-performance tensor operators. We introduce a tensor-operator auto-generation framework with a one-line user prompt (QiMeng-TensorOp), which enables LLMs to automatically exploit hardware characteristics to generate tensor operators with hardware primitives, and tune parameters for optimal performance across diverse hardware. Experimental results on various hardware platforms, SOTA LLMs, and typical tensor operators demonstrate that QiMeng-TensorOp effectively unleashes the computing capability of various hardware platforms, and automatically generates tensor operators of superior performance. Compared with vanilla LLMs, QiMeng-TensorOp achieves up to $1291 \times$ performance improvement. Even compared with human experts, QiMeng-TensorOp could reach $251 \%$ of OpenBLAS on RISC-V CPUs, and $124 \%$ of cuBLAS on NVIDIA GPUs. Additionally, QiMeng-TensorOp also significantly reduces development costs by $200 \times$ compared with human experts.