Vortex: Efficient Sample-Free Dynamic Tensor Program Optimization via Hardware-aware
Strategy Space Hierarchization

Tensor programs act as an operator-level abstraction, they are widely used across different tasks in neural network computation (pytorch, ; tensorflow, ; tvm, ; ansor, ). Conventional tensor programs inherently incorporate static shape information as an integral part of their input. On the contrary, dynamic-shape tensor programs enable processing tensor programs with unknown shapes (dynamic_survey, ). Dynamic-shape tensor programs have arisen in response to a twofold demand, driven by intrinsic data format dynamism and system execution and scheduling strategies.

The intrinsic data format dynamism is one of the fundamental driving forces demanding the integration of dynamic-shape tensor programs. For example, in natural language processing (NLP) (bert, ), the inherent variability in sequence lengths is a characteristic feature that traditional static-shape tensor programs find challenging to accommodate. The diverse lengths of sentence tasks necessitate a flexible framework capable of dynamically adapting to varying input sizes. In computer vision (CV) tasks, conventional methods rely on fixed image sizes for tensor program input, which limits flexibility (mobilenet, ). Recent developments (fast-rcnn, ; scale_fact_dect, ; dynamic_stride_net, ; dynamic_survey, ) have introduced dynamic-shape tensors as input, enhancing support for more advanced detection and tracking tasks. Additionally, in the field of graph neural network (GNN) (gnn_survey, ; GNNAdvisor, ), the dynamic nature of graph structures, marked by varying numbers of vertices and edges, necessitates the use of dynamic-shape tensor programs. These examples highlight the essential need for systems that seamlessly manage the intrinsic variability in diverse data formats.

System execution and scheduling strategies further underscore the importance of dynamic-shape tensor programs. For example, dynamic adjustment of batch sizes in the execution of neural networks introduces variability that demands adaptability in the underlying tensor program (dvabatch, ; lazybatch, ; orca, ). The ability to efficiently handle varying batch sizes is crucial for optimizing resource utilization and achieving optimal performance in real-world applications.

Many optimizations have been proposed for dynamic-shape tensor programs, all following a similar workflow (nimble, ; dietcode, ). We term this methodology the sample-driven approach. In this subsection, we first demonstrate its workflow and then discuss the inherent limitations.

As depicted in Figure 2, the current approach typically employs a sample list to depict the dynamic-shape parameters, combining the preknown static shape and the tensor program to staticize the dynamic-shape tensor program into a shape-generic search space. An auto-tuning module, adapted from the static-shape compiler(tvm, ; ansor, ; tensorir, ), is subsequently utilized. The auto-tuning module takes the tensor program with unspecified tile parameters, proceeding to investigate the high-performance tiling configurations within the comprehensive search space. This process produces fine-tuned micro-kernels for each input sample. At runtime, a decision-tree-based selector is employed to choose the appropriate micro-kernel based on the runtime shape. The kernel constructor finalizes the process by setting kernel launch parameters and incorporating padding to extend the applicability of these micro-kernels to runtime shapes not included in the sample list, utilizing the pre-compiled micro-kernels for runtime execution.

However, this approach faces significant limitations due to its reliance on a predetermined sample list for dynamic-shape parameters. This limitation restricts the compiler’s flexibility and overall functionality. In situations where input shapes are not included in the sample list, this sample-driven dynamic compilation optimization does not consistently provide high-performance computing support. This is a critical shortfall, as the diversity and variability of real-world data often extend beyond the scope of the predefined sample list.

To empirically validate this limitation, we conduct an experiment targeting the first general matrix multiply (GEMM) operation of the Bert model (bert, ). This GEMM operation entails the multiplication of two matrices, $A$ and $B$. In this context, $M$ denotes the number of rows in matrix $A$, computed as the product of batch size and sequence length. $N$ represents the number of columns in matrix $B$, fixed at 768, and $K$ corresponds to the number of columns in matrix $A$ (and rows in matrix $B$), fixed at 2304. Utilizing DietCode’s default sample configuration, tests are performed with a fixed batch size of 16 and sequence length varying from 5 to 128 in the increment step of 19.

