FusionStitching: Boosting Memory Intensive Computations for
Deep Learning Workloads

JIT fusion is the state-of-the-art technique to reduce context switch and memory access overhead of intensive memory operations for machine learning workloads. As the state-of-the-art fusion engine, XLA only supports thread-local data transferring for fusion, which relies on index analyzing and re-computation to improve thread locality. A bad case is to put a reduction in the middle of a fusion pattern. If different threads require the same value of reduction result, each thread needs to recompute the reduction independently and redundantly. To avoid this, XLA does not fuse operations that lead to middle-reductions. As a result, this strategy misses many optimization exploration space.

Figure 1 shows how XLA and FusionStitching perform fusion for layer normalization. XLA forms 4 fusions, where fusion.3 and fusion.7 end with reductions and fusion.2 ends with expensive operation. XLA does not do other fusion to avoid re-computation overhead, with the drawback of CPU-GPU context switch and off-chip memory access overhead. While FusionStitching generates only one kernel with all intermediate results residing in on-chip memory and involves no CPU-GPU switching. No re-computation is introduced in FusionStitching .

Static libraries are not feasible to solve this problem as a single memory intensive operation is lightweight and the combinations of them vary in different models. TVM [11, 55] mainly focuses on compute intensive operations but not fusion optimization, whereas XLA only provides simple rule-based fusion strategies for memory intensive ops. In addition, TVM relies on pre-defined schedules for code generation and usually requires a long time for tuning.

In this work, we have analyzed a variety of machine learning workloads (Section 7). we get two main observations.

