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

TL;DR

QiMeng-TensorOp tackles the challenge of automatically generating high-performance tensor operators at the hardware-primitive level across heterogeneous platforms. It combines hardware-aware prompts, explicit hardware-factor extraction, and an LLM-assisted MCTS-driven auto-tuning loop to produce assembly-level or PTX-level kernels tailored to target devices. The approach delivers dramatic performance gains—up to $1291\times$ over vanilla prompts and substantial advantages over manually optimized libraries on several platforms—while reducing development costs by up to $200\times$. This work enables portable, high-performance tensor operators for evolving hardware and paves the way for rapid deployment of optimized kernels across CPU, RISC-V, ARM, and GPU architectures.

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.

Figures (5)

• Figure 1: Comparison of tensor operator optimization paradigms
• Figure 2: QiMeng-TensorOp overview. The proposed framework takes a user's one-sentence description as input and generates high-performance tensor operators using hardware primitives via three automatic steps. Step 1 activates LLMs’ comprehension of tensor operator optimization and extracts target hardware factors for subsequent generation. In Step 2, sketches and kernels are generated to form tensor operators. Sketch generation uses optimizations such as tiling, and hardware-primitive-level kernels like PACK and COMPUTE employ instruction-level optimization. In Step 3, Auto-Tuning uncovers subtle optimization opportunities and further enhances performance. (A detailed showcase please refer to Appendix C.)
• Figure 3: Conv operator performance comparison
• Figure 4: (a) Ablation results of QiMeng-TensorOp components for GEMM on K1 CPU. (b) Performance comparison of tuning methods on GEMM of size 1024 on K1 CPU.
• Figure 5: (a) Development cost comparison on GEMM of size 1024 (b) GEMM performance comparison between QiMeng-TensorOp and TVM with different search counts on A76 CPU.