diff --git a/nn_dataflow/core/map_strategy.py b/nn_dataflow/core/map_strategy.py index 00ee303..6e7982f 100644 --- a/nn_dataflow/core/map_strategy.py +++ b/nn_dataflow/core/map_strategy.py @@ -201,10 +201,13 @@ def gen_nested_loop_desc(self): # Loop occupancies affect accesses. aocc = [util.prod(data_loops[dce].take(locc)) for dce in range(de.NUM)] - # Replication uses the single DRAM, gbuf, itcn. - for mhe in [me.DRAM, me.GBUF, me.ITCN]: + # Replication uses the single DRAM, gbuf. + for mhe in [me.DRAM, me.GBUF]: uaccess[mhe] = tuple(a * n * o for a, n, o in zip(access_unitpass[mhe], rcnt, aocc)) + # Itcn access is replicated across all PEs. + uaccess[me.ITCN] = tuple(a * rsz * o for a, o + in zip(access_unitpass[me.ITCN], aocc)) # Replication uses different PEs. regf scales with op replication, # i.e., affected by all loop occupancies. Also consider external # occupancy.