Cocco: Hardware-Mapping Co-Exploration towards Memory Capacity-Communication Optimization

On a neural network accelerator the on-chip memory is both the thing that saves you from expensive trips to DRAM and the thing that eats the silicon — a survey of sixteen commercial parts finds it occupying anywhere from 4 to 79 percent of the chip, which tells you nobody agrees on how much to build; Cocco (Tsinghua) argues the reason is that how much memory you need depends on how you carve the network up, and how you carve it up depends on how much memory you have, so the two have to be searched together.

Why this problem is worth solving

The technique everyone reaches for has a natural stopping point that nobody had located.

The other half is a decision usually made before the software exists.

What people did before, and where it broke

• Tile each layer to fit the buffer and reorder the loops for reuse. The mainstream of accelerator scheduling, refined for years. Where it broke. Reuse within a single layer is bounded. The output of a layer is written out and read straight back by the next one, which fusing avoids entirely and layer-level scheduling cannot.
• Enumerate every possible grouping and take the best. What the earliest fusion work does, later sped up by compressing the search state. Where it broke. Even the improved version faces exponential states for irregular network shapes, so it cannot finish on large models — and it certainly cannot be run once per candidate hardware configuration.
• Greedily merge the pair of groups that helps most, until nothing helps. Borrowed from a well-known compiler’s function grouping. Where it broke. It gets stuck in local optima, and more fundamentally its merge decisions are computed against a fixed hardware configuration, so it cannot participate in a hardware search at all.
• Order the layers by depth and solve with dynamic programming. The approach taken for irregular networks. Where it broke. It can only group layers that are contiguous in depth order, which is a genuinely restricted subset of the possible groupings, and its state transitions depend on a predefined buffer size — so again, no co-exploration.
• Sweep the buffer size, and run a partition search for each candidate. The obvious way to combine the two, and the paper implements it as a baseline in random and grid variants. Where it broke. Each capacity is searched from scratch, so nothing learned about good groupings at one size carries to another, and the search direction of a grid sweep determines how long convergence takes.
• The gap. Every partitioning method has a fixed buffer size wired into its decision rule or its state, which is precisely why the hardware and the mapping have always been chosen in sequence. The coupling is not subtle — bigger buffers permit bigger groups, and bigger groups are what justify bigger buffers.

How Cocco searches both at once

The idea

Nobody agrees how much memory an accelerator should have, and the survey makes that vivid. Across sixteen commercial neural processors, on-chip memory occupies anywhere from about 4 percent of the die to nearly 79 percent, with everything in between. That is not a spread you see when a question has been settled. The reason it stays open is that the answer depends on the workload and on how the workload is scheduled: more memory lets more consecutive layers execute as one unit, which is worth a great deal — fusing cuts external memory traffic by 42 to 75 percent and bandwidth demand by 27 to 68 percent on the models measured here. But the paper also shows the gains from ever-larger groups flattening quickly, so the useful question is not whether to fuse but where the curve bends, and that point moves with the network. Exploiting inter-layer reuse to escape the ceiling of what a single layer can offer is the same starting point as the memory-bound fusion work covered elsewhere on this site, taken here toward hardware sizing rather than compiler codegen.

Every existing partitioning method has a buffer size wired into it, which is precisely why the two were never searched together. This is the structural observation that justifies the framework, and it holds against each baseline for a different reason. The greedy method computes each merge’s benefit against a specific hardware configuration, so its decisions are only meaningful for that one. The dynamic programming method’s state transitions depend on a predefined capacity, so the recurrence itself changes when the capacity does. Enumeration is exponential in graph size and could not be repeated per candidate even if it finished. And the natural workaround — sweep capacities, run a partition search under each — is measured here and found wanting, because it searches each size from scratch and cannot carry what it learned about good groupings from one capacity to the next. Treating the partition and the capacity as a single object to be searched fixes all of this at once. Recognising that two choices usually made in sequence actually interact, and that co-optimising them beats optimising either, is the same finding as the format-and-schedule work covered elsewhere on this site.

The honest reading of the headline is that its two ends measure very different things. The paper reports lowering cost by between 1.89 and 50.33 percent, which is a wide range for a reason worth knowing. The 1.89 percent is against simulated annealing on one network — that is, against a real search method also doing co-optimisation with the same operators, where the claim is stability rather than quality. The 50.33 percent is against a fixed large buffer on another network, which is to say against not searching the hardware at all. Between those two poles sit the two-step sweeps. The paper is not hiding this; the table shows every comparison. But a reader who takes the range as a single number will overestimate what the genetic machinery buys over the nearest alternative and underestimate what co-exploration buys over none. The same candour appears in the multi-core study, which reports energy going up from one core to two because of communication overhead. Reporting where a coordinated optimisation helps a lot and where it barely helps at all is the same discipline as the coordinated graph transformation work covered elsewhere on this site.

Executing one group of layers

The result is either stalls or a buffer full of data cached in case somebody needs it. With more layers and more branches this compounds, because every producer is guessing at the union of what its consumers might want rather than being told.

