feat: account for spatial-overflow rounds in symbolic latency model

[okaikov]

Hi,
While working with some hardware configurations with specific fanout limits, I noticed an interesting edge case in the symbolic latency model.

Currently, ComputeStats.combine_spatial in _stats.py uses max_nonzero() to determine max_latency. This works perfectly when the spatial bound is within the hardware fanout. However, when the bound exceeds the fanout (e.g., bound=512, fanout=240), the hardware must perform multiple sequential rounds to complete the operation. In this case, the latency should be scaled by ⌈bound/fanout⌉.

At the moment, AF models this as a single parallel batch (1 round), which leads to an underestimation of the total latency, even though the total_ops (MAC count) remains correct.

I've also noticed that attempting to model this manually by wrapping a Spatial loop in a Temporal loop on the same rank variable causes a crash in the mapping validator during tensor-view reordering.

I believe updating the symbolic engine to naturally handle these overflow rounds would be a great addition to the tool's accuracy. Would you be open to a change in _stats.py to account for this?

Suggested change (conceptual):
//We need to determine the number of rounds required based on the spatial bound
//and the hardware fanout.
rounds = math.ceil(current_spatial_bound / hardware_fanout)
self.max_latency = rounds * max_nonzero(self.max_latency, other.max_latency)

[tanner-andrulis]

This sounds like a spatial loop in a temporal loop. What specifically is the setup and error you are getting there?

[okaikov]

Thanks for the question — happy to clarify. I want to separate two distinct issues I conflated a bit in the original write-up.

---

Issue 1 — latency underestimation (the main bug)

This is independent of any Temporal/Spatial nesting. A single Spatial loop whose bound exceeds the hardware fanout produces wrong latency. Minimal reproducer with m=8, fanout=4 (so ⌈8/4⌉ = 2 sequential rounds needed):

# arch:     RegFile with spatial pe, fanout=4
# workload: C[m] = sum_k A[m,k]*B[k],  m=8, k=8
# mapping:
#   MainMemory  keep [A, B, C]
#   Spatial(m, tile_shape=1, name='pe', component=RegFile)  # needs 8 units, fanout=4 → 2 rounds
#   RegFile     keep [A, B, C]
#   Temporal(k, tile_shape=1)
#   Compute Matmul0

spec   = Spec.from_yaml('arch.yaml', 'workload.yaml', 'mapping.yaml')
result = evaluate_mapping(spec)
print(result.latency())  # prints 8 — should be 16 (2 rounds × 8 k-iterations)

total_ops=64 is correct; only max_latency is wrong. The root cause is in analyze_spatial: combine_spatial receives the child result after repeat_spatial(shape_repeats), which multiplies total_ops but leaves max_latency unchanged, then takes max_nonzero(0, child_latency) — giving 1 round regardless of overflow. The missing factor is ⌈shape_repeats / component_spatial_dim.fanout⌉, which needs to scale max_latency before or inside combine_spatial.

---

Issue 2 — crash when attempting the workaround

To your point about "a spatial loop inside a temporal loop": I did try that, and Temporal(m, tile=4) wrapping Spatial(m, tile=1) does work and gives the correct latency of 16 — so that path is viable as a manual workaround. I mis-stated this in the original report, sorry.

The crash I hit is a different scenario: pairing Spatial(m, tile=1) with Temporal(m, tile=1) to cover the same rank produces two loops with tile_shape=1 on the same rank variable:

# mapping:
#   MainMemory  keep [A, B, C]
#   Spatial(m, tile_shape=1, name='pe', component=RegFile)
#   Temporal(m, tile_shape=1)          ← both tile=1 on 'm'
#   RegFile     keep [A, B, C]
#   Temporal(k, tile_shape=1)
#   Compute Matmul0

Traceback (most recent call last):
File "model/main.py", line 317, in _assert_valid_pmapping
rank_variables.remove(node.rank_variable)
KeyError: 'm'

_assert_valid_pmapping iterates all loops with tile_shape == 1 and calls set.remove() for each rank variable. When two loops share the same rank variable and both have tile_shape=1, remove() is called twice and throws on the second call. Replacing remove() with discard() would make it safe — though the real fix is for the symbolic engine to handle overflow natively so the workaround isn't needed at all.

[tanner-andrulis]

Issue 1 is intended behavior. The loop has 8 iterations, so it takes 2 timesteps. The real problem is that the written mapping is not valid for the hardware; it uses 8 PEs when the hardware only has 4.

Issue 2 is a bug (thank you! will fix!), but the mapping is still not what you're looking for. The mapping you would like is:

# For rank size M=8
# mapping:
#   MainMemory  keep [A, B, C]
#   Temporal(m, tile_shape=4)
#   Spatial(m, tile_shape=1, name='pe', component=RegFile)
#   RegFile     keep [A, B, C]
#   Temporal(k, tile_shape=1)
#   Compute Matmul0

The outer m loop has tile shape 4, so when the inner loop runs, it has divides the size-4 tile into 4 size-1 tiles to activate all 4 PEs. Alternatively, the following would work:

# For rank size M=8
# mapping:
#   MainMemory  keep [A, B, C]
#   Spatial(m, tile_shape=2, name='pe', component=RegFile)
#   Temporal(m, tile_shape=1)
#   RegFile     keep [A, B, C]
#   Temporal(k, tile_shape=1)
#   Compute Matmul0

Like the above, this divides the full shape 8 into 4 size-2 tiles, one for each PE, then runs two temporal iterations beneath.

[okaikov]

Hey @tanner-andrulis , thanks for the detailed clarification — the correct mapping patterns make sense!

One thing I wanted to follow up on though: if the mapping is truly invalid (spatial bound > fanout with no outer temporal loop), shouldn't total_ops also be wrong?

Taking the minimal example at face value — m=8, fanout=4, no outer temporal loop — the hardware fires 4 PEs once and covers only 4 of the 8 m-values. So the actual MAC count should be: 4 PEs × 8 k-iterations = 32 MACs

But the tool reports total_ops=64, which implies it is modeling 2 overflow rounds internally — just not reflecting that in latency. So the two metrics are currently inconsistent with each other:

Metric	Tool reports	What the mapping actually describes
total_ops	64 ✅ (implies 2 rounds)	32
latency	8 ❌ (implies 1 round)	16

It seems like the tool is counting MACs as if overflow rounds happen, but latency as if they don't.

This makes me think there are really two valid paths forward:

1. Reject the mapping at validation time (consistent with "the mapping is invalid"), and ensure total_ops also reflects that — not silently compute a number that assumes overflow rounds.
2. Model overflow rounds consistently in both total_ops and latency.

Happy to be corrected if I'm misreading how total_ops is computed here — but the inconsistency seemed worth flagging!

[tanner-andrulis]

Bug report appreciated! I've pushed a fix for Issue 2, and am currently working on getting checks for Issue 1 so that it will raise an error & give fix instructions if the mapping is invalid.

For the most recent message, I'm implementing 1. Reject the mapping , but not for the reason of overflow rounds.

Right now, total ops is derivable from the workload (how many computations does the Einsum need?) and compute latency depends on the number of utilized hardware instances. The mapping written uses 8 PEs, so the mapping describes a latency of 8, and has no temporal loops, so there is only one round. However, this one round incorrectly uses more PEs than the hardware has available. This is why we need to throw out the mapping; it is invalid on the hardware.

Note that, in general, AccelForge is designed to avoid any implicit behaviors. The mapping explicitly describes how the workload is scheduled spatially and temporally on the hardware, and things like overflow rounds would be an implicit behavior.

[tanner-andrulis]

Pushed fixes to both issues. Thank you!