面向TPU粗粒度指令的自动张量化方法

Automatic Tensorization for TPU Coarse-grained Instructions

张量化是通过调用硬件特定指令对张量运算进行加速的过程。TPU支持多种粗粒度指令,可表示神经网络级别的算子,且没有明确的运算规模限制。现有张量化方法对于粗粒度指令需要手写大量的IR匹配片段,且难以实现灵活的双缓存(ping-pong buffer)形式的指令并行优化,不利于扩展至TPU场景。为此,提出了一种面向TPU粗粒度指令的自动张量化方法——Tir2TPU。首先,基于TensorIR抽象语法树的分析对运算程序进行指令替换。其次,设计了一种模拟硬件行为的并行模型以实现指令并行优化。最后,构建了基于TPU硬件特征的程序调度空间以实现快速自动调优。实验对矩阵乘法等5种机器学习模型中常用的算子进行了性能评估。实验结果表明,Tir2TPU自动优化生成的算子与TPU自有编译器相比可取得最高3.1倍、平均1.78倍的运算加速,并且可取得平均90%的手工优化性能。

syntax tree.Secondly,it also utilizes a parallel model that simulates hardware behavior to generate parallel instruction flow.Finally,Tir2TPU combines a hardware-centric schedule space based on TPU features,which greatly accelerates auto-tuning process.The performance of Tir2TPU is evaluatedon 5 commonly used operators in machine learning models.Experimental results show that Tir2TPU can achieve up to 3×and an average of 1.78×speedup compared to TPU’s compiler,and consistently delivers 90%performance compared to manually optimized operators.