Why it works

• Deriving from consumers avoids producing what nobody asked for. Pushing tiles forward means a producer must guess the union of its consumers’ demands and cache the surplus → whereas walking backwards lets each layer be told exactly what is required and how often → so the buffer holds what is needed rather than what might be.
• A least common multiple turns mismatched rates into a repeating schedule. Consumers with different kernel sizes and strides consume at different rates → so advancing the producer by the least common multiple makes one producer step correspond to whole numbers of consumer steps → and taking the co-prime solution gives the smallest unit that repeats, which is what the hardware issues.
• Regions plus sliding reuse remove recomputation entirely. Each layer owns a main and a side region, and window data still needed is kept in place → so intermediates inside a group are never written out and never recomputed → and only layers required by a later group ever reach external memory.
• The two decisions constrain each other, so a joint encoding beats a sweep. Bigger buffers permit bigger groups and bigger groups justify bigger buffers → so searching capacities independently discards everything learned about good groupings at the previous size → while a candidate carrying both lets crossover blend capacities and groupings together.
• Repairing invalid candidates beats rejecting them. A grouping paired with too small a capacity would ordinarily be discarded → but splitting the offending group turns it into a valid nearby candidate instead → which keeps the population useful and is why the search converges in tens of thousands of samples.

What it costs (trade-offs and scope)

• The headline range spans two incomparable comparisons. The 1.89 percent figure is against simulated annealing on one model — a real search method using the same operators — while the 50.33 percent is against a fixed large buffer on a different model, meaning against no hardware search at all. Both are in the paper’s tables, but the range should not be read as one number.
• Simulated annealing is close, and the claim against it is stability. The paper says the annealing baseline is not as stable across benchmarks rather than that it is much worse, and the smallest reported margin is exactly that comparison.
• Adding cores makes energy worse. Going from one core to two increases energy in most cases because of communication overhead, which the paper reports plainly. The compensating benefit is that per-core memory requirements fall as cores share weights.
• The best-known answer is unavailable where it matters most. Exhaustive enumeration finds the true optimum on the four simpler models and cannot finish on the four larger irregular ones — which are exactly the models where Cocco’s advantage over greedy and dynamic programming is claimed.
• The curve being climbed is already flattening. Fusing layers buys 42 to 75 percent of external memory traffic, but the paper’s own motivating measurement shows larger groups adding little beyond that, so the search operates in a region where differences are small.
• Cost is whatever the preference parameter says it is. The objective adds buffer size to a weighted mapping cost, and all headline percentages come from one setting of that weight. The sensitivity study shows the chosen capacity moving substantially as it changes.
• Everything is simulated. Costs come from a simulator built on two established accelerator evaluators, with area and energy from a 12 nanometre library. The one piece implemented in real silicon is the buffer region manager, in an in-house test chip, at 0.18 percent of core area.
• The searches are not cheap in absolute terms. Fifty thousand samples for co-exploration and four hundred thousand for partitioning alone — efficient relative to enumeration, not free.
• Some costs are assumed away. Scalar operations such as activation functions are treated as hidden in the pipeline and ignored, and the study is centred on a single accelerator core with multi-core treated as an extension.

The evaluation

Check yourself

Sources

The paper is “Cocco: Hardware-Mapping Co-Exploration towards Memory Capacity-Communication Optimization” by Zhanhong Tan, Zijian Zhu and Kaisheng Ma at Tsinghua University, published at ASPLOS 2024 Volume 1, pages 69-84, DOI 10.1145/3617232.3624865. It is open access under a Creative Commons Attribution licence.

This explainer is written from the full text, read from the authors’ copy on arXiv (2402.00629) alongside the published version. Everything above comes from the paper’s own sections 2 through 5: the survey of sixteen commercial neural processors with their memory capacities and area ratios, the note that on-chip memory costs one to two square millimetres per megabyte at 12 nanometres and that a large array costs dozens of times a multiply to read, and the measurement that fusing layers reduces external memory access by 42.3 to 74.7 percent and bandwidth by 26.8 to 67.8 percent while larger groups add little; the section 3 execution scheme, including the failure mode of the production-centric approach, the three stages that size output tiles then derive producer requirements in reverse topological order using least common multiples, the non-uniqueness of the update-count solution and the co-prime choice giving the minimal elementary operation, the main and side buffer regions with sliding-window reuse and no recomputation, and the region manager register file measured at 0.18 percent of core area in the authors’ test chip; the section 4 problem formulation of graph partition and its combination with buffer sizing under a preference parameter, together with the specific reasons enumeration, greedy merging and depth-ordered dynamic programming each cannot co-explore; the genetic encoding, crossover and four mutation operators, the repair of oversized groups during evaluation, and tournament selection; and the section 5 evaluation covering partition quality against enumeration where it can finish, the co-exploration comparison against fixed, two-step and annealing baselines with its full range of margins, the shared versus separate buffer comparison, the convergence and sampling analysis, the preference parameter sensitivity, and the multi-core and batch studies including the increase in energy from one core to two.