GTA: Generating high-performance tensorized program with dual-task scheduling

Anxing Xie

doi:10.1016/j.sysarc.2025.103359

Return

GTA: Generating high-performance tensorized program with dual-task scheduling

2025-03-01

PRE

DOI：10.1016/j.sysarc.2025.103359

Original

Original request for help

Save

Abstract

En 中文

Generating high-performance tensorized programs for deep learning accelerators (DLAs) is crucial for ensuring the efficient execution of deep neural networks. But, producing such programs for different operators across various DLAs is notoriously challenging. Existing methods utilize hardware abstraction to represent acceleration intrinsics, enabling end-to-end automated exploration of the intrinsics mapping space. However, their limited search space and inefficient exploration strategies often result in suboptimal tensorized programs and significant search time overhead. In this paper, we propose GTA, a framework designed to generate high-performance tensorized programs for DLAs. Unlike existing deep learning compilers, we first coordinate intrinsic-based mapping abstraction with rule-based program generation strategy, followed by the application of resource-constrained rules to eliminate ineffective tensor program candidates from the search space. Second, we employ a dual-task scheduling strategy to allocate tuning resources across multiple subgraphs of deep neural networks and their mapping candidates. As a result, GTA can find high-performance tensor programs that are outside the search space of existing state-of-the-art methods. Our experiments show that GTA achieves an average speedup of more than 1.88x over AMOS and 2.29x over Ansor on NVIDIA GPU with Tensor Core, as well as 1.49x over Ansor and 2.76x over PyTorch on CPU with AVX512.

Keywords:

Mapping

Code generation

Compiler optimization

Tensor computation