用概率编程优化RISC-V向量指令,让AI模型跑得更快更省空间。
Tensor Program Optimization for the RISC-V Vector Extension Using Probabilistic Programs
- 基于TVM的元调度框架,用概率程序自动调优AI算子在RISC-V向量单元上的实现。
- 相比GCC自动向量化平均快46%,比muRISCV-NN快29%,代码体积更小。
- 已在FPGA和商用RISC-V芯片上验证,适合嵌入式AI部署,开源可扩展。
RISC-V为从嵌入式设备到高性能计算集群的应用提供了灵活可扩展的平台,其向量扩展(RVV)在加速AI工作负载方面备受关注。然而,非专家程序员要高效利用RISC-V CPU的向量单元,通常依赖编译器的自动向量化功能或手动编写库(如muRISCV-NN)。现有自调优框架尚未与RISC-V RVV集成,严重限制了复杂AI工作负载的高效部署。本文提出一种基于TVM编译器的工作流,将RISC-V RVV扩展集成至TVM的MetaSchedule框架——一个用于张量操作调优的概率程序框架。我们在FPGA上实现了多个RISC-V SoC,并在其中对广泛AI工作负载进行调优。结果表明,相比GCC自动向量化,本方案平均降低46%执行延迟;相比muRISCV-NN,提升29%。此外,生成的二进制文件代码内存占用更小,更适合嵌入式设备。我们还在商用支持RVV 1.0的RISC-V SoC上评估,发现本方案找到的映射平均比LLVM快35%。项目已开源,欢迎社区拓展至其他RISC-V扩展。
原文摘要 · Abstract (English)
RISC-V provides a flexible and scalable platform for applications ranging from embedded devices to high-performance computing clusters. Particularly, its RISC-V Vector Extension (RVV) becomes of interest for the acceleration of AI workloads. But writing software that efficiently utilizes the vector units of RISC-V CPUs without expert knowledge requires the programmer to rely on the autovectorization features of compilers or hand-crafted libraries like muRISCV-NN. Smarter approaches, like autotuning frameworks, have been missing the integration with the RISC-V RVV extension, thus heavily limiting the efficient deployment of complex AI workloads. In this paper, we present a workflow based on the TVM compiler to efficiently map AI workloads onto RISC-V vector units. Instead of relying on hand-crafted libraries, we integrated the RVV extension into TVM's MetaSchedule framework, a probabilistic program framework for tensor operation tuning. We implemented different RISC-V SoCs on an FPGA and tuned a wide range of AI workloads on them. We found that our proposal shows a mean improvement of 46% in execution latency when compared against the autovectorization feature of GCC, and 29% against muRISCV-NN. Moreover, the binary resulting from our proposal has a smaller code memory footprint, making it more suitable for embedded devices. Finally, we also evaluated our solution on a commercially available RISC-V SoC implementing the RVV 1.0 Vector Extension and found our solution is able to find mappings that are 35% faster on average than the ones proposed by LLVM. We open-sourced our proposal for the community to expand it to target other RISC-V extensions.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。