Decoupling a tensor's persistence backing from its consumption source: DRAM reads on a fused external output that is consumed from SRAM

[okaikov]

Decoupling a tensor's persistence backing from its consumption source: DRAM reads on a fused external output that is consumed from SRAM

Summary

When a tensor is produced and reduced on-chip (SRAM), consumed by the next Einsum from SRAM, and also written to DRAM once as a required external output, AF charges a full DRAM read of that tensor in addition to the DRAM write. The read appears even though the consumer reads the tensor from SRAM. It looks like AF ties a tensor's persistence backing and its consumption source to the same memory, so forcing the external-output write to DRAM (dram.keep) also makes the consumer's fetch resolve to DRAM.

Expected: DRAM read (AO) = 0. Observed: DRAM read (AO) ≈ DRAM write ≈ tensor size.

Is this decoupling expressible today, or is it a modeling limitation?

Use case

A tensor AO that is simultaneously:

1. an intermediate — produced by Einsum N, consumed by Einsum N+1;
2. a required external output — it must be written to DRAM once (checkpoint / offload / activation dump);
3. consumed on-chip — Einsum N+1 reads it from SRAM (the DRAM copy is a write-only drain, never read back).
   In hardware these are three roles of one tensor: SRAM is where it is produced, reduced, and consumed; DRAM is where a copy is persisted and never read. AF appears to be able to express "persist to DRAM, consume from DRAM" and "consume from SRAM, no persistence," but not "persist to DRAM and consume from SRAM, with the DRAM side write-only."

Expected vs. observed

For AO (materialized size 2048 bits in the repro below):

Metric	Expected	Observed
DRAM write (AO)	≈ size (one spill)	2048
DRAM read (AO)	0	2048
SRAM read (AO), consumer Einsum	≈ size × reuse	32768
DRAM read (AO), producer Einsum	0	0

The two numbers that isolate the issue: the consumer reads AO from SRAM (32768), and the producer has 0 DRAM reads (so this is not a reduction/accumulation artifact) — yet DRAM still shows a read of AO equal to its full size. There is no round-trip to attribute it to (consumption is on-chip) and no accumulation to attribute it to (producer DRAM read is zero). It is a backing-fetch of a tensor the consumer never fetches from DRAM.

Minimal reproduction

import accelforge as af
from accelforge.frontend.arch import Arch, Memory, Compute
 
# AO is produced+reduced in SRAM, consumed by the next einsum FROM SRAM,
# and also written to DRAM once as a required external output. Expect DRAM reads of AO = 0.
wl = af.Workload(rank_sizes={"M":16,"N0":16,"N1":16,"N2":16}, bits_per_value={"All":8}, einsums=[
    {"name":"producer","tensor_accesses":[
        {"name":"I","projection":["m","n0"]},{"name":"W0","projection":["n0","n1"]},
        {"name":"AO","projection":["m","n1"],"output":True}]},       # AO: intermediate + external output
    {"name":"consumer","tensor_accesses":[
        {"name":"AO","projection":["m","n1"]},{"name":"W1","projection":["n1","n2"]},
        {"name":"O","projection":["m","n2"],"output":True}]},
])
 
def mem(name, keep, may_keep="All", nrfa=None):
    t = {"keep":keep, "may_keep":may_keep}
    if nrfa is not None:
        t["no_refetch_from_above"] = nrfa
    return Memory(name=name, size="inf", leak_power=0, area=0, tensors=t,
        actions=[{"name":"read","energy":1,"throughput":"inf"},
                 {"name":"write","energy":1,"throughput":"inf"}])
 
comp = Compute(name="MAC", leak_power=0, area=0,
               actions=[{"name":"compute","energy":1,"throughput":1}])
 
# AO forced to DRAM (external output) AND kept in SRAM (produced+consumed on-chip); no re-fetch from above.
arch = Arch(nodes=[mem("dram", "~Intermediates | AO"),
                   mem("sram", "~dram | AO", nrfa="AO"),
                   comp])
 
m = af.Spec(arch=arch, workload=wl).map_workload_to_arch(print_progress=False)
d = m.to_dict()
g = lambda c: (d[c][0] if isinstance(d[c], (list, tuple)) else d[c])
 