Figure 3 shows that DietCode exhibits a significant performance discrepancy for dynamic parameters not included in the sample list, as compared to the results achieved with the vendor’s library cuBLAS (cublas, ). This is due to the absence of specifically optimized micro-kernels for these shapes and the increased inefficiency from padding loss. Moreover, this method’s rigidity becomes evident when considering the need to modify the sample list to accommodate various computational scenarios. Such modifications necessitate re-tuning of the system, which worsens its inflexibility.

The mismatch between runtime-used and offline-sampled shapes diminishes the efficiency of traditional sample-driven methodology in dynamic-shape compilation. Such methodologies not only fall short in delivering high performance but also result in considerable tuning overhead. However, we observe that generating the sample list for tuning micro-kernels is not mandatory for supporting the execution of a dynamic tensor program. Sample-driven methodology treats the hardware platform as a black box, where micro-kernels are only tuned based on performance feedback of sampled shapes. This approach overlooks the rich vein of prior knowledge available within the hardware itself.

Figure 4 demonstrates the architecture of current mainstream hardware deployments, such as central processing units (CPUs) (intel_xeon_whitepaper, ) and graphics processing units (GPUs) (a100_whitepaper, ). There exist inherent similarities among these hardware, each of which has a hierarchical structure. This hierarchy is distinctly multi-level, wherein each level comprises a predetermined quantity of computational or storage units. For instance, each CPU core has its own L1 cache and ALUs, while all the CPU cores share the L2/L3 cache and DRAM. As for GPUs, each streaming multiprocessor (SM) also has its own computing units (CUDA cores, Tensor Cores) and L1 Cache, while all SMs share the L2 cache and DRAM. At each level, the resources available for executing a dynamic tensor program operator are constrained by the inherent limitations of the hardware design.

To investigate the impact of hardware limits on kernel performance, we experiment by generating configurations with different resource usages for the common-used matrix multiplication. We collect these configurations during Ansor (ansor, )’s tuning process. Figure 5 shows the floating point operations per second (FLOPS) and the corresponding resource usage of the collected configurations. As observed, the operator performance seriously declines as the used resources exceed the hardware’s upper limit. This implies that configurations with hardware-unfriendly parameters consistently underperformed, allowing us to preemptively prune inefficient configurations from the large strategy space. Consequently, this approach maintains a streamlined strategy space for micro-kernel generation and eliminates the need to create a predefined sample list.

Moreover, the dependency on hardware hierarchy motivates us to generate the micro-kernels in a bottom-up manner. As the hardware hierarchy enables us to prune inefficient configurations based on the hardware limits level by level, only a limited number of micro-kernels is ultimately kept to support the dynamism of the tensor program. This approach enables comprehensive support for all potential input shapes in dynamic tensor programming.

In this subsection, we use the example of GEMM tensor program on GPUs as a paradigmatic example of recursive execution patterns on hierarchical hardware platforms. GEMM, as a classic operator in deep learning, is mathematically defined as $C=A\times B$, where $A$ and $B$ are the input matrices and $C$ is the output matrix.

As shown in Figure 7, it is natural to use recursive for-loops to represent GEMM execution flows. This approach breaks down the GEMM tensor program into a series of recursive loops set. These recursive loop sets elegantly map the high-level GEMM tensor program onto specific hardware levels, ranging from off-chip to on-chip memory and registers. Specifically, the uppermost recursion, OffChip_Mem_GEMM, processes end-to-end matrix blocks in off-chip memory. At intermediate levels, OnChip_Mem_GEMM handles smaller matrix blocks in on-chip memory. The innermost recursion, Register_GEMM, focuses on individual elements in registers, utilizing specific calculation instructions, such as Fused-Multiply-Add (FMA) (a100_whitepaper, ) instructions on GPUs.

