Feature Request: loop_order constraint on arch components to model fixed-dataflow hardware

[okaikov]

Summary

Most real DNN accelerators have a fixed internal dataflow — a hardware-determined loop execution order that cannot be changed at runtime (e.g., a raster scan: X → Y → Z → N, or a weight-stationary systolic array: K → C → X → Y). Currently, AccelForge has no way to express this in the arch definition. FFM generates and explores all valid loop orders during template generation, unaware that most of them are illegal on the target hardware. The arch API should provide a mechanism to declare the fixed loop order so FFM only searches over tile sizes within that constraint.

Motivation

The problem

Consider a custom DL accelerator with a MAC array that executes a fixed raster-scan dataflow internally:
for N: ← outermost (batch)
for Z: ← channels
for Y: ← spatial height
for X: ← spatial width (innermost — fastest changing)
MAC()

This order is etched into the hardware datapath. A mapping that iterates Z inside Y is physically impossible on this chip — it would require rewiring the memory bus at runtime.

The user wants to:
Fix the loop order (N → Z → Y → X) in the arch, as a hardware property, not a workload property.
Let FFM search for optimal tile sizes within that fixed order.
Get correct PPA estimates that reflect the actual hardware execution pattern.

What currently exists
The arch API provides Memory.tensors.tile_shape (Comparison objects) and Memory.tensors.no_refetch_from_above to constrain which loops go above or below a storage node — but nothing that fixes the order of temporal loops within a slot.

spec.mapper.explore_loop_orders = False reduces the number of loop orders FFM explores (it yields only one canonical ordering for partially-relevant rank variables), but:

It does not guarantee a specific user-defined order.
It applies globally, not per-component or per-rank.
It has no semantic connection to any arch property.

The only available workaround is the private/undocumented _pmapping_row_filter_function parameter in map_workload_to_arch(), which is a post-hoc filter applied after template generation — not a first-class arch constraint:

Current workaround — fragile and disconnected from the arch

def raster_scan_filter(row) -> bool:
... # inspect generated mapping and discard non-conforming templates

result = spec.map_workload_to_arch(
_pmapping_row_filter_function=raster_scan_filter
)

This approach:
Is not expressed in the arch, so the constraint is invisible to anyone reading the arch definition.
Does not prevent FFM from generating invalid templates — it only discards them after the fact, wasting search budget.
Uses an internal/unstable API.
Requires users to understand FFM's internal mapping representation to implement correctly.

[tanner-andrulis]

I believe that tensor_order_options under Storage objects is the property you would like. You can make something output-stationary by adding [Output, Input | Weight].

Can you give an example of a specific architecture whose loop order costraint is not captured by tensor_order_options? (Please just give an example this issue reads like a wall of Claude)

[okaikov]

tensor_order_options controls the relative order of tensor storage nodes, not temporal loop nodes. They are orthogonal concerns.

Here's a minimal counter-example. Consider a conv tile with rank variables X (width), Y (height), Z (channels).

My hardware executes a fixed raster-scan: Z (outer) → Y → X (inner). A mapping that is physically impossible on this chip is:

Temporal(Y, tile_shape=1) # Y outermost
Temporal(Z, tile_shape=1) # Z middle
Temporal(X, tile_shape=1) # X innermost
SRAM: Output, Input, Weight # tensor_order_options = [Output, Input|Weight]
op: MAC
This mapping satisfies tensor_order_options: Output is above Input and Weight. But the temporal loop order (Y outermost) is illegal on my hardware — the controller physically cannot iterate Y outside Z.

Now consider the only legal mapping on this chip:

Temporal(Z, tile_shape=1) # Z outermost
Temporal(Y, tile_shape=1) # Y middle
Temporal(X, tile_shape=1) # X innermost
SRAM: Output, Input, Weight # same tensor_order_options
op: MAC
Both mappings satisfy the same tensor_order_options, but only one is legal. The constraint "X must iterate inside Y and Z" cannot be expressed by storage node ordering.

The ask is for a way to specify, at the Arch level, "rank variable X must be the innermost temporal loop" (or a full fixed order like Z → Y → X). This is a hardware dataflow property, not a tensor placement property.

[tanner-andrulis]

I'm going to need the full workload, mapping, and arch specifications. For the workload, please share the Einsum, for the mapping, give all nodes (the storage nodes, loop nodes, and compute node), and for the arch, give all memory levels.

The reason we need specificity here is that "constraining temporal loops" is not a fully-defined problem because there may be multiple levels of tiling for every rank variable, so we would need to know what the loop order is AND which loops to constrain, of which there may be multiple options for each rank variable.

[okaikov]

[sorry for the multiple edits...]

Hi Tanner,