def s(component, act, es=None):
    return sum(g(c) for c in m.columns
               if f"<SEP>action<SEP>{component}<SEP>AO<SEP>{act}" in c
               and (es is None or c.startswith(es + "<SEP>")))
 
print("DRAM write (AO)              :", s("dram", "write"))
print("DRAM read  (AO)              :", s("dram", "read"), " <-- expected 0")
print("SRAM read  (AO), consumer    :", s("sram", "read", "consumer"), " <-- consumer reads on-chip")
print("DRAM read  (AO), consumer    :", s("dram", "read", "consumer"))
print("DRAM read  (AO), producer    :", s("dram", "read", "producer"))

Output:

DRAM write (AO)              : 2048
DRAM read  (AO)              : 2048  <-- expected 0
SRAM read  (AO), consumer    : 32768  <-- consumer reads on-chip
DRAM read  (AO), consumer    : 2048
DRAM read  (AO), producer    : 0

Root cause (as far as we traced it)

The DRAM read is attributed at the point where AO is fetched from its backing store for consumption. In model/_looptree/accesses.py (around the parent-buffer accounting), a buffer whose parent is the backing store has its reads elided, and the reduction/consumption read is otherwise attributed to the tensor's backing. Because AO must persist to DRAM, dram.keep makes DRAM the highest keeper and therefore the backing store; the consumer's single fetch of AO is then attributed to DRAM even though the mapping consumes AO from SRAM.

In other words, AF seems to assume one backing store per tensor, serving both persistence and consumption. When those two roles live in different memories — persist to DRAM, consume from SRAM — the model cannot represent the DRAM side as write-only, and a full-size backing read appears.

This is the standard activation-checkpoint / weight-offload pattern (persist a fused activation to DRAM while continuing to consume it on-chip), so it is likely to recur outside our setup.

What we tried

• sram.keep including AO (accumulation and consumption confirmed on-chip — consumer SRAM read is present and full).
• no_refetch_from_above="AO" (limits refetch; the residual one-pass DRAM read remains).
• Restricting may_keep / raising DRAM read cost (no effect — the read is not an opportunistic cache choice, it is the backing fetch).
• Node reordering and back narrowing (moves which node is the backing, but whichever memory persists AO becomes the consumption backing too).
  None of these decouple the two roles: as long as DRAM is where AO persists, DRAM is also where its consumer fetch is charged.

Question / requested capability

Is there a way to declare, for a single tensor, that DRAM is its persistence backing (mandatory write) while SRAM is its consumption source (all reads on-chip, DRAM read count = 0)? If this is expressible with existing keep / back / reservation semantics, we would appreciate a pointer to the correct spelling. If not, we would like to flag it as a modeling gap: AF cannot currently represent a write-only DRAM persist for a tensor that is consumed from SRAM.

Environment

• AccelForge 1.0.460
• Python 3.12.3
• Python API (not YAML)

[tanner-andrulis]

Right now Accelforge enforces, for each tensor, a single backing store, and for all Einsums this tensor starts and ends in the single backing store. Your results appear to be consistent with the backing store being DRAM and the two Einsums communicating with DRAM.

We do support the case you're speaking of (two Einsums communicate the tensor via SRAM, and we also want to write it back to DRAM), but this would require two backing stores, and a third Einsum. The Einsums would be the following:

einsums:
# First Einsum
- einsum: AO[m,n1] = I[m,n0] * W0[n0,n1]

# Copy the intermediate tensor into DRAM. AO_DRAM will be in DRAM, while AO will continue living in SRAM.
- einsum: AO_DRAM[m,n1] = AO[m,n1]
  is_copy_operation: True

# Second Einsum
- einsum: O[m,n2] = AO[m,n1] * W1[n1,n2]

Then you'd constrain AO to be backed in SRAM and AO_DRAM to be backed in DRAM.

Would this satisfy the problem you're trying to address?

(also FYI I'm working on a commit that'll rearrange all of the workload examples so if you're using any of those then copy the workloads you're using to somewhere local).