This top-down recursive approach is intuitive and crucial for optimizing tensor programs on modern hardware architectures. It gives us an opportunity to decouple optimization space between different layers from each other, ensuring that each recursive layer is finely tuned to the unique capabilities and limitations of the corresponding hardware level, whether it be off-chip memory, on-chip memory, or registers. Furthermore, this approach provides a clear structure for optimizing different aspects of the tensor program independently at each level.

Recognizing the hierarchical nature of tensor program execution, we also acknowledge the variations between different hardware and tensor programs. Specifically, CPUs and GPUs exhibit distinct computation modes and memory access controls, where CPUs are optimized for multi-threaded parallelism, and GPUs excel in Warp, CTA, and Grid-level parallelism. Moreover, different tensor programs demonstrate unique loop patterns; for instance, the loop characteristics of Convolution markedly differ from those in GEMM.

These variations inspire our design of a unified abstraction, which delineates a universal approach for representing tensor program executions across various hardware platforms. Algorithm 1 elaborates on this abstraction. It maintains the layer-wise recursive structure as demonstrated in GEMM (Figure 7) as the core, and enables custom loop mapping and execution stages for various tensor programs.

To achieve a universal and customizable representation, we classify loops within each hierarchical level into three distinct sets in Algorithm 1. The Parallel Loop Set is designed for parallel execution; the Temporal Spatial Loop Set manages temporal non-reduction operations; and the Temporal Reduction Loop Set focuses on temporal reduction operations. Each level, identified as level N, abstracts the execution into three stages: Load, rKernel(N-1), and Store. These stages serve as flexible interfaces, allowing for tailored execution based on the specific requirements of each hardware level.

Table 1 illustrates how this unified abstraction rKernel is implemented across different hardware configurations, highlighting its versatility. Our focus lies in scrutinizing recursive execution patterns among different hierarchies. For CPUs, at the lowest level (L0), the rKernel abstraction allows for direct data transfer from Global memory to Registers or from CacheBuffer to Registers, depending on the specific needs of the computation. A “CacheBuffer" is defined as a memory buffer, sized within the L2 cache limits, to ensure consistent caching of its contents in the L2 cache for efficient data access and processing (onednn, ). Additionally, store operations at this level also reflect this adaptability, offering the choice of transferring data back to Global memory or CacheBuffer. As we progress to level L1, the abstraction provides options for either transferring data from Global memory to CacheBuffer or performing no operation, signifying a versatile approach to data handling. The highest level (L2) in CPUs focuses on the multi-thread mechanism at the process level, capitalizing on the CPU’s capabilities for multi-core parallel processing.

Similarly, for GPUs, including both Cuda Cores and Tensor Cores, rKernel adapts to different operational requirements. At L0, there’s an option for loading data either from Global memory or Shared Memory to Registers, and similarly, storing data either back to Global memory or to Shared Memory. This flexibility is crucial for optimizing memory utilization in the highly parallel environments of GPUs. At the L1 level, similar to CPUs, the abstraction can facilitate data transfer from Global to Shared Memory or no operation, enabling efficient resource management. The L2 level focuses on Grid-level operations, enhancing the scalability across the GPU’s multiple streaming multiprocessor (SM) architecture.

rKernel achieve a hierarchy abstraction of execution patterns that apply universally to various tensor programs and hardware types. This approach ensures a tailored strategy space for each hardware hierarchy level and facilitates the universal optimizations for dynamic-shape tensor programs.

This subsection introduces an innovative bottom-up method for generating candidates tailored to align with top-down recursive abstraction. This method becomes particularly crucial when execution parameter information is lacking, posing a significant challenge in every hierarchical layer for identifying suitable shape candidates for micro-kernel creation. Our approach centers on two key processes:

Firstly, we utilize the hardware’s parameter information to determine constraints for the candidate range. As highlighted in §2.3, our empirical studies have identified a notable decrease in hardware execution efficiency when the utilization at any hardware level is extremely low or high. This insight allows us to deduce a feasible range for candidate shapes, based on hardware utilization metrics.

