Compound operators, such as Log_softmax and RMSNorm, have been widely studied to enhance performance in deep neural networks (DNNs). Nonetheless, these operators often suffer from high hardware adaptation costs and limited optimization effects. AI compilers optimize them through operator splitting and successive operator fusion strategies. However, prior studies indiscriminately split all compound operators and failed to fully explore the fusion search space, incurring inefficient fusion schemes. To overcome these limitations, we propose a Co-Optimization Framework for Operator Splitting and Fusion (CoSF). In the operator splitting phase, we analyze memory reuse levels among operators and classify the compound operators into three types, according to their data locality. Then, we propose a fusion-aware splitting strategy. For each type of compound operator, it evaluates the successive fusion benefits after splitting the compound operator and automatically generates operator splitting strategies. In the operator fusion phase, to reduce the massive computation graph resulting from operator splitting, we propose a dominator tree-based graph partitioning algorithm to efficiently partition the computation graph. We then employ dynamic programming for each partitioned subgraph to generate an optimized fusion strategy. Finally, we propose a hardware-agnostic evaluation model to select the most effective fusion solution from multiple candidates. Experimental results demonstrate that CoSF achieves a 1.3–3.4 \(\times \) speedup on GPU and a 1.59–3.93 \(\times \) speedup on CPU compared to TVM, Pytorch, and TF-XLA.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

CoSF: A \(\underline{Co}\) -Optimization Framework for Operator \(\underline{S}\) plitting and  \(\underline{F}\) usion

  • Wei Li,
  • Ao Ren,
  • Qingqiu Lan,
  • Haining Fang,
  • Zhenyu Wang,
  • Yujuan Tan,
  • Kan Zhong,
  • Duo Liu

摘要

Compound operators, such as Log_softmax and RMSNorm, have been widely studied to enhance performance in deep neural networks (DNNs). Nonetheless, these operators often suffer from high hardware adaptation costs and limited optimization effects. AI compilers optimize them through operator splitting and successive operator fusion strategies. However, prior studies indiscriminately split all compound operators and failed to fully explore the fusion search space, incurring inefficient fusion schemes. To overcome these limitations, we propose a Co-Optimization Framework for Operator Splitting and Fusion (CoSF). In the operator splitting phase, we analyze memory reuse levels among operators and classify the compound operators into three types, according to their data locality. Then, we propose a fusion-aware splitting strategy. For each type of compound operator, it evaluates the successive fusion benefits after splitting the compound operator and automatically generates operator splitting strategies. In the operator fusion phase, to reduce the massive computation graph resulting from operator splitting, we propose a dominator tree-based graph partitioning algorithm to efficiently partition the computation graph. We then employ dynamic programming for each partitioned subgraph to generate an optimized fusion strategy. Finally, we propose a hardware-agnostic evaluation model to select the most effective fusion solution from multiple candidates. Experimental results demonstrate that CoSF achieves a 1.3–3.4 \(\times \) speedup on GPU and a 1.59–3.93 \(\times \) speedup on CPU compared to TVM, Pytorch, and TF-XLA.