HAOT: Heterogeneous Hardware-Aware Tensor Computation Optimization Framework via Transformer

Efficient execution of tensor computations on different hardware, such as CPUs, GPUs, and spatial accelerators, poses substantial challenges due to divergent memory hierarchies, compute models, and architectural constraints. However, existing auto-tuning approaches typically fail to integrate hardware-aware formulations, leading to inefficient search processes and suboptimal utilization of hardware resources. To address these challenges, we introduce HAOT, a novel hardware-aware tensor computation optimization framework that combines a learning-based scheduling policy with a hardware analytical model. Unlike traditional methods, HAOT not only adapts the learning-based scheduling strategy but also dynamically constrains the search space based on the availability of hardware resources, ensuring both the feasibility and optimal use of resources. Experiments across a range of hardware platforms show that HAOT achieves up to 1.4× average speedup on individual tensor operations and a 1.8× improvement in end-to-end deep learning workloads, outperforming state-of-the-art baselines. Ablation studies demonstrate the complementary roles of the Transformer-based policy model and the hardware analysis model, each contributing to generating high-quality schedule decisions. Furthermore, HAOT exhibits superior sample efficiency, requiring significantly fewer search trials to converge to the optimal schedule.