Secondly, we follow a key design principle: ensuring the shape size of candidates in an upper layer is an integer multiple of the shape size in the lower layer. This approach aims to minimize padding loss during the construction of micro-kernels. As shown in Figure 8, if the sizes at one level are not multiples of those below, it results in more padding losses and inefficiencies at higher levels. Conversely, constructing candidates as integer multiples from one level to the next predominantly confines padding loss only to the outermost execution level, adhering to runtime requirements.

The overall process is detailed in Algorithm 2. The core function, GenerateCandidatesForLayer, is designed to operate distinctly based on the hierarchical layer it addresses. It begins by acquiring the hardware specifications for a given layer (denoted as L). This is achieved through the GetHardwareInfo and InitCands function, which retrieves essential hardware constraints that critically influence the parameter space. To validate candidate feasibility, the key is assessing memory usage against layer-specific limits and accounting for hardware constraints, such as a GPU’s 1024 threads-per-block maximum. For the initial layer (L = 0), the function employs FilterByISA to refine the candidate set according to the hardware platform’s Instruction Set Architecture (ISA) compatibility. For instance, on Intel CPUs, the FilterByISA function considers the granularity constraints of AVX512 (intel_xeon_whitepaper, ). Similarly, on GPU with Tensor Cores, the function assesses the constraints imposed by the Matrix Multiply-Accumulate (MMA) instruction (a100_whitepaper, ). These considerations ensure alignment with hardware capabilities.

Our algorithm utilizes FilterByMultiples, a method inspired by the classic sieve approach (sieve, ), to filter candidates. This function iteratively processes previous layer candidates (prevCands), generating multiples constrained via GenerateMultiples within the current layer’s candidate range. This approach ensures comprehensive exploration of viable parameter sets and maintains filtering efficiency. Additionally, we employ a mapping mechanism that uses a table to record the links between each candidate in the current layer and its possible match candidates in the previous layer, which is crucial for the subsequent analysis module.

In this subsection, we present a novel analyzing method which effectively incorporates analytical and empirical methodologies. This approach is designed to optimize the trade-off between efficiency and accuracy in strategy analysis.

Despite the contrasting architectures of GPUs and CPUs, $Vortex$’s abstraction consistently represents both, as exemplified in Table 1. One of the key aspects contributing to this universal representation is the definition of data structures.

The rKernel data structure, as illustrated in Figure 10, is a cornerstone of $Vortex$, tailored to encapsulate and streamline the complex processes involved. At its core lies the layer_meta_info class, pivotal for orchestrating the optimization strategy of each hierarchical layer. Within this class, the layer_depth attribute determines the layer’s position within the hierarchical structure, which is a crucial factor for the recursive optimization process. Moreover, the map<axis, LOOP_TYPE> provides a strategic mapping of loops to their respective types, including Spatial (S), Temporal Parallel (TP), and Temporal Reduction (TR). This nuanced approach to loop optimization aligns seamlessly with the unique characteristics of each tensor program. The ANALYZE_TYPE enum, encompassing empirical and analytical options, facilitates the selection of the appropriate optimization analysis method. Finally, the functional pointers—load_func, store_func, and compute_func—are critical in dynamically managing the various stages of computation. This data structure encapsulates the necessary elements to navigate dynamic-shape tensor program optimization challenges adeptly.

rKernel serves as a recursive-based template for dynamic tensor program. When dealing with a fixed hardware and a particular operator, it necessitates users to initialize different levels of layer_meta_info. The development effort required for this task is minimal. This is because, for both CPU and GPU, we set the hierarchy level to three, which does not impose a significant burden. Furthermore, the dynamic parameters of the template enable consistent computational support for various runtime shapes. Owing to its versatility, development costs are apportioned across various computing scenarios. We built $Vortex$ on top of TVM (tvm, ). We harness the ‘tensorize’ primitive to implement the load, store, and compute functions, and leverage TVM’s robust and versatile code generation capabilities, we seamlessly target both CPU and GPU platforms. In Figure 11, we illustrate a representative GEMM kernel code generation for GPUs, showcasing our methodology’s depth and adaptability in optimizing dynamic-shape tensor programs for GPU architectures.

