Optimizing Tensor Train Decomposition in DNNs for RISC-V Architectures Using Design Space Exploration and Compiler Optimizations

TL;DR

This work tackles the challenge of deploying DNNs on resource-constrained RISC-V devices by integrating Tensor Train Decomposition-based low-rank factorization with a design-space exploration framework and RISC-V-tailored compiler optimizations. It introduces a layer-wise LRF DSE methodology that first prunes the TT decomposition design space via shape alignment and FLOPs/memory heuristics, then further narrows candidates by inferred latency and hardware scalability, and finally applies aggressive compiler optimizations to accelerate the custom T3F Einsum layers. Experimental results show TT-decomposed FC layers run, on average, 3x faster than IREE and 8x faster than Pluto for the same factorized models, with end-to-end speedups averaging around 12x versus uncompressed IREE. The combination of design-space pruning and compiler-level optimizations enables efficient deployment of DNNs on edge devices powered by RISC-V, with strong potential for broader applicability to other architectures.

Abstract

Deep neural networks (DNNs) have become indispensable in many real-life applications like natural language processing, and autonomous systems. However, deploying DNNs on resource-constrained devices, e.g., in RISC-V platforms, remains challenging due to the high computational and memory demands of fully connected (FC) layers, which dominate resource consumption. Low-rank factorization (LRF) offers an effective approach to compressing FC layers, but the vast design space of LRF solutions involves complex trade-offs among FLOPs, memory size, inference time, and accuracy, making the LRF process complex and time-consuming. This paper introduces an end-to-end LRF design space exploration methodology and a specialized design tool for optimizing FC layers on RISC-V processors. Using Tensor Train Decomposition (TTD) offered by TensorFlow T3F library, the proposed work prunes the LRF design space by excluding first, inefficient decomposition shapes and second, solutions with poor inference performance on RISC-V architectures. Compiler optimizations are then applied to enhance custom T3F layer performance, minimizing inference time and boosting computational efficiency. On average, our TT-decomposed layers run 3x faster than IREE and 8x faster than Pluto on the same compressed model. This work provides an efficient solution for deploying DNNs on edge and embedded devices powered by RISC-V architectures.

Figures (18)

• Figure 1: Parameters and FLOPs percentage of FC and non-FC parts for various DNN models
• Figure 3: Visual illustration of Listing 2
• Figure 4: The proposed DSE methodology
• Figure 5: Scatterplots illustrating FLOPs and memory for all shape permutations across three distinct configurations for a CNN layer with $(M,N) = (9216,4096)$. The proposed aligned shape permutation is highlighted in red
• Figure 6: Scatterplots illustrating FLOPs and Memory for all shape permutations across three distinct configurations for a LLM layer with $(M,N) = (2048,2048)$. The proposed aligned shape permutation is highlighted in red

Theorems & Definitions (9)

• Proposition 1
• proof
• Remark 1
• Proposition 2
• proof
• Definition 1
• Proposition 3
• proof
• Proposition 4