Severe Context Switch Overhead Context switch overhead between CPU and GPU is a severe performance issue. Table 2 shows the breakdown analysis of a wide range of machine learning applications. The count of GPU kernel calls ($\#$) for TensorFlow implementation can reach up to 10,406, which leads to large kernel launch overhead. Even fused with XLA, the kernel calls can be still up to 6,842. As a result, both scheduling and pre-/post-processing time on CPU brought by machine learning framework dominate the execution time for some models (like BERT inference, DIEN, ASR, and CRNN).

Large Portion of Memory-intensive Ops Also in Table 2, the execution time of memory intensive ops can be up to 40% in the overall time for some models. Off-chip memory traffic time usually occupies a large portion of the overall execution time of memory intensive ops. By fusing a large quantity of operations, our system is able to reduce CPU-GPU context switch overhead and leverage high speed on-chip memory to reuse intermediate data between ops.

Although data reuse is a potential opportunity to fuse a large scope of ops together to avoid redundant re-computation, there are several main challenges for both applying reuse given a fusion pattern and deciding op fusion strategy.

As for applying reuse for a given fusion pattern to generate kernel:

• •

  The first challenge is how to evaluate reuse benefit. Reuse is not always better than re-computation. Reuse will lead to extra inter-thread data communication, whereas re-computation introduces redundant compute overhead but without inter-thread communication. Besides, introducing reuse with shared memory may hurt kernel occupancy, thus hurts parallelism.

• •

  The second challenge is how to apply reuse. There are two types of data reuse: intra-warp reuse and intra-block reuse. Intra-warp/block reuse means that threads within a warp/block reuses intermediate data produced by other threads within the same warp/block. The trade-off is that, intra-warp reuse has less memory access overhead while intra-block reuse has better parallelism.

For a given machine learning model, the challenge is to decide what ops should be fused together. XLA also tries to solve this challenge, but with limited exploration space. One main problem is that, forming a fusion pattern by data reuse is not always better than separate kernels. Reuse requires data locality within thread-block or warp, which can potentially limit parallelism. Take reduction as an example, reuse of reduction result requires to compute reduction within a block or warp. However, some reductions perform better with several blocks work together, which has higher parallelism. Not fusing such a reduction with its consumers may result in better performance, even with context switch and off-chip memory access overhead. Intra-block reuse may further hurt parallelism as it requires extra shared memory. Another typical problem is that, when operation A can be fused with either operation B or operation C, while B and C cannot be fused together due to some limitation, it is not always easy to decide which one to fuse for A. Rule-based approaches, like XLA, fail to find effective fusion plans for varied models.

As described before, FusionStitching widens fusion possibilities by introducing data reuse, a rarely used method in state-of-the-art JIT fusion techniques. Data reuse can be applied when different threads processing operation $B$ requires the same value generated by operation $A$. We assume $B$ is the consumer of $A$. The threads requiring the same value can read the same data with inter-thread communication, but do not have to re-compute it redundantly like in XLA. We observe that tensor shapes of a sequence of memory intensive operations always shrink and broaden frequently due to operations like reduce, broadcast, gather, and slice. Thereby the opportunity to explore data reuse is large.

In our work, we explore both intra-warp reuse and intra-block reuse, corresponding to warp composition and block composition in figure 3. Take reduction as an example, intra-warp reuse means that each warp does reduction for a row of data and stores result in the register of the first lane of the warp. Consumers of the reduction read data with register-shuffle from the first lane. Intra-block reuse does reduction for the row with all threads in the block and stores results in shared memory. Consumers of the reduction read data from shared memory. This approach enables stitching operations with on-chip memory while avoiding re-computation. Data locality should be guaranteed for correct data reuse. Intra-warp reuse requires warp level data locality, and intra-block reuse requires block level data locality.

We divide optimization process into two stages in FusionStitching : fusion exploration and code generation. The two stages are conceptually independent but highly related.

As shown in figure 2, FusionStitching consists of fusion explorer and code generator. Fusion explorer generates possible fusion patterns that may enjoy data reuse. It then applies beam search to form several candidate fusion plans with these fusion patterns, and selects the best fusion plan from candidate plans with a cost model. Code generator generates GPU kernels for each fusion pattern produced by fusion explorer. We pre-define a set of schedules (i.e., templates) for each kind of operators. A schedule describes code generation about how an operation runs in parallel. Meanwhile, it indicates whether threads of an operation reuse data with other threads and the reuse scope, and whether this operation produces data for consumer’s reuse. FusionStitching divides operations of a fusion pattern into several groups, and enumerates all valid combinations of the schedules in the groups. With a cost model that estimates the performance of each schedule and launch dimension setting, FusionStitching selects the best configuration of the fusion pattern and generates GPU kernel code.

We design a two-level cost model in FusionStitching . Fusion explorer needs to search in large search space and applies delta-evaluator (5.4), which is fast but less accurate. Code generator operates on merged GPU kernels and needs more accurate performance estimation, and thus we apply latency-evaluator (4.3) which is more accurate but slower.

We first describe how code generator works in section 4 and then describe fusion explorer in section 5.

We study about four kernel composition schemes, which indicate main behaviors of common memory intensive ops. Different scheme indicates different data dependence and parallel behaviors of kernels to fuse, ranging from no dependence to complex cross-thread dependence, and from uniformed parallelism to non-homogeneous computations. Figure 3 illustrates the four kind of composition schemes.

Kernel Packing packs computations of ops with no data dependence. This scheme is instrumental in reducing context switch overhead of kernel launch and framework scheduling. It also reduces loop control overheads in instruction level. To reduce control flow overhead, we perform aggressive loop fusion [5, 18] to merge as many element-wise ops as possible into a single loop structure when these ops have the same parallelism dimension.

Thread Composition fuses dependent ops and transfer intermediate results via registers within a local thread context. This is the fusion scheme XLA applies. This may introduce redundant computations between threads requiring the same value.

Warp Composition fuses dependent operators and apply intra-warp reuse. This scheme employs register shuffle to communicate between threads within a warp. A typical case is warp reduction and its consumers, which applies warp level reduction and leverages registers to transfer intermediate results to threads of its consumers.

Block Composition applies intra-block reuse and unlocks the potential to enable composing non-homogeneous computations into large fused kernels, as long as these computations can communicate within block level. It makes use of shared memory to transfer intermediate results for data reuse.

Warp composition and block composition are flexible schemes as they allow different ops to execute in independent schedules in the fused kernel. The only requirement is to keep warp/block locality between producers and consumers. These two schemes are essential to compose a broad range of op kinds with various parallelism characteristics and dependence relationships efficiently.

As far as we know, FusionStitching is the first work to thoroughly study about all above composition techniques for just-in-time compilation of memory intensive operations for machine learning workloads. XLA only realizes kernel packing and thread composition. We do not stitch ops that involves inter-block communications as it results in global memory level synchronization and introduces high overhead.

Code generation in FusionStitching is not as straightforward as XLA, because of the trade-off of composition schemes we discussed above. We use a cost model based tuning approach in FusionStitching , and ease the process by op grouping.

According to the four types of composition schemes, we pre-define a set of schedules for each kind of operators (light element-wise, expensive element-wise, and reduction ops). Specifically, we define a single template for light element-wise ops representing kernel packing and thread composition. Note that we find kernel packing and thread composition can be described with the same schedule. We define three schedules for expensive element-wise ops and reduction ops. The first one represents kernel packing and thread composition. The second one stores intermediate result to the register of the first lane in each warp, representing warp composition. The third one stores intermediate result to shared memory, representing block composition.

FusionStitching enumerates all schedule combinations of ops. It estimates these combinations according to cost model (sec 4.3). To ease the enumeration process, we divide ops in a fusion pattern into several groups. Each group applies one schedule and different groups transfer intermediate data with warp/block level reuse.

We group ops according to op types. Given a fusion pattern, we enumerate a small set of grouping strategies. We call sub-root as the output op of a group, and root as the output of the fusion. We first identify sub-roots and generate groups rooted from sub-roots and root. Reduce ops are always regarded as sub-root. Expensive element-wise ops are enumerated to both sub-roots and non sub-roots. Other ops are neither sub-roots. This leads to a set of groups rooted from reduction, expensive element-wise and root ops.

The insight of the grouping approach is that, the schedule of non sub-roots can be determined by the schedule of sub-roots by tensor indices propagation. Thus we only need to enumerate the schedule of sub-roots and root ops, and we can propagate to get schedules of all ops.

FusionStitching enumerates grouping strategies, and emulates schedules of every sub-root/root op and launch dimension of the fused kernel. As data reuse requires correct data locality in the reuse scope, schedules do not match data locality requirement are discarded. After estimating the performance of each enumeration with latency-evaluator, FusionStitching selects code generation strategy with the best estimated performance.

FusionStitching requires an accurate estimation for kernel performance for code generation strategy. Fortunately, as the searching space of code generation is not very large, we can tolerate a relative slow cost model.

We build latency-evaluator as in equation 1, where $L$ is the estimated execution cycles of a fused kernel.

$N_{wave}$ means how many waves of warps that will be processed by a GPU, noting that the warps for a large GPU kernel will be spitted into several waves to be executed on a GPU where warps in the same wave executes concurrently. $L_{warp}$ means the latency (cycles) a single warp spends in the fused kernel. The multiply of $N_{wave}$ and $L_{warp}$ stands for the total cycles to execute the fused kernel.

We estimate $N_{wave}$ with the total number of warps to issue (${N_{warp}}$) and the occupancy of the fused kernel (${Occupancy}$). $N_{warp}$ is calculated with launch dimension and tensor shape. We calculate $Occupancy$ with launch dimension, shared memory usage (Sec. 4.4) and estimated register usage. We estimate the register usage by analyzing the life time of every intermediate value and get the maximum register usage for a thread. Also, we estimate shared memory usage by life time analyzing. This approach is accurate enough for us to calculate $Occupancy$.

As for $L_{warp}$, we use the reported $CPI$ numbers for different types of ops [22, 21] and multiply it with the total instruction count ($N_{instruction}$) we estimated.

It is essential to use shared memory moderately as large amount of shared memory usage hurts kernel parallelism, especially for large granularity compositions. To use as much shared memory as possible while not hurting parallelism, we explore a dataflow based shared memory sharing technique. The insight is that, FusionStitching reuses previous allocated shared memory as much as possible to reduce unnecessary shared memory allocation.

We use dominance tree algorithm[12] for shared memory dataflow analysis. The approach takes a computation graph and shared memory requests as input, and outputs an allocation map. To optimize shared space sharing, we traverse ops of the computation graph in topological order. When an op does not need shared space, previous allocation information will be propagated forward. If an op needs shared space, we merge allocation information of all its operands, test the dominance relation to check if we can share any previously allocated space for current op, and reuse the space if possible.

Beside the reuse of intermediate result between producers and consumers with shared memory and register shuffles, FusionStitching also reduces thread local redundant calculations. Thread local redundant calculations mainly come from memory access index calculations and thread local intermediate values. It is because different parts in a fusion kernel may use different schedules, and index and some intermediate values are generated independently within each schedule, even in the same thread. Before generating the code, FusionStitching first analyzes the overall index and intermediate value characteristics and then organizes the output code to reuse computations and data as much as possible.

We formulate fusion exploration as a subgraph searching problem. For computation graph $G=(V,E)$, where $V$ and $E$ are sets of vertices and edges respectively. We define a fusion pattern $P_{i}=(V_{i},E_{i})$ as a subgraph of $G$, with $V_{i}\subseteq V$, $E_{i}\subseteq E$. A fusion plan is a set of disjoint fusion patterns $S=\{P_{0},\cdots,P_{k-1}\}$. The score function of $P_{i}$ is annotated as $f(P_{i})$. The higher the performance is, the larger $f(P_{i})$ is. So the goal of computation fusion problem is to find fusion plan $S$ with maximal $\sum_{i=1}^{k}f(P_{i})$.

A Brute-force way to enumerate all fusion patterns has a complexity of up to $O(2^{V})$ without pruning. We proposed an approximate dynamic programming approach with complexity of $O(V+E)$ to find good fusion patterns.

The basic idea of fusion exploration is that, we generate candidate-patterns starting from each vertex in post-order in the graph, and select and compose final fusion plan with these candidate patterns. The candidate-patterns starting from vertex $V_{i}$ is the set of fusion patterns whose producer node is $V_{i}$. We describe how we generate candidate-patterns for each vertex in this section and describe how to compose the final fusion plan in section 5.3

Given a computation graph $G$, we get a topological sorting list. We generate candidate-patterns for vertices in post-order, from the last vertex to the first vertex. This is to guarantee that each searching results in a fusion pattern with current vertex as producer.

Equation 2 shows the approximate dynamic programming process. $C_{i}$ is the set of all consumers of $V_{i}$. PatternReduction is the approach to get candidate patterns of $V_{i}$ according to all its consumers’ candidate pattens.

We describe PatternReduction approach with an example shown in Figure 4, where $V_{0}$ is the root vertex who produces the output of whole graph. Assume that candidate-patterns for vertices before $V_{8}$ have already been generated. (We only explore top $k$ patterns in candidate-patterns for each vertex according to score function $f$. We set $k$ as 3 in this case.) We will generate candidate-patterns for $V_{8}$ with its consumers’ information. A naive approach is to append $V_{8}$ to all possible combinations of patterns in its consumers’ candidate-patterns sets, and select the top 3 patterns within the appended combinations. However, the combinations will be huge when the consumer number and top $k$ setting is large, noting that JIT approach requires timely optimization. Instead, we design an approximate divide-and-conquer process to find top 3 patterns with limited complexity as PatternReduction.

We first divide the consumers of $V_{8}$ into several groups and find candidate-patterns for $V_{8}$ considering these groups independently, and finally compose final candidate-patterns by reducing the above results of all groups. We assume the group number in this case is 2 (group {$V_{5}$, $V_{6}$} and group {$V_{7}$}). For group {$V_{5},V_{6}$}, we enumerate all possible combinations of patterns in candidate-patterns of $V_{5}$ and $V_{6}$. Specifically, there are 15 possible combinations between $V_{5}$ and $V_{6}$, including empty set. We append $V_{8}$ to each of the 15 combinations and select the top 3 patterns as the temporary candidates associated to group {$V_{5},V_{6}$}. We get another top 3 patterns considering group {$7$} as temporary candidates. We finally select the final top 3 patterns as candidate-patterns for $V_{8}$ from all above 6 temporary candidates. Note we validate top patterns according to score function $f$.

The PatternReduction process above is recursive if consumers number of a vertex is very large. We first divide the consumers into two groups. To get candidate patterns of a group, we divide the group further until it is small enough.

A constraint of a fusion pattern is that, no cyclic dependence is allowed. Figure 6 shows an example that a cyclic dependence occurs after fusion. FusionStitching discards patterns with cyclic during the searching process. Meanwhile, FusionStitching only explores fusion patterns that the code generator can process. It does not form fusion patterns with cross-block communication requirement.

Remote Fusion To further reduce the context switch overhead between CPU and GPU, we try to merge fusion patterns that are not adjacent in the graph after above procedures. As is shown in Figure 5, we add a virtual vertex $h$ as the producer for all vertices and apply PatternReduction. We finally get the candidate-patterns of $V_{h}$, which includes the fusion of remote patterns that are not adjacent. The remote pattern fusion helps to reduce generated kernels and thus reduces the context switch overhead. Remote fusion results in kernel packing for code generation.

All candidate-patterns of all vertices, which may overlap with each other, form a new set $E$. The insight of forming good fusion plan is to select non-overlapped patterns from $E$ with high score of $f$. Note $f$ means how much cost the pattern saves (sec 5.4).

FusionStitching uses beam search [17] to generate top-3 (width of the beam) candidate fusion plans, and finally selects the best plan within the 3 candidates with latency-evaluator (sec 4.3).

Specifically, FusionStitching maintains 3 buffer sets to store candidate fusion plans. It traverses from the producer vertex to the consumer vertex in op graph $G$, and try to append each candidate pattern of each vertex to each buffer set in turn if it introduces no overlapping. It may result in more than 3 temporal sets after appending in each iteration, and will select top-3 with highest accumulated $f$ as new buffer sets.

After traversing all vertices, each buffer set forms a candidate fusion plan. FusionStitching will use the more accurate cost model latency-evaluator to estimate which candidate plan performs the best. The best one is the final fusion plan.

FusionStitching applies delta-evaluator to form the score function $f$. One insight is that, we only need to estimate the performance gain or loss when forming a fusion pattern, but do not require accurate estimation of the overall execution time. With this insight, the score function $f$ represents the performance gain or loss of a fusion pattern only.

There are three main factors in delta-evaluator: reduced memory access latency, reduced CPU-GPU context switch overhead, and performance penalty of kernel fusion. The score function $f$ is the summary of these three parts, shown in equation 3.

We estimate reduced memory access latency ($T_{reduced\_mem}$) with two factors. The first is the amount of memory traffics between operators to be fused, which can be calculated according to the shape and type of input and output tensors. And the second is the change of memory type to store the intermediate values between operations. We build a regression model to predict the reduced memory access latency when changing the memory type from global memory to register or shared memory, when given memory traffic amount. The regression model is based on latency data we collected offline on various GPU vendors with various amount of data traffic.

$T_{reduced\_calls}$ can be easily estimated by multiplying the number of fused kernels with a fixed value representing average CPU-GPU context switch time we collected.

The performance penalty ($T_{penalty}$) is estimated by a simplified version of latency-evaluator (sec 4.3). As fusion plan exploration faces a very large searching space, using latency-evaluator directly takes too long time for JIT optimization. The most time consuming part of latency-evaluator is the estimation of $Occupancy$. To estimate $Occupancy$ fast, we use a fixed number (16) as registers number, and use the maximal shared memory usage in and between any ops within a fusion pattern as the overall shared memory usage. Life time analyzing of registers and shared memory is discarded in delta-evaluator.

In this section, we evaluate FusionStitching using a variety of machine learning applications with different characteristics in different fields. Table 1 summarizes the various fields of the evaluated applications and the characteristics of these applications. These applications range from images (CRNN[37]), speech (ASR[52]), NLP (Transformer[47], BERT[15]), to internet scale E-commerce search and recommendation systems (DIEN[57]). The building blocks of these workloads include perceptron, attention, convolution, RNN and a broad range of memory intensive operators.

To demonstrate the benefits of FusionStitching over previous work, we compare it with the default TensorFlow implementation and XLA (up-to-date with community functions). All evaluation results are collected on NVIDIA V100 GPU with 16 GB device memory. The server runs Red Hat Enterprise Linux 7.2 with CUDA toolkit 10.2 and cuDNN 7.6.

We evaluate the speedup of FusionStitching by comparing inference cost or the training time of one iteration for TF, XLA and FusionStitching with the same batch-size. During our test, the accuracy in each iteration of training and the result of inference are the same with TF and XLA. We repeat 10 times and use the average performance to validate speedup. As for training workloads, we collect the execution time from the 11th iteration to the 20th (guaranteed to be stable), to avoid the initialization overhead of the early training iterations.

We show the speedup of FusionStitching in figure 7, where the execution time of TensorFlow is normalized to 1. Compared to TensorFlow, our approach achieves up to $2.42\times$ speedup, with 1.66$\times$ on average. Compared to XLA, our approach achieves up to $2.21\times$ speedup, with 1.45$\times$ on average. Note XLA shows performance degradation for DIEN (both training and inference), while FusionStitching does not show negative optimization in any of these cases.

We also test the inference workloads on NVIDIA T4 GPU and get the similar speedup.

We apply FusionStitching in production and measure the performance benefits. It shows that FusionStitching saves about 7,000 GPU hours for about 30,0000 tasks in a month. The performance result on the cluster shows FusionStitching is robust. Note that XLA, as state-of-the-art JIT engine, cannot be enabled default as it suffers from negative optimizations for many cases.

Overall, the performance benefit of FusionStitching comes from both reduced CPU-GPU context switch and off-chip memory access. FusionStitching gives a good solution to the problems we found (sec 2.2). We provide performance breakdown information in the following section.

Table 2 shows the kernel breakdown information, including execution time (T) of memory intensive ops (Mem), compute intensive ops (Math), CPU time (CPU, kernel launch and framework scheduling), CUDA memcpy/memset activities (Cpy) and kernel call times (#). Note that the breakdown profiling process is different with the process to measure the end-to-end performance. This is because profiling introduces some overhead and makes the end-to-end time not accurate. We profile Mem, Math, Cpy time directly, and get CPU time by dividing other time from end-to-end time.

Before we analyze the effect of our technique, we need to point out that XLA affects the behavior of Matrix operations (GEMM and GEMV). It tends to fuse GEMVs into GEMM, and GEMMs to larger GEMM when there are Matrices sharing common input. Some other algebra transformation and loop-invariant code motion also reduces GEMM count. The difference in GEMM count in table 2 is caused by such reasons. Meanwhile, XLA affects the runtime behavior of TensorFlow and leads to more or less CUDA memcpy/memset activities. FusionStitching is built upon XLA and exhibits the same behavior.

According to Table 2, we have the following observations.

Reduced context switch overhead. FusionStitching effectively reduces the memory intensive kernel calls of all workloads, which results in reduced kernel launch and framework scheduling overhead. As is shown in Table 2, the number of memory intensive kernel calls with FusionStitching is 38.0% of that with XLA in average, ranging from 27.8% to 48.4%. Meanwhile, FusionStitching reduces CUDA memcpy/memset activities (Cpy time) than XLA, with 34.3% decrease in average. This is because we fuse more ops together into larger kernels than XLA and many CUDA memcpy/memset are combined together. The CPU time difference in table 2 indicates the reduced time due to the decrease of kernel calls and CUDA memcpy/memset activities. FusionStitching achieves up to 61.0% saving of the $CPU$ time comparing with XLA, 41.0% in average.

Take DIEN-train as an example, the kernel call number for memory intensive ops is 2109 with FusionStitching , and 6842 with XLA. Meanwhile, the CUDA memcpy/memset activities is reduced to 1395, comparing to 1996 with XLA. The final $CPU$ time with FusionStitching is significantly less than both TF and XLA, thanks for the reduced kernel calls. Note that XLA increases CUDA memcpy/memset activities and results in severe performance drop here. FusionStitching avoids the increased memcpy/memset calls due to larger kernel granularity and do not suffer from the drawback. DIEN-infer has the similar behavior. Optimizing kernel fusions while considering runtime behaviors (like memcpy activities) could be a future research topic.

Reduced memory intensive op execution time. FusionStitching reduces the total execution time for memory-intensive operations. The speedup of memory-intensive ops for all workloads is $1.39\times$ in average comparing with XLA, and up to $1.74\times$. The performance speedup mainly comes from reduced global memory access. By fusing memory-intensive operations aggressively, the intermediate values can be cached in registers and shared memory.

Take CRNN as an example, it reads 667.6 MB global memory with XLA, while FusionStitching reduces the traffic to 225.8 MB. About 66% global memory access has been reduced for memory intensive ops. The execution time of all memory intensive computations thus achieves $1.48\times$ speedup than XLA.

Overall speaking, FusionStitching supports more complex fusion patterns than XLA with effective kernel generation, which relaxes the fusion conditions and thus reduces the final kernel numbers and intermediate global memory transactions. These two factors are essential for the performance of memory intensive ops (sec 2.2). Meanwhile, with well controlled performance estimation and reduced runtime memcpy activities than XLA, FusionStitching is less likely to result in bad case about optimization.

We take the Layer Normalization case in Figure 1 again to analyze the fusion result of XLA and FusionStitching . Note Layer Normalization is a very common component in deep learning models.

XLA forms four different fusion kernels. While FusionStitching is able to fuse them all and generate a more efficient kernel. We collect the performance of all kernels of the two version. The single kernel with FusionStitching achieves a speedup of $1.23\times$ comparing with the sum of all 4 kernels with XLA. For this test, we do not count the context switch overhead for the 4 kernels in XLA version, which further hurts the performance of XLA.

There are two factors that prevents XLA to fuse operators with a larger granularity in this case. The first reason is reduce op in xla-fusion.7 and xla-fusion.3. As we discussed before, XLA does not explore reuse of intermediate results of producer op. Fusing reduce op in the middle of a kernel results redundant computation of the same value in different threads, thus hurts the performance. The second reason is expensive element-wise ops with small tensor shape (xla-fusion.2). XLA does not tend to fuse expensive ops processing small tensors in the middle of a kernel as the fusion will introduce redundant computations of expensive instructions.

Instead, FusionStitching finds notable potential for larger granularity fusion in this case and fuses all operators into one kernel. It applies data reuse to prevent redundant computation. In this way, the intermediate global memory transaction and CPU-GPU context switch is avoided.

Similar with XLA, FusionStitching is designed for tune-once-run-many-times scenarios, which is a basic characteristic of deep learning workloads. For the training that could take up to several days, FusionStitching only needs to run in the first training iteration. For an inference task, the executable kernels can be prepared once and run many times in the future.

We measure the one-time overhead introduced by FusionStitching compared to XLA for training. The value is the JIT compilation time of FusionStitching minus that with XLA. Results show that the extra overhead is less than 30 minutes for the workloads we evaluated in this paper, which is far less than the overall training time.

As for cost models, we tried to use complete latency-evaluator to estimate $T_{penalty}$ in delta-evaluator. The results show a much longer tuning time, but do not show better performance of tuning results. It demonstrates the simplified cost model works good for fusion exploration.

A problem of both FusionStitching and XLA is that, they cannot handle dynamic shapes, appears in some deep learning workloads, with low tuning overhead. The reason is that the design of XLA service framework is not friendly to dynamic shape, while FusionStitching is implemented based on XLA service framework. This implementation problem does not affect the insight that FusionStitching shows.