The integration of offline and runtime components in $Vortex$ involves several key steps. During runtime, $Vortex$ employ analytical cost models to estimate the execution costs associated with different candidate solutions with runtime shape information. Subsequently, $Vortex$ select the most suitable micro-kernel candidates based on these optimal execution cost estimations.

Notably, our selection process accommodates the dynamic nature of hardware platforms. For instance, in the case of GPUs, the presence of a larger MMA instruction (a100_whitepaper, ) padding in the tensor core necessitates adaptive hardware solutions. We provide implementations for both CUDA cores and Tensor cores, allowing us to choose the appropriate backend hardware based on the runtime input shapes adaptively, further optimizing execution efficiency. The performance benefits of this adaptive strategy are thoroughly explored in §7.4. Additionally, by considering the computational shape of the selected micro-kernels in conjunction with the runtime shape, we collectively compute runtime-specific computational details, such as grid configurations. This comprehensive approach ensures the versatility and generality of our computational framework, facilitating efficient and adaptable runtime operations.

In this subsection, we present the evaluation of single dynamic-shape tensor program, specifically assessing GEMM and Convolution operators on CPU and GPU platforms. It is notable that DietCode is limited to GPU CUDA Cores and requires pre-determination of dynamic shape samples. We leverage the parameters from Table 3 and Table 4 as sample sets for DietCode’s offline compilation process. Importantly, the latency measurement for $Vortex$ encompasses both the operator execution time on the hardware platforms and the runtime overhead from $Vortex$’s cost model.

The evaluation results, shown in Figure 12, demonstrate $Vortex$’s performance across various configurations. The x-axis outlines the number of floating-point operations (FLOPs) in the workloads, including all GEMM test cases from Table 3 and convolution from Table 4, while the y-axis represents the speedups. $Vortex$ consistently achieves a generalized performance speedup, demonstrating improvements across different hardware setups, operators, and tensor shapes. To further quantify the effectiveness of $Vortex$, we emphasize two metrics: the percentage of cases where $Vortex$ shows performance improvement (defined as cases where the speedup is greater than one) and the average speedup across all cases. Table 5 presents these detailed results. Overall, this comprehensive evaluation confirms that $Vortex$ provides robust and efficient acceleration results.

Additionally, we explore the impact of DietCode’s reliance on sample-specific performance. As shown in Table 6, we configure the M dimension dynamically in DietCode, sampling and compiling it within the range [128, 256). The results show a performance decline when deviating from this range, highlighting DietCode’s limited flexibility.

Figure 13 presents a comprehensive evaluation of the end-to-end performance of classic language models and CNN models. We compare $Vortex$ against existing state-of-the-art solutions, using oneDNN and cuBLAS/cuDNN as the normalized baselines for CPU and GPU evaluations, respectively.

$Vortex$ demonstrates significant performance improvements across various tasks. $Vortex$ achieves notable average speedups of $2.91\times$ for BERT, $2.63\times$ for BERT-Large, and $2.94\times$ for GPT-2, across different baselines and hardware configurations. For CNNs, $Vortex$ achieves remarkable average speedups of $2.01\times$ for AlexNet, $2.13\times$ for ResNet, and $3.24\times$ for GoogleNet. Furthermore, $Vortex$ demonstrates considerable performance improvements over multiple existing solutions. Specifically, $Vortex$ achieves average $1.73\times$, $4.26\times$, $1.43\times$, $1.71\times$, $3.32\times$ and $4.13\times$ over oneDNN, ONNX Runtime, cuBLAS, cuDNN, CUTLASS, DietCode, respectively.

The evaluation results underscore that the performance of $Vortex$ varies with different hardware and models, reflecting the distinct execution characteristics inherent to each environment. Notably, $Vortex$ consistently outperforms the baseline across a wide range of input shapes and model types, demonstrating its exceptional adaptability and efficiency in real-world dynamic-shape DNN computation.