Balanced Segmentation for Multi-Accelerator Inference

• The paper introduces profiling-based, differentiable, and receptive field-aware segmentation methods that balance workload across multiple accelerators for reduced latency and energy consumption.
• Focusing on methodology and key results, it employs binary search and dynamic programming to minimize makespan and achieves up to 8× latency reduction and 50× energy savings.
• The approach addresses communication overhead, halo management, and quantization challenges to ensure functionally correct, cost-effective deployments on heterogeneous hardware.

Balanced Segmentation for Multi-Accelerator Inference refers to the suite of algorithmic strategies and system-level optimizations for partitioning deep neural network inference workloads—especially convolutional neural networks (CNNs)—across multiple cooperating hardware accelerators (edge TPUs, heterogeneous system-on-chips, or distributed edge servers). The primary objective is to maximize throughput or minimize energy by distributing computation such that each accelerator operates near its performance/energy frontier and hardware bottlenecks (memory, bandwidth, quantization granularity) are respected. Approaches span profiled layer-segmentation with memory constraints, channel-wise mixed-precision assignment, receptive field- and halo-aware fused-block scheduling, and differentiable mapping optimized at training time.

1. Partitioning Formulations and Objectives

Partitioning CNN inference for multi-accelerator execution arises as a constrained scheduling problem, where the decision variables consist of segment boundaries (layer-wise or channel-wise), sub-layer mapping to each accelerator, and sometimes spatial splits in feature domains. The objective can be formalized as minimizing the makespan of the slowest accelerator (min-max per-segment inference time), balancing energy consumption, or optimizing Pareto fronts between accuracy and deployment cost.

In the profiler-driven approach (Villarrubia et al., 2 Mar 2025), the task is to partition the ordered layers of a pre-trained CNN into $N$ contiguous segments $(S_1,\ldots,S_N)$ for $N$ identical accelerators (e.g., Edge TPUs), such that:

• $\max_j \sum_{i \in S_j} T_i$ (the per-segment compute time) is minimized,
• $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$ (the per-accelerator on-chip memory cap is observed).

In differentiable mapping approaches (ODiMO) (Risso et al., 2023, Risso et al., 2024), the partitioning extends to assigning output channels of each layer to heterogeneous accelerators, taking into account per-accelerator precision, latency, energy, quantization-induced accuracy loss, and hardware-specific operator support.

2. Profiling-Based Layer Segmentation

Profiling-based segmentation begins with instrumented end-to-end inference on the target accelerator, capturing per-layer empirical execution time $T_i$ and peak on-chip memory footprint $M_i$. Analytical models, parameterized for FLOPs, memory bandwidth, and data precision, supplement these measurements:

• $T_i \approx F_i/P + (W_i+O_i)/B$ (FLOP- and bandwidth-based latency).
• $M_i = R_i \times S$ (activation shape and data type size).

Using this data, the segmentation problem is solved via binary search over makespan $T^*$: for a given $(S_1,\ldots,S_N)$0, a greedy left-to-right scan attempts to allocate contiguous layer chunks without exceeding $(S_1,\ldots,S_N)$1 or the memory cap per accelerator. If a feasible $(S_1,\ldots,S_N)$2-way split exists, $(S_1,\ldots,S_N)$3 is lowered. Otherwise, $(S_1,\ldots,S_N)$4 is increased. This approach achieves $(S_1,\ldots,S_N)$5 computational complexity (Villarrubia et al., 2 Mar 2025).

Heuristics further speed up the process, e.g., sorting layers by compute density $(S_1,\ldots,S_N)$6 and adding a small slack $(S_1,\ldots,S_N)$7 to avoid excessive fragmentation due to minor memory overflows.

3. Fine-Grained and Differentiable Partitioning

Fine-grained mapping addresses the limitations of layer-level splits in heterogeneous SoCs, enabling assignment of individual output channels to accelerators of varying capability or precision. This is formulated using relaxed assignment variables $(S_1,\ldots,S_N)$8 for each layer $(S_1,\ldots,S_N)$9, output channel $N$0, and compute unit $N$1, constrained so $N$2. Training-time optimization jointly updates network weights and mapping assignments to minimize a composite loss: $N$3 subject to accuracy constraints (Risso et al., 2023, Risso et al., 2024).

The cost term $N$4 encodes smooth hardware latency or energy models, with quantization levels matched to accelerator capabilities via fake-quantization in the forward/backward passes. After convergence, assignments are discretized and the network retrained with the final mapping.

This approach enables the discovery of Pareto-optimal trade-offs between accuracy and efficiency, as shown by ODiMO reducing energy and latency by up to $N$5 and $N$6, respectively, on contemporary SoCs (Risso et al., 2024).

4. Receptive Field-Aware Spatial and Fused-Layer Segmentation

For distributed (edge/fog) multi-accelerator inference, partitioning along the spatial dimensions of feature maps or fusing blocks of layers is effective. To guarantee bit-exactness (i.e., no accuracy loss from partitioning), receptive field-aware arithmetic delineates viable partition points: spatial splits are only valid where activations' receptive fields do not cross the subdomain boundary, preventing the loss of convolutional context (Li et al., 2022).