We need to enforce a specific temporal loop order in our hardware mapping (our controller hardcodes nz outermost, then ny, then nx). After investigation we believe tensor_order_options cannot express this constraint. Below is a minimal reproduction using AF's Python API.

---

Architecture

import accelforge as af
from accelforge.frontend.arch import Memory, Compute, Action, ArchNodes, Tensors

dram = Memory(
    name="DRAM",
    size=float("inf"),
    leak_power=0,
    area=0,
    actions=[
        Action(name="read", energy=1.0, throughput=float("inf")),
        Action(name="write", energy=1.0, throughput=float("inf")),
    ],
    tensors=Tensors(keep="~Intermediates", may_keep="All"),
)

gb = Memory(
    name="GlobalBuffer",
    size=float("inf"),
    leak_power=0,
    area=0,
    actions=[
        Action(name="read", energy=0.1, throughput=float("inf")),
        Action(name="write", energy=0.1, throughput=float("inf")),
    ],
    tensors=Tensors(
        keep="All",
        tensor_order_options=[["Output"], ["Input", "Weight"]],
    ),
)

mac = Compute(
    name="MAC",
    leak_power=0,
    area=0,
    actions=[Action(name="compute", energy=0.01, throughput=1.0)],
)

arch = af.Arch(nodes=ArchNodes([dram, gb, mac]))
Workload

from accelforge.frontend.workload import Workload, Einsum, TensorAccess

workload = Workload(
    einsums=[
        Einsum(
            name="test_einsum",
            renames={"NX": "x", "NY": "y", "NZ": "z"},
            tensor_accesses=[
                TensorAccess(name="Input", projection={"NX": "x", "NY": "y"}, bits_per_value=8),
                TensorAccess(name="Weight", projection={"NZ": "z"}, bits_per_value=8),
                TensorAccess(name="Output", projection={"NX": "x", "NY": "y", "NZ": "z"}, output=True, bits_per_value=32),
            ],
            rank_sizes={"NX": 8, "NY": 4, "NZ": 2},
        )
    ]
)

[okaikov]

Mapping A (hardware-illegal on our controller) — ny outermost:

from accelforge.frontend.mapping import Mapping, Temporal, Storage, Nested, Compute as MappingCompute

mapping_a = Mapping(nodes=[
Storage(tensors=["Input", "Weight", "Output"], component="DRAM"),
Nested(nodes=[
Storage(tensors=["Input", "Weight", "Output"], component="GlobalBuffer"),
Temporal(rank_variable="ny", tile_shape=1),
Temporal(rank_variable="nz", tile_shape=1),
Temporal(rank_variable="nx", tile_shape=1),
MappingCompute(einsum="test_einsum", component="MAC"),
])
])
Mapping B (hardware-legal) — nz outermost:

mapping_b = Mapping(nodes=[
Storage(tensors=["Input", "Weight", "Output"], component="DRAM"),
Nested(nodes=[
Storage(tensors=["Input", "Weight", "Output"], component="GlobalBuffer"),
Temporal(rank_variable="nz", tile_shape=1),
Temporal(rank_variable="ny", tile_shape=1),
Temporal(rank_variable="nx", tile_shape=1),
MappingCompute(einsum="test_einsum", component="MAC"),
])
])
Verification

spec_a = af.Spec(arch=arch, workload=workload, mapping=mapping_a)
spec_b = af.Spec(arch=arch, workload=workload, mapping=mapping_b)

Both return PASS

Conclusion

Both mappings have Output storage above Input/Weight in GlobalBuffer, so both satisfy tensor_order_options. AF structurally validates both as equally valid. The only difference is the order of two Temporal nodes inside Nested.

tensor_order_options constrains storage node ordering, not temporal loop ordering. On our hardware, only Mapping B is physically possible.

[tanner-andrulis]

That looks like input-stationary. Putting tensor_order_options: [[Input, Output | Weight]] will force the nz loop to be at the top for any mapping coming out of the mapper. For the relative ordering of ny and nx, order does not matter because permuting them will not change access counts for the mapping.

