Yuanbo Wen 0001

dblp:262/3144-1 · DBLP profile ↗
← Back
33ranked-venue papers
2as first author
32since 2021 · last 2026
0000-0002-7775-2724ORCID · conflict

Domains — the database's venue-derived domains; a paper can count in several

Systems, architecture and hardware · 18 · 1 first-author · 17 since 2021Artificial intelligence and machine learning · 14 · 1 first-author · 14 since 2021Software engineering, systems software and programming languages · 6 · 6 since 2021Graphics, computer vision, multimedia, augmented reality and games · 4 · 4 since 2021Applied, interdisciplinary, general and emerging computing · 1 · 1 since 2021
YearPublicationVenuePosition
2026 QiMeng-Kernel: Macro-Thinking Micro-Coding Paradigm for LLM-Based High-Performance GPU Kernel Generation
abstract
Developing high-performance GPU kernels is critical for AI and scientific computing, but remains challenging due to its reliance on expert crafting and poor portability. While large language models (LLMs) offer promise for automation, both general-purpose and finetuned LLMs suffer from two fundamental and conflicting limitations: correctness and efficiency. The key reason is that existing LLM-based approaches directly generate the entire optimized low-level programs, requiring exploration of an extremely vast space encompassing both optimization policies and implementation codes. To address the challenge of exploring an intractable space, we propose Macro Thinking Micro Coding (MTMC), a hierarchical framework inspired by the staged optimization strategy of human experts. It decouples optimization strategy from implementation details, ensuring efficiency through high-level strategy and correctness through low-level implementation. Specifically, Macro Thinking employs reinforcement learning to guide lightweight LLMs in efficiently exploring and learning semantic optimization strategies that maximize hardware utilization. Micro Coding leverages general-purpose LLMs to incrementally implement the stepwise optimization proposals from Macro Thinking, avoiding full-kernel generation errors. Together, they effectively navigate the vast optimization space and intricate implementation details, enabling LLMs for high-performance GPU kernel generation. Comprehensive results on widely adopted benchmarks demonstrate the superior performance of MTMC on GPU kernel generation in both accuracy and running time. On KernelBench, MTMC achieves near 100% and 70% accuracy at Levels 1-2 and 3, over 50% than SOTA general-purpose and domain-finetuned LLMs, with up to 7.3× speedup over LLMs, and 2.2× over expert-optimized PyTorch Eager kernels. On the more challenging TritonBench, MTMC attains up to 59.64% accuracy and 34× speedup. All models and datasets will be made publicly available.
Xinguo Zhu, Shaohui Peng, Jiaming Guo, Yunji Chen, Qi Guo 0001, Yuanbo Wen 0001, Hang Qin, Ruizhi Chen, Qirui Zhou, Ke Gao 0012, Ling Li 0001
AAAI6
2026 QiMeng-PRepair: Precise Code Repair via Edit-Aware Reward Optimization
abstract
Changxin Ke, Rui Zhang, Jiaming Guo, Yuanbo Wen, Li Ding, Shuo Wang, Xuyuan Zhu, Xiong Peng, Di Huang, Zidong Du, Xing Hu, Qi Guo, Yunji Chen. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2026.
Changxin Ke, Rui Zhang 0040, Jiaming Guo, Yuanbo Wen 0001, Xuyuan Zhu, Xiong Peng, Zidong Du, Xing Hu 0001, Qi Guo 0001, Yunji Chen
ACL (1)4
2026 LLM-SYCL: Automated SYCL Generation from CUDA via Search-Driven LLM Translation
Zhou Liang, Yuanbo Wen 0001, En Shao, Guangming Tan
APPT3
2026 QiMeng-Tensify: Scaling Up Tensor Computation Optimization via Architecture-Aware LLM-Guided MCTS
Shouyang Dong, Jun Bi, Yuanbo Wen 0001, Xiyue Yu, Jianxing Xu, Guanglin Xu, Ling Li 0001, Xuehai Zhou, Tianshi Chen 0002, Qi Guo 0001
ISCA3
2026 FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor-Vector Parallelism
abstract
The attention mechanism is central to modern deep learning, particularly in large language models (LLMs), but suffers from quadratic computational complexity. To accelerate attention computation on GPUs, fused attention techniques (e.g., FlashAttention) consolidate the matrix multiplication (GEMM) and softmax computations into a single kernel. However, these operations remain computationally decoupled: the GEMM leverages high-performance tensor units (Tensor Cores), while the softmax executes on slower vector units (CUDA cores). This imbalance induces severe vector intervals—periods where tensor units sit idle awaiting vector unit completion—significantly underutilizing tensor units. Furthermore, ongoing hardware advancements delivering faster tensor units exacerbate this bottleneck.
Jianxing Xu, Yuanbo Wen 0001, Jun Bi, Ruibai Xu, Guanglin Xu, Rui Zhang 0040, Wei Li 0008, Ling Li 0001, Tianshi Chen 0002, Qi Guo 0001, Yunji Chen
PPoPP2
2026 AGON: Automated Design Framework for Customizing Processors From ISA Documents
abstract
Customized processors are essential for domain-specific applications such as the Internet of Things (IoT) and multi-media embedded systems, yet their design often requires extensive expert intervention. Traditional approaches, including hardware design using encapsulated abstractions (e.g., Chisel) and high-level synthesis (HLS) from languages like C or SystemC, reduce some manual efforts but remain either costly or suboptimal. Recent explorations into leveraging Large Language Models (LLMs) to generate RTL from natural language specifications have shown promise, but these methods still struggle with generating complex and high-performance processors mainly due to the complicated low-level details in the RTL code. In this work, we introduce AGON, a novel framework designed to facilitate the development of customized processor RTL from instruction set architecture (ISA) documents using LLMs. The framework comprises two layers: a functional description layer and a hardware implementation layer. At the functional layer, AGON employs a nano-operator (nOP)-based Intermediate Representation (IR) that abstracts basic instruction operations, thereby reducing the semantic gap between natural language and RTL code. This abstraction significantly shortens the descriptive code required for LLM generation, improving the generation accuracy in single-pass. At the hardware layer, AGON offers three abstraction levels (i.e. instruction, ISA, and processor) along with rule-based primitives to systematically lower the nOP-based IR into a fully optimized processor implementation. This decoupled design not only ensures correctness-by-construction but also enables automated, PPA-aware performance optimization. We evaluate AGON by designing high-performance out-of-order processors that correctly execute practical programs. Experimental results demonstrate that processors generated with AGON achieve an average speedup of 4.51× on specific tasks compared to expert-designed general-purpose CPUs while requiring minimal design effort.
Chongxiao Li, Pengwei Jin, Tianyun Ma, Husheng Han, Shuyao Cheng, Yifan Hao 0001, Yongwei Zhao 0001, Guanglin Xu, Zidong Du, Rui Zhang 0040, Xiaqing Li, Yuanbo Wen 0001, Xing Hu 0001, Qi Guo 0001
IEEE Trans. Comput. Aided Des. Integr. Circuits Syst.13
2026 DASA: Distribution-Aware Sparse Attention for Accelerating Diffusion Transformer
abstract
Diffusion Transformers (DiTs) have demonstrated remarkable success in text-to-video generation. However, the self-attention mechanism in DiTs imposes significant computational and memory burdens, particularly when handling long patch sequences like high-resolution or long-time videos. While sparse attention shows promise in reducing self-attention costs, existing approaches struggle to deliver performance gains due to the unique challenges in DiTs,i.e., varied sparse patterns across layers and timesteps, and the cumulative nature of inference errors over timesteps. In this paper, we propose DASA, an algorithm-hardware co-design that effectively addresses these challenges of attention sparsification in DiTs. Specifically, leveraging the insight that the generation quality is primarily influenced by overall distribution drift rather than changes in specific values, we introduce a novel Distribution-Aware Filtering (DAF) mechanism for sparsification. To further accelerate the process, we design a specialized Filtering Unit that enables fast candidate selection based on the proposed DAF mechanism. Experimental results show that DASA achieves 2.52× speed up compared to A100 GPU, and up to 1.22× speedup over state-of-the-art accelerators for self-attention computation.
Tianyun Ma, Jiaming Guo, Xinkai Song, Husheng Han, Pengwei Jin, Xiangtao Guan, Yifan Hao 0001, Yuanbo Wen 0001, Shuyao Cheng, Zidong Du, Rui Zhang 0040, Xing Hu 0001, Qi Guo 0001
IEEE Trans. Comput. Aided Des. Integr. Circuits Syst.10
2025 QiMeng-GEMM: Automatically Generating High-Performance Matrix Multiplication Code by Exploiting Large Language Models
abstract
As a crucial operator in numerous scientific and engineering computing applications, the automatic optimization of General Matrix Multiplication (GEMM) with full utilization of ever-evolving hardware architectures (e.g. GPUs and RISC-V) is of paramount importance. While Large Language Models (LLMs) can generate functionally correct code for simple tasks, they have yet to produce high-performance code. The key challenge resides in deeply understanding diverse hardware architectures and crafting prompts that effectively unleash the potential of LLMs to generate high-performance code. In this paper, we propose a novel prompt mechanism called QiMeng-GEMM which enables LLMs to comprehend the architectural characteristics of different hardware platforms and automatically search for the optimization combinations for GEMM. The key of QiMeng-GEMM is a set of informative, adaptive, and iterative meta-prompts. Based on this, a searching strategy for optimal combinations of meta-prompts is used to iteratively generate high-performance code. Extensive experiments conducted on 4 leading LLMs, various paradigmatic hardware platforms, and representative matrix dimensions unequivocally demonstrate QiMeng-GEMM’s superior performance in auto-generating optimized GEMM code. Compared to vanilla prompts, our method achieves a performance enhancement of up to 113×. Even when compared to human experts, our method can reach 115% of cuBLAS on NVIDIA GPUs and 211% of OpenBLAS on RISC-V CPUs. Notably, while human experts often take months to optimize GEMM, our approach reduces the development cost by over 240×.
Qirui Zhou, Yuanbo Wen 0001, Ruizhi Chen, Ke Gao 0012, Weiqiang Xiong, Ling Li 0001, Qi Guo 0001, Yunji Chen
AAAI2
2025 Mosaic: Exploiting Instruction-Level Parallelism on Deep Learning Accelerators with iTex Tessellation
abstract
Deep learning has achieved great success in numerous application areas at the cost of high computational complexity. To meet the ever-increasing computational demand, commodity hardware platforms (e.g., CPUs and GPUs) offer abundant computing resources including scalar, vector, and tensor units for deep learning that could execute in parallel. However, existing top-down tiling-based deep learning compilers often generate a homogeneous mapping from the given tensor computation task to hardware arithmetic instructions, failing to utilize different computing units simultaneously to achieve higher performance.
Jianxing Xu, Yuanbo Wen 0001, Ruibai Xu, Tingfeng Ruan, Jun Bi, Rui Zhang 0040, Xinkai Song, Yifan Hao 0001, Xing Hu 0001, Zidong Du, Chongqing Zhao, Jiang Jie, Qi Guo 0001
ASPLOS (2)2
2025 Automated Superscalar Processor Design by Learning Data Dependencies
abstract
Automated processor design, which can significantly reduce human efforts and accelerate design cycles, has received considerable attention. While recent advancements have automatically designed single-cycle processors that execute one instruction per cycle, their performance cannot compete with modern superscalar processors that execute multiple instructions per cycle. Previous methods fail on superscalar processor design because they cannot address inter-instruction data dependencies, leading to inefficient sequential instruction execution. This paper proposes a novel approach to automatically designing superscalar processors using a hardware-friendly model called the Stateful Binary Speculation Diagram (State-BSD). We observe that processor parallelism can be enhanced through on-the-fly inter-instruction dependent data predictors, reusing the processor's internal states to learn the data dependency. To meet the challenge of both hardware-resource limitation and design functional correctness, State-BSD consists of two components: 1) a lightweight state-selector trained by simulated annealing method to detect the most reusable processor states and store them in a small buffer; and 2) a highly precise state-speculator trained by BSD expansion method to predict the inter-instruction dependent data using the selected states. It is the first work to achieve the automated superscalar processor design, i.e. QiMeng-CPU-v2, which improves the performance by about 380x than the state-of-the-art automated design and is comparable to human-designed superscalar processors such as ARM Cortex A53.
Shuyao Cheng, Rui Zhang 0040, Wenkai He, Pengwei Jin, Chongxiao Li, Zidong Du, Xing Hu 0001, Yifan Hao 0001, Guanglin Xu, Yuanbo Wen 0001, Ling Li 0001, Qi Guo 0001, Yunji Chen
IJCAI10
2025 QiMeng-TensorOp: One-Line Prompt is Enough for High-Performance Tensor Operator Generation with Hardware Primitives
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× 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× compared with human experts.
Xuzhi Zhang, Shaohui Peng, Qirui Zhou, Yuanbo Wen 0001, Qi Guo 0001, Ruizhi Chen, Xinguo Zhu, Weiqiang Xiong, Haixin Chen, Congying Ma, Ke Gao 0012, Yunji Chen, Ling Li 0001
IJCAI4
2025 Cambricon-SR: An Accelerator for Neural Scene Representation with Sparse Encoding Table
abstract
Neural Scene Representation (NSR) is a promising technique for representing real scenes.By learning from dozens of 2D photos captured from different viewpoints, NSR computes the 3D representation of real scenes.However, the performance of NSR processing running on GPU is insufficient for applications.Cambricon-R achieves high performance of more than 60 scenes per second, but at the cost of modeling quality.
Tianbo Liu 0006, Xinkai Song, Zhifei Yue, Xing Hu 0001, Zhuoran Song, Yuanbo Wen 0001, Yifan Hao 0001, Wei Li 0008, Zidong Du, Rui Zhang 0040, Jiaming Guo, Shaohui Peng, Guangzhong Sun, Qi Guo 0001, Tianshi Chen 0002
ISCA7
2025 MigGPT: Harnessing Large Language Models for Automated Migration of Out-of-Tree Linux Kernel Patches Across Versions
abstract
Out-of-tree kernel patches are essential for adapting the Linux kernel to new hardware or enabling specific functionalities. Maintaining and updating these patches across different kernel versions demands significant effort from experienced engineers. Large language models (LLMs) have shown remarkable progress across various domains, suggesting their potential for automating out-of-tree kernel patch migration. However, our findings reveal that LLMs, while promising, struggle with incomplete code context understanding and inaccurate migration point identification. In this work, we propose MigGPT, a framework that employs a novel code fingerprint structure to retain code snippet information and incorporates three meticulously designed modules to improve the migration accuracy and efficiency of out-of-tree kernel patches. Furthermore, we establish a robust benchmark using real-world out-of-tree kernel patch projects to evaluate LLM capabilities. Evaluations show that MigGPT significantly outperforms the direct application of vanilla LLMs, achieving an average completion rate of 72.59\% ($\uparrow 50.74\%$) for migration tasks.
Pucheng Dang, Dong Li 0008, Yuanbo Wen 0001, Qi Guo 0001, Xing Hu 0001
NeurIPS5
2025 QiMeng-NeuComBack: Self-Evolving Translation from IR to Assembly Code
abstract
Compilers, while essential, are notoriously complex systems that demand prohibitively expensive human expertise to develop and maintain. The recent advancements in Large Language Models (LLMs) offer a compelling new paradigm: Neural Compilation, which could potentially simplify compiler development for new architectures and facilitate the discovery of innovative optimization techniques. However, several critical obstacles impede its practical adoption. Firstly, a significant lack of dedicated benchmarks and robust evaluation methodologies hinders objective assessment and tracking of progress in the field. Secondly, systematically enhancing the reliability and performance of LLM-generated assembly remains a critical challenge. Addressing these challenges, this paper introduces NeuComBack, a novel benchmark dataset specifically designed for IR-to-assembly compilation. Leveraging this dataset, we first define a foundational Neural Compilation workflow and conduct a comprehensive evaluation of the capabilities of recent frontier LLMs on Neural Compilation, establishing new performance baselines. We further propose a self-evolving prompt optimization method that enables LLMs to iteratively evolve their internal prompt strategies by extracting insights from prior self-debugging traces, thereby enhancing their neural compilation capabilities. Experiments demonstrate that our method significantly improves both the functional correctness and the performance of LLM-generated assembly code. Compared to baseline prompts, the functional correctness rates improved from 44% to 64% on x86_64 and from 36% to 58% on aarch64, respectively. More significantly, among the 16 correctly generated x86_64 programs using our method, 14 (87.5%) surpassed clang-O3 performance. These consistent improvements across diverse architectures (x86_64 and aarch64) and program distributions (NeuComBack L1 and L2) validate our method's superiority over conventional approaches and its potential for broader adoption in low-level neural compilation.
Hainan Fang, Yuanbo Wen 0001, Jun Bi, Tonghui He, Yanlin Tang, Jiaming Guo, Rui Zhang 0040, Qi Guo 0001, Yunji Chen
NeurIPS2
2025 QiMeng-MuPa: Mutual-Supervised Learning for Sequential-to-Parallel Code Translation
abstract
The rise of GPU-based high-performance computing (HPC) has driven the widespread adoption of parallel programming models such as CUDA. Yet, the inherent complexity of parallel programming creates a demand for the automated sequential-to-parallel approaches. However, data scarcity poses a significant challenge for machine learning-based sequential-to-parallel code translation. Although recent back-translation methods show promise, they still fail to ensure functional equivalence in the translated code. In this paper, we propose \textbf{QiMeng-MuPa}, a novel \textbf{Mu}tual-Supervised Learning framework for Sequential-to-\textbf{Pa}rallel code translation, to address the functional equivalence issue. QiMeng-MuPa consists of two models, a Translator and a Tester. Through an iterative loop consisting of Co-verify and Co-evolve steps, the Translator and the Tester mutually generate data for each other and improve collectively. The Tester generates unit tests to verify and filter functionally equivalent translated code, thereby evolving the Translator, while the Translator generates translated code as augmented input to evolve the Tester. Experimental results demonstrate that QiMeng-MuPa significantly enhances the performance of the base models: when applied to Qwen2.5-Coder, it not only improves Pass@1 by up to 28.91\% and boosts Tester performance by 68.90\%, but also outperforms the previous state-of-the-art method CodeRosetta by 1.56 and 6.92 in BLEU and CodeBLEU scores, while achieving performance comparable to DeepSeek-R1 and GPT-4.1. Our code is available at \url{https://github.com/kcxain/mupa}.
Changxin Ke, Rui Zhang 0040, Guangli Li, Yuanbo Wen 0001, Shuoming Zhang, Ruiyuan Xu, Jiaming Guo, Chenxi Wang 0005, Ling Li 0001, Qi Guo 0001, Yunji Chen
NeurIPS6
2025 AI Computing Systems for Large Language Models Training
Yuanbo Wen 0001, Han-Qi Lyu, Chang Liu 0021, Rui Zhang 0040, Xia-Qing Li, Chao Wang 0003, Zidong Du, Qi Guo 0001, Ling Li 0001, Xue-Hai Zhou, Yun-Ji Chen
J. Comput. Sci. Technol.2
2025 Efficient and Fast High-Performance Library Generation for Deep Learning Accelerators
abstract
The widespread adoption of deep learning accelerators (DLAs) underscores their pivotal role in improving the performance and energy efficiency of neural networks. To fully leverage the capabilities of these accelerators, exploration-based library generation approaches have been widely used to substantially reduce software development overhead. However, these approaches have been challenged by issues related to sub-optimal optimization results and excessive optimization overheads. In this paper, we proposeHeronto generate high-performance libraries of DLAs in an efficient and fast way. The key is automatically enforcing massive constraints through the entire program generation process and guiding the exploration with an accurate pre-trained cost model.Heronrepresents the search space as a constrained satisfaction problem (CSP) and explores the space via evolving the CSPs. Thus, the sophisticated constraints of the search space are strictly preserved during the entire exploration process. The exploration algorithm has the flexibility to engage in space exploration using either online-trained models or pre-trained models. Experimental results demonstrate thatHeronaveragely achieves 2.71$\times$speedup over three state-of-the-art automatic generation approaches. Also, compared to vendor-provided hand-tuned libraries,Heronachieves a 2.00$\times$speedup on average. When employing a pre-trained model,Heronachieves 11.6$\times$compilation time speedup, incurring a minor impact on execution time.
Jun Bi, Yuanbo Wen 0001, Xiaqing Li, Yongwei Zhao 0001, Enshuai Zhou, Xing Hu 0001, Zidong Du, Ling Li 0001, Huaping Chen 0001, Tianshi Chen 0002, Qi Guo 0001
IEEE Trans. Computers2
2025 VariPar: Variation-Aware Workload Partitioning in Chiplet-Based DNN Accelerators
abstract
Chiplet-based DNN accelerators have been extensively explored to save design and manufacturing costs. Previous works regard all chiplets as identical and employ uniform workload partitioning strategies. These workload partitioning strategies overlook various real-world factors that contribute to remarkable performance variations among chiplets, including manufacturing process variation, thermal condition, physical placement, and power supply condition. When considering these performance variations, a variation-aware workload partitioning can achieve superior performance. This paper introduces VariPar, a systematic framework to employ variation-aware partitioning strategy in chiplet-based DNN accelerators. VariPar models performance variations for each chiplet and partition workloads accordingly. VariPar includes a simulator with multi-factor variation modeling and a heuristic search engine to generate near-optimal partitioning within a reasonable time. Experiment results show that VariPar achieves 1.45× performance and 1.82× energy efficiency improvement on average when compared to uniform partitioning strategy.
Yongwei Zhao 0001, Mo Zou, Yang Liu 0466, Yifan Hao 0001, Xiaqing Li, Rui Zhang 0040, Yuanbo Wen 0001, Xing Hu 0001, Zidong Du, Qi Guo 0001, Tianshi Chen 0002
IEEE Trans. Comput. Aided Des. Integr. Circuits Syst.8
2025 Harmonia: A Unified Architecture for Efficient Deep Symbolic Regression
abstract
Symbolic regression (SR), the process of formulating a mathematical expression based on observed data points, is a fundamental task in artificial intelligence but is often hindered by its intense computational demands. Deep-learning-based SR methods (DSR) aim to alleviate these demands by breaking down the SR process into two stages: 1) neural network (NN) inference and 2) Broyden-Fletcher–Goldfarb-Shanno (BFGS) optimization. Although NN accelerators can expedite the NN stage, the performance of the BFGS optimization is compromised due to its poor performance for the variety of transcendental functions. Moreover, the distinct computational characteristics of NN inference and BFGS cause not only low hardware utilization but also significant area waste. To address these issues, we propose Harmonia, a unified architecture with the neural transcendental function unit (NTFU) and the Unified Array for efficient DSR. The NTFU utilizes the radial basis function network (RBFN) as a universal approximator for various transcendental functions, which significantly reduces the heavy transcendental function computation cost. We further propose an efficient training algorithm called random nonlinear optimization (RNO) to obtain a lightweight RBFN without accuracy loss. Moreover, Harmonia supports configurable dataflow which integrates the two computing stages into the Unified Array. Experimental results show that Harmonia achieves hardware utilization of 83.83%, on average. Compared to the GPU baseline, Harmonia achieves$4.8\times $speedup and$47.6\times $energy saving, alongside considerable low area cost.
Tianyun Ma, Yuanbo Wen 0001, Xinkai Song, Pengwei Jin, Husheng Han, Ziyuan Nan, Zhongkai Yu, Shaohui Peng, Yongwei Zhao 0001, Huaping Chen 0001, Zidong Du, Xing Hu 0001, Qi Guo 0001
IEEE Trans. Comput. Aided Des. Integr. Circuits Syst.2
2024 TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing
abstract
Heterogeneous collaborative computing with NPU and CPU has received widespread attention due to its substantial performance benefits. To ensure data confidentiality and integrity during computing, Trusted Execution Environments (TEE) is considered a promising solution because of its comparatively lower overhead. However, existing heterogeneous TEE designs are inefficient for collaborative computing due to fine and different memory granularities between CPU and NPU. 1) The cacheline granularity of CPU TEE intensifies memory pressure due to its extra memory access, and 2) the cacheline granularity MAC of NPU escalates the pressure on the limited memory storage. 3) Data transfer across heterogeneous enclaves relies on the transit of non-secure regions, resulting in cumbersome re-encryption and scheduling.
Husheng Han, Xinyao Zheng, Yuanbo Wen 0001, Yifan Hao 0001, Erhu Feng, Ling Liang 0003, Jianan Mu, Xiaqing Li, Tianyun Ma, Pengwei Jin, Xinkai Song, Zidong Du, Qi Guo 0001, Xing Hu 0001
ASPLOS (4)3
2024 Revisiting Automatic Pipelining: Gate-level Forwarding and Speculation
abstract
Pipelining is a widely applied micro-architectural performance optimization and requires non-trivial designs for better execution throughput. The key to pipeline throughput optimization is to resolve data hazards caused by read-after-write (RAW) dependencies, which are traditionally tackled by forwarding and speculation to avoid pipeline stalls. However, existing approaches are conducted based on high-level dataflow analysis, with potential loss of optimization opportunities for lack of analysis of the netlist structures.
Shuyao Cheng, Chongxiao Li, Zidong Du, Rui Zhang 0040, Xing Hu 0001, Xiaqing Li, Guanglin Xu, Yuanbo Wen 0001, Qi Guo 0001
DAC8
2024 AutoOS: Make Your OS More Powerful by Exploiting Large Language Models
abstract
With the rapid development of Artificial Intelligence of Things (AIoT), customizing and optimizing operating system (OS) kernel configurations for various AIoT application scenarios is crucial for maximizing system performance. However, existing approaches falter due to the overwhelming problem complexity (i.e., over 15,000 configuration options in the Linux kernel), together with the huge evaluation costs and error-prone options that may result in OS boot-up failure, which all make it an unresolved problem to optimize the Linux kernel automatically. In this paper, we introduce AutoOS, a novel framework exploiting Large Language Models for customizing and optimizing OS kernel configurations automatically for various AIoT application scenarios.Inspired by the inherently directory-structured kernel configuration process, we first formulate our research problem as optimizing on a dynamic tree. We then propose a novel framework integrating a state machine-based traversal algorithm as the observe-prune-propose-act-correct loop, which can effectively refine the optimization space and ensure a successful OS boot-up.Experimental results show that AutoOS can automatically customize and optimize the OS kernel configurations without human effort. More importantly, AutoOS even achieves better performance by up to 25% than vendor-provided configuration.
Huilai Chen, Yuanbo Wen 0001, Limin Cheng, Shouxu Kuang, Ling Li 0001, Rui Zhang 0040, Xinkai Song, Wei Li 0008, Qi Guo 0001, Yunji Chen
ICML2
2024 Prompt-based Visual Alignment for Zero-shot Policy Transfer
abstract
Overfitting in RL has become one of the main obstacles to applications in reinforcement learning(RL). Existing methods do not provide explicit semantic constrain for the feature extractor, hindering the agent from learning a unified cross-domain representation and resulting in performance degradation on unseen domains. Besides, abundant data from multiple domains are needed. To address these issues, in this work, we propose prompt-based visual alignment (PVA), a robust framework to mitigate the detrimental domain bias in the image for zero-shot policy transfer. Inspired that Visual-Language Model (VLM) can serve as a bridge to connect both text space and image space, we leverage the semantic information contained in a text sequence as an explicit constraint to train a visual aligner. Thus, the visual aligner can map images from multiple domains to a unified domain and achieve good generalization performance. To better depict semantic information, prompt tuning is applied to learn a sequence of learnable tokens. With explicit constraints of semantic information, PVA can learn unified cross-domain representation under limited access to cross-domain data and achieves great zero-shot generalization ability in unseen domains. We verify PVA on a vision-based autonomous driving task with CARLA simulator. Experiments show that the agent generalizes well on unseen domains under limited access to multi-domain data.
Haihan Gao, Rui Zhang 0040, Qi Yi, Hantao Yao, Haochen Li 0002, Jiaming Guo, Shaohui Peng, Yunkai Gao 0001, QiCheng Wang, Xing Hu 0001, Yuanbo Wen 0001, Zidong Du, Ling Li 0001, Qi Guo 0001, Yunji Chen
ICML11
2024 Cambricon-D: Full-Network Differential Acceleration for Diffusion Models
abstract
Diffusion models have made significant progress in current image generation tasks, thus becoming a prominent area of research. Diffusion models necessitate repetitive iterations on minimally altered input data across timesteps, each timestep requiring the recalculation of the entire model, resulting in a remarkable computational redundancy and substantial hardware expenditures.Performing differential computing on input data seems to be a feasible approach for addressing such computational redundancy and improving hardware efficacy. However, non-linear operations (particularly activation functions) necessitate the merging of deltas (i.e., differential values) with raw inputs repeatedly to ensure computational correctness, leading to significant memory access for loading raw inputs, which fragmentedly blocks the forwarding of deltas throughout the network and undermines performance.To solve this problem, we propose Cambricon-D, a fullnetwork differential computing architecture with concise memory access. While maintaining the computational efficiency brought by differential computing, Cambricon-D employs a sign-mask dataflow, which requires only the loading of 1-bit signs (instead of large bitwidth raw inputs), thereby facilitating the seamless forwarding of deltas and effectively mitigating memory access overheads. Experimental results show that, compared to Diffy, Cambricon-D’s dataflow reduces 66% ~ 82% off-chip memory access. In total, Cambricon-D achieves 1.46× ~ 2.38× speedup over A100 on various diffusion models with different resolutions.
Weihao Kong, Yifan Hao 0001, Qi Guo 0001, Yongwei Zhao 0001, Xinkai Song, Xiaqing Li, Mo Zou, Zidong Du, Rui Zhang 0040, Chang Liu 0021, Yuanbo Wen 0001, Pengwei Jin, Xing Hu 0001, Wei Li 0008, Zhiwei Xu 0002, Tianshi Chen 0002
ISCA11
2024 Cambricon-C: Efficient 4-Bit Matrix Unit via Primitivization
abstract
Deep learning trends to use low precision numeral formats to cope with the ever-growing model sizes. For example, the large language model LLaMA2 has been widely deployed in 4-bit precision. With larger models and fewer unique values caused by low precision, an increasing proportion of arithmetic in matrix multiplication is repeating. Although discussed in prior works, such value redundancy has not been fully exploited, and the cost to leverage the value redundancy often offsets any advantages. In this paper, we propose to primitivize the matrix multiplication, that is decomposing it down to the 1-ary successor function (a.k.a. counting) to merge repeating arithmetic. We revisited various techniques to propose Cambricon-C SA, a 4-bit primitive matrix multiplication unit that doubles the energy efficiency over conventional systolic arrays. Experimental results show that Cambricon-C SA can achieve$\mathbf{1}.\mathbf{95}\times$energy efficiency improvement compared with MAC-based systolic array.
Yongwei Zhao 0001, Yifan Hao 0001, Yuanbo Wen 0001, Yuntao Dai, Xiaqing Li, Yang Liu 0466, Rui Zhang 0040, Mo Zou, Xinkai Song, Xing Hu 0001, Zidong Du, Huaping Chen 0001, Qi Guo 0001, Tianshi Chen 0002
MICRO4
2023 Heron: Automatically Constrained High-Performance Library Generation for Deep Learning Accelerators
abstract
Deep Learning Accelerators (DLAs) are effective to improve both performance and energy efficiency of compute-intensive deep learning algorithms. A flexible and portable mean to exploit DLAs is using high-performance software libraries with well-established APIs, which are typically either manually implemented or automatically generated by exploration-based compilation approaches. Though exploration-based approaches significantly reduce programming efforts, they fail to find optimal or near-optimal programs from a large but low-quality search space because the massive inherent constraints of DLAs cannot be accurately characterized.
Jun Bi, Qi Guo 0001, Xiaqing Li, Yongwei Zhao 0001, Yuanbo Wen 0001, Enshuai Zhou, Xing Hu 0001, Zidong Du, Ling Li 0001, Huaping Chen 0001, Tianshi Chen 0002
ASPLOS (3)5
2023 BALTO: fast tensor program optimization with diversity-based active learning
Jun Bi, Xiaqing Li, Qi Guo 0001, Rui Zhang 0040, Yuanbo Wen 0001, Xing Hu 0001, Zidong Du, Xinkai Song, Yifan Hao 0001, Yunji Chen
ICLR5
2023 Cambricon-R: A Fully Fused Accelerator for Real-Time Learning of Neural Scene Representation
abstract
Neural scene representation (NSR) initiates a new methodology of encoding a 3D scene with neural networks by learning from dozens of photos taken from different camera positions. NSR not only achieves significant improvement in the quality of novel view synthesis and 3D reconstruction but also reduces the camera cost from the expensive laser cameras to the cheap color cameras on the shelf. However, performing 3D scene encoding using NSR is far from real-time due to the extremely low hardware utilization (only utilization of hardware peak performance), which greatly limits its applications in real-time AR/VR interactions
Xinkai Song, Yuanbo Wen 0001, Xing Hu 0001, Tianbo Liu 0006, Haoxuan Zhou, Husheng Han, Tian Zhi, Zidong Du, Wei Li 0008, Rui Zhang 0040, Chen Zhang 0001, Lin Gao 0004, Qi Guo 0001, Tianshi Chen 0002
MICRO2
2023 Emergent Communication for Rules Reasoning
abstract
Research on emergent communication between deep-learning-based agents has received extensive attention due to its inspiration for linguistics and artificial intelligence. However, previous attempts have hovered around emerging communication under perception-oriented environmental settings, that forces agents to describe low-level perceptual features intra image or symbol contexts. In this work, inspired by the classic human reasoning test (namely Raven's Progressive Matrix), we propose the Reasoning Game, a cognition-oriented environment that encourages agents to reason and communicate high-level rules, rather than perceived low-level contexts. Moreover, we propose 1) an unbiased dataset (namely rule-RAVEN) as a benchmark to avoid overfitting, 2) and a two-stage curriculum agent training method as a baseline for more stable convergence in the Reasoning Game, where contexts and semantics are bilaterally drifting. Experimental results show that, in the Reasoning Game, a semantically stable and compositional language emerges to solve reasoning problems. The emerged language helps agents apply the extracted rules to the generalization of unseen context attributes, and to the transfer between different context attributes or even tasks.
Yifan Hao 0001, Rui Zhang 0040, Enshuai Zhou, Zidong Du, Xishan Zhang, Xinkai Song, Yuanbo Wen 0001, Yongwei Zhao 0001, Xuehai Zhou, Jiaming Guo, Qi Yi, Shaohui Peng, Ruizhi Chen, Qi Guo 0001, Yunji Chen
NeurIPS8
2023 ANPL: Towards Natural Programming with Interactive Decomposition
abstract
Though LLMs are capable of generating plausible programs, it’s challenging to interact with the LLMs further to revise the program, especially if the user’s specific requirements are different from the initial proposal. In this paper, we introduce ANPL, an interactive programming system that ensures users can always refine the generated code towards their specific programmatic intents via structured decompositions. Borrowing the paradigm of sketching from program synthesis, an ANPL program consists of a set of input-outputs that it must satisfy, a “sketch” — control/data flow expressed in precise code (e.g. Python), and “holes” — sub-modules to be implemented by the LLM specified with natural language. The user revises an ANPL program by either modifying the sketch, changing the language used to describe the holes, or providing additional input-outputs to a particular hole, turning it into a sub-ANPL program that can be solved recursively. This workflow allows the users to offload programming burdens to the LLM as much as possible while retaining the ability to pinpoint and resolve bugs locally, without exposing the rest of the program to the LLM. We deploy ANPL on the Abstraction and Reasoning Corpus (ARC), a set of unique tasks that are challenging for state-of-the-art AI systems, showing it outperforms baseline programming systems that (a) without the ability to decompose tasks interactively and (b) without the guarantee that the modules can be correctly composed together. Additional evaluations on APPS, HumanEval, and real-world programming tasks have validated that the ANPL framework is applicable to multiple programming domains. We release the ANPL solutions to the ARC tasks as a dataset, providing insights into how humans decompose novel tasks programmatically.
Ziyuan Nan, Xing Hu 0001, Pengwei Jin, Shaohui Peng, Yuanbo Wen 0001, Rui Zhang 0040, Zidong Du, Qi Guo 0001, Yewen Pu, Yunji Chen
NeurIPS6
2022 BabelTower: Learning to Auto-parallelized Program Translation
abstract
GPUs have become the dominant computing platforms for many applications, while programming GPUs with the widely-used CUDA parallel programming model is difficult. As sequential C code is relatively easy to obtain either from legacy repositories or by manual implementation, automatically translating C to its parallel CUDA counterpart is promising to relieve the burden of GPU programming. However, because of huge differences between the sequential C and the parallel CUDA programming model, existing approaches fail to conduct the challenging auto-parallelized program translation. In this paper, we propose a learning-based framework, i.e., BabelTower, to address this problem. We first create a large-scale dataset consisting of compute-intensive function-level monolingual corpora. We further propose using back-translation with a discriminative reranker to cope with unpaired corpora and parallel semantic conversion. Experimental results show that BabelTower outperforms state-of-the-art by 1.79, 6.09, and 9.39 in terms of BLEU, CodeBLEU, and specifically designed ParaBLEU, respectively. The CUDA code generated by BabelTower attains a speedup of up to 347x over the sequential C code, and the developer productivity is improved by at most 3.8x.
Yuanbo Wen 0001, Qi Guo 0001, Xiaqing Li, Jianxing Xu, Yanlin Tang, Yongwei Zhao 0001, Xing Hu 0001, Zidong Du, Ling Li 0001, Chao Wang 0003, Xuehai Zhou, Yunji Chen
ICML1
2022 Enabling One-Size-Fits-All Compilation Optimization for Inference Across Machine Learning Computers
abstract
Machine Learning Computers (MLCs) with tensor functional units (e.g., NVIDIA's Tensor Core, Google's TPU and Habana's Tensor Processor Core) have emerged significantly over recent years. The broad diversity of MLCs makes it hard to deploy machine learning workloads with optimized performance. Though deep learning compilers (e.g., TVM) are effective to produce optimized code for different hardware back-ends, when deploying to a new MLC, it is tedious to implement platform-specific compilation optimizations by thoroughly understanding system/architectural details. To address this problem, we propose a holistic approach to achieve one-size-fits-all compilation optimization across different MLCs or inference. The key observation is that diverse MLCs share multiple key architectural characteristics for tensor processing, which can be generalized for conducting cross-platform compilation optimizations. Concretely, we propose the Tensor Abstract Machine (TAM), which features such common architectural characteristics, as the abstraction of a broad range of MLCs. To leverage architectural characteristics of the TAM, we propose the Tensor Scheduling Language (TSL) consisting of tensor computation description and tensor scheduling primitives for implementing operations with portable optimization. Experimental results demonstrate that the code generated from the same optimization schedule achieves 1.05x to 2.05x better performance than hand-tuned libraries and deep learning compilers across different platforms.
Yuanbo Wen 0001, Qi Guo 0001, Zidong Du, Jianxing Xu, Xing Hu 0001, Wei Li 0008, Rui Zhang 0040, Chao Wang 0003, Xuehai Zhou, Tianshi Chen 0002
IEEE Trans. Computers1
2020 Addressing Irregularity in Sparse Neural Networks Through a Cooperative Software/Hardware Approach
abstract
Neural networks have become the dominant algorithms rapidly as they achieve state-of-the-art performance in a broad range of applications such as image recognition, speech recognition, and natural language processing. However, neural networks keep moving toward deeper and larger architectures, posing a great challenge to hardware systems due to the huge amount of data and computations. Although sparsity has emerged as an effective solution for reducing the intensity of computation and memory accesses directly, irregularity caused by sparsity (including sparse synapses and neurons) prevents accelerators from completely leveraging the benefits, i.e., it also introduces costly indexing module in accelerators. In this article, we propose a cooperative software/hardware approach to address the irregularity of sparse neural networks efficiently. Initially, we observe the local convergence, namely larger weights tend to gather into small clusters during training. Based on that key observation, we propose a software-based coarse-grained pruning technique to reduce the irregularity of sparse synapses drastically. The coarse-grained pruning technique, together with local quantization, significantly reduces the size of indexes and improves the network compression ratio. We further design a multi-core hardware accelerator, Cambricon-SE, to address the remaining irregularity of sparse synapses and neurons efficiently. The novel accelerator have three key features: 1) selector modulesto filter unnecessary synapses and neurons, 2) compress/decompress modules for exploiting the sparsity in data transmission (which is rarely studied in previous work), and 3) a multi-core architecture with elevated throughput to meet the real-time processing requirement. Compared against a state-of-the-art sparse neural network accelerator, our accelerator is 1.20x and 2.72x better in terms of performance and energy efficiency, respectively. Moreover, for real-time video analysis tasks, Cambricon-SE can process 1080p video at the speed of 76.59 fps.
Tian Zhi, Xuda Zhou, Zidong Du, Qi Guo 0001, Shaoli Liu, Bingrui Wang, Yuanbo Wen 0001, Chao Wang 0003, Xuehai Zhou, Ling Li 0001, Tianshi Chen 0002, Ninghui Sun, Yunji Chen
IEEE Trans. Computers8