Partition variables $N$7 denote what fraction of a fused block $N$8 is assigned to edge server $N$9, ensuring all outputs' receptive fields are covered. Communication is reduced by exchanging only halo regions necessary for adjacent servers' next-layer inputs, rather than complete sub-outputs.

Dynamic Programming for Fused-layer Parallelization (DPFP) (Li et al., 2022) seeks the optimal segmentation by

• Precomputing the cost $\max_j \sum_{i \in S_j} T_i$0 of assigning contiguous layers $\max_j \sum_{i \in S_j} T_i$1 as a fused block,
• Minimizing total inference time via DP over all valid choices of block boundaries and assignments.

Table: Representative End-to-End Results for Fused-Layer Partitioning (Li et al., 2022)

Platform	ES count	DPFP Inference Time (ms)	Speedup over Standalone
RTX 2080 Ti	2	2.34	up to 73%
RTX 2080 Ti	7	1.67	up to 73%
AGX Xavier	7	8.72	strong scaling, diminishing returns with ES $\max_j \sum_{i \in S_j} T_i$2 7

Communication volume is reduced by approximately $\max_j \sum_{i \in S_j} T_i$3 compared to naive layer-level scattering approaches, with inference speedup saturating as ES count exceeds $\max_j \sum_{i \in S_j} T_i$4.

5. Experimental Results and Observed Speedups

Profiling-based balanced segmentation on multiple Edge TPUs achieves up to $\max_j \sum_{i \in S_j} T_i$5 speedup versus the official compiler-generated pipeline, and up to $\max_j \sum_{i \in S_j} T_i$6 versus single-TPU baselines, enabled by finer load balance and the overlapping of data transfers with computation (Villarrubia et al., 2 Mar 2025).

ODiMO-based approaches demonstrate substantial latencies/energy improvements, with heterogeneous-SoC deployment yielding

• Up to $\max_j \sum_{i \in S_j} T_i$7 energy reduction or $\max_j \sum_{i \in S_j} T_i$8 latency reduction (CIFAR-10/Tiny-ImageNet/DIANA SoC) at $\max_j \sum_{i \in S_j} T_i$9 drop in accuracy (Risso et al., 2023).
• For the Darkside SoC, latency is reduced by up to $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$0, and energy savings up to $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$1 have been demonstrated (Risso et al., 2024).

Fused-layer DP-based segmentation yields $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$2 ms inference time on VGG-16 ($\sum_{i \in S_j} M_i \leq M_{\text{cap}}$3 ES, RTX 2080 Ti, 100 Gbps interconnect), and consistently achieves reliability $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$4 under severe uplink variation (Li et al., 2022).

6. Design Trade-Offs, Limitations, and Best Practices

Balanced segmentation exposes trade-offs among compute-communication overlap, hardware memory limits, and quantization precision:

• Segment boundary placement should avoid splitting after wide layers with large activation footprints, as this increases communication cost.
• Level of parallelism: For $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$5, network transfer latency may dominate further segmentation benefits; excessive fragmentation induces "pipeline bubbles" due to non-overlapped communication.
• Hardware heterogeneity: Channel-level or spatial splits with quantization/fake-quantization during mapping enable accuracy-energy/latency Pareto balancing, but may yield convergence issues for extreme $\sum_{i \in S_j} M_i \leq M_{\text{cap}}$6 or highly non-uniform accelerator support.
• Precision maintenance: Layer/channel assignment should respect precision/quantization sensitivity—input/output layers are typically assigned to high-precision compute units.
• Halo management: For spatial-split strategies, only minimal region exchange is optimal; further increasing block number can diminish speedup due to communication overhead.

ODiMO's channel-reordering post-processing (ensuring contiguous sub-layer execution on each accelerator) and DPFP's receptive field based boundaries are key for observation of full functional correctness (Risso et al., 2023, Li et al., 2022).

7. Extensions and Applicability across Platforms

Balanced segmentation methodologies are generalizable to diverse hardware configurations:

• Profiling-driven contiguous segmentation applies to identical multi-accelerator arrays (e.g., multiple TPUs), with pipeline overlapping extending efficacy to bandwidth-limited USB/PCIe-attached devices (Villarrubia et al., 2 Mar 2025).
• Differentiable mapping is extensible to arbitrary SoCs/hardware accelerators, assuming analytic/simulated cost models, per-operator/quantization support enumerated, and differentiable channel-wise assignment variables (Risso et al., 2024, Risso et al., 2023).
• Receptive field-aware segmentation with halo-aware communication is directly relevant for spatially partitioned distributed inference in collaborative edge/fog architectures (Li et al., 2022).

A common limitation is the current assumption of either strict contiguity (profiled methods) or small-N parallelism; extending to non-contiguous, globally asynchronous, or multi-branch/trunked networks requires further research.

---

References

• "Balanced segmentation of CNNs for multi-TPU inference" (Villarrubia et al., 2 Mar 2025)
• "Optimizing DNN Inference on Multi-Accelerator SoCs at Training-time" (Risso et al., 2024)
• "Precision-aware Latency and Energy Balancing on Multi-Accelerator Platforms for DNN Inference" (Risso et al., 2023)
• "Receptive Field-based Segmentation for Distributed CNN Inference Acceleration in Collaborative Edge Computing" (Li et al., 2022)