# Tensor Program Optimization for the RISC-V Vector Extension Using Probabilistic Programs

arXiv: https://arxiv.org/abs/2507.01457 · HTML: https://mlsanity.com/paper/2507.01457/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 19 references that are in the map.
- First posted on arXiv 2 Jul 2025.

## Builds on (most cited first, 11 of 11)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://mlsanity.com/paper/1511.06434/) (2015)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [MLPerf Tiny Benchmark](https://mlsanity.com/paper/2106.07597/) (2021)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [A "New Ara" for Vector Computing: An Open Source Highly Efficient RISC-V V 1.0 Vector Processor Design](https://mlsanity.com/paper/2210.08882/) (2022)
- [Spatz: Clustering Compact RISC-V-Based Vector Units to Maximize Computing Efficiency](https://mlsanity.com/paper/2309.10137/) (2023)
- [Instruction Scheduling in the Saturn Vector Unit](https://mlsanity.com/paper/2412.00997/) (2024)

## Built on by (most cited first, 1 of 1)

- [TASTE: Throughput-Aware Batch Size Tuning for On-Device Edge Learning](https://mlsanity.com/paper/2609.07444/) (2026)

## Abstract (arXiv)

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.