I know that in these examples, you are writing your own mappings, but neither given mapping makes much sense, because putting the storage nodes at the top for all three tensors makes them use tiles that are too big to get reuse. Moreover, if all the storage nodes are at the top, then for the given workload (with no partially-relevant ranks; defined in https://arxiv.org/abs/2602.15172), the loop order will not affect the mapping's access counts at all because we're reserving all tensors in the global buffer for all computations below, regardless of the order of those computations.

Take a look at https://arxiv.org/abs/2602.15172 Sections 4.1 and 4.2, which explains that in the first mapping, we should be moving the "Input" and "Output" storage nodes down, and in the second mapping, we should be moving the "Weight" and "Output" storage nodes down. The following are the mappings when using the optimal storage node placement:

# Mapping B (illegal)

mapping_a = Mapping(nodes=[
    Storage(tensors=["Input", "Weight", "Output"], component="DRAM"),
    Storage(tensors=["Weight"], component="GlobalBuffer"),
    Temporal(rank_variable="ny", tile_shape=1),
    Storage(tensors=["Input"], component="GlobalBuffer"),
    Temporal(rank_variable="nz", tile_shape=1),
    Temporal(rank_variable="nx", tile_shape=1),
    Storage(tensors=["Output"], component="GlobalBuffer"),
    MappingCompute(einsum="test_einsum", component="MAC"),
])

# Mapping B (legal)

mapping_b = Mapping(nodes=[
    Storage(tensors=["Input", "Weight", "Output"], component="DRAM"),
    Storage(tensors=["Input"], component="GlobalBuffer"),
    Temporal(rank_variable="nz", tile_shape=1),
    Storage(tensors=["Weight"], component="GlobalBuffer"),
    Temporal(rank_variable="ny", tile_shape=1),
    Temporal(rank_variable="nx", tile_shape=1),
    Storage(tensors=["Output"], component="GlobalBuffer"),
    MappingCompute(einsum="test_einsum", component="MAC"),
])

We can see that A does not respect tensor_order_options: [[Input, Output | Weight]], while B does.

[okaikov]

Thanks for the detailed explanation and the paper reference — really helpful context.

For the single-einsum case you described, I think you're right that tensor_order_options may be sufficient for energy estimation when there are no partially-relevant ranks.

Our main use case is actually fused multi-einsum workloads (cascaded convolutions), which is where we rely most on the auto-mapper. I wonder if loop order becomes non-semantic there? Since it affects how long intermediate tensors stay live between producer and consumer einsums, I'd expect the optimal tiling for one loop order to give a different energy than the same tiling under a different order — meaning a mapping optimized for a loop order our hardware can't execute might not reflect what we'd actually get.

Also curious about the ny vs nx case — you mentioned their relative order doesn't affect access counts, which makes sense energy-wise. But is there a way to constrain it for hardware correctness, even if PPA is the same? I noticed Timeloop has architecture_constraints with a temporal permutation field — would something like that be feasible to expose in AccelForge's Python API?

Happy to put together a more concrete fused workload example if that would help!

[tanner-andrulis]

This is a dense response, so I've included a few subheadings. In short, AccelForge does not differentiate equivalent loop orders, and if we were to constrain loop orders, we need a way of knowing what loops to constrain, which is a difficult problem. For this reason, we focus on constraining storage nodes, rather than loops. tensor_order_options lets you pick what tensors are stationary, and if there are violating loops, no_refetch_from_above and no_resend_to_below may also help you constrain those loops away. If you have an example where it is clear (and generalizes) what loops need to be constrained, these loops result in non-equivalent mappings, and this is not captured in existing constraints, then we can update the API to include that constraint.

AccelForge does not differentiate equivalent loop orders
In general, if there is not a storage node or reservation node between two loops in a mapping, then they can be rearranged without affecting any of the mapping's stats. This is true both for single-Einsum pmappings and for the full multi-Einsum LoopTree. The mapper leverages this fact heavily, and during both pmapping generation & fusion stages it treats multiple loop orders as equivalent.

So, to frame your issue here, I would say that you want constrain away cases where the loop order is invalid for your hardware and there are storage/reservation nodes between the violating loops. It would be helpful to see an example of this (in either a fused or non-fused mapping).

I know you also would like to "constrain it for hardware correctness, even if PPA is the same", but note that AccelForge is a PPA tool, and heavily relies on the fact that we can avoid skipping decisions that lead to same-or-worse PPA. From AccelForge's point of view, there is no difference between the nx and ny loop orderings, and so it would be very difficult to constrain that order in AccelForge. It would be easier if you imagined those mappings as equivalent and rearranged the loops as needed when programming on your target hardware.

It is difficult to tell what loops to constrain
Timeloop has a fundamentally different and more limited approach than AccelForge. Timeloop, for each memory level, puts all tensors in a single storage node and then creates a block of loops beneath. Within and only within the loops of one memory level may one constrain the loop order. AccelForge lets mappings create multiple separate storage nodes, one for each tensor, for each memory level, opening a larger mapspace (https://ieeexplore.ieee.org/document/9360462/). However, without this pre-blocking that Timeloop does (which AccelForge intentionally avoids because it leads to worse mappings), AccelForge does not have Timeloop's notion of a block of loops connected to a single memory level.

So, if we were to write a general loop order constraint, it would need to be applicable in general to a mapping with a stack of storage nodes (with perhaps >1 per memory level, storage nodes for a given memory level not necessarily back-to back, and memory levels not necessarily in order), and multiple loops for each rank variable. This is a difficult problem, and one we have not come up with a good abstraction for yet.