Replace StaticWeightStream with ObjectFifos in MobileNet (#3129)

Author: James Yen

programming_examples/ml/mobilenet/bottleneck/post_l1.py

@@ -19,9 +19,9 @@
 
 from aie.iron import Kernel, ObjectFifo, Worker
 from aie.iron.controlflow import range_
+from aie.iron.dataflow.endpoint import ObjectFifoEndpoint
 
 from bottleneck._common import i8, load_wts
-from lowlevel_dma import StaticWeightStream
 from network_spec import block as nsblock
 
 
@@ -49,25 +49,26 @@ def post_l1(act_in, sf, *, placement, data_dir):
     # Input:  (7,1,80) int8   Output: (1,1,960) int8
     # ------------------------------------------------------------------
     # 76800 B of L1 weights are too large for the compute tile. Stage them on
-    # MemTile(4,1) and stream 9600 B chunks via StaticWeightStream.
+    # MemTile(4,1) and stream 9600 B chunks via ObjectFifo.
     PostOutputSplit = 8
     PostRepeatChannels = post_L1_InH  # = 7
     post_l1_wts_full_sz = post_L1_OutC * post_L1_InC  # 76800 B on MemTile
     post_l1_wts_chunk = post_l1_wts_full_sz // PostOutputSplit  # 9600 B per chunk
 
     post_l1_wts_data = load_wts(data_dir, "post_conv_chain.txt", post_l1_wts_full_sz)
 
-    post_l1_pb = StaticWeightStream(
-        obj_type=i8((post_l1_wts_full_sz,)),
-        initial_value=post_l1_wts_data,
+    post_l1_wts_of = ObjectFifo(
+        i8((post_l1_wts_full_sz,)),
+        depth=1,
         name="post_L1_wts",
-        recv_type=i8((post_l1_wts_chunk,)),
+        consumer_obj_type=i8((post_l1_wts_chunk,)),
+        init_values=[post_l1_wts_data.reshape(post_l1_wts_full_sz)],
         repeat_count=PostRepeatChannels,
-        memtile_placement=placement["memtile"],
-        compute_placement=placement["compute"],
-        mem_lock_id=2,
-        comp_lock_id=0,
     )
+    # Pin the producer (source of weight data) to a MemTile. Normally a
+    # Worker sets its fifo endpoint implicitly, but this fifo has no
+    # producing Worker.
+    post_l1_wts_of.prod().endpoint = ObjectFifoEndpoint(placement["memtile"])
 
     # Round-trip avgpool output through L3 (DDR) so it can be re-broadcast to
     # all 4 PostL2 FC tiles — a direct compute→4-compute fan-out exceeds
@@ -138,7 +139,7 @@ def post_l1_fn(
         fn_args=[
             act_in.cons(),
             act_out_post_avgpool_shim.prod(),
-            post_l1_pb,
+            post_l1_wts_of.cons(),
             k_post_l1,
             post_L1_InW,
             post_L1_InH,

programming_examples/ml/mobilenet/bottleneck/post_l2.py

@@ -18,9 +18,9 @@
 
 from aie.iron import Kernel, ObjectFifo, Worker
 from aie.iron.controlflow import range_
+from aie.iron.dataflow.endpoint import ObjectFifoEndpoint
 
 from bottleneck._common import i8, load_wts
-from lowlevel_dma import StaticWeightStream
 from network_spec import block as nsblock
 
 
@@ -30,8 +30,8 @@ def post_l2(act_in, sf, *, placement, data_dir):
     Args:
         act_in: ObjectFifo  — host-scratch fill of the avgpool output (uint16).
         sf: dict            — full scale-factor mapping; uses sf["POST"]["FC1"], ["FC2"].
-        placement: dict     — PLACEMENT["post_l2"] with keys "fc1_memtiles",
-                              "fc2_memtiles", "compute", "join_memtile".
+        placement: dict     — PLACEMENT["post_l2"] with keys "wts_memtiles",
+                              "compute", "join_memtile".
         data_dir: str       — directory holding FC{1,2}_{0..3}_chain.txt.
 
     Returns:
@@ -66,9 +66,6 @@ def post_l2(act_in, sf, *, placement, data_dir):
         depth=2,
     )
 
-    # FC weights: ping-pong across two adjacent MemTiles per compute tile (FC1
-    # on the even-column MemTile, FC2 on the odd-column MemTile that hosts the
-    # MM2S DMA — keeps DMA off (4,1) which post_L1 owns).
     PostOutputSplitL2 = 40
     fc_full_per_tile = post_L2_InC * fc_out_per_tile  # 409600 B per FC half
     fc_recv_per_tile = fc_full_per_tile // PostOutputSplitL2  # 10240 B on compute
@@ -83,8 +80,7 @@ def post_l2(act_in, sf, *, placement, data_dir):
         tile=placement["join_memtile"],
     )
 
-    fc1_memtiles = placement["fc1_memtiles"]
-    fc2_memtiles = placement["fc2_memtiles"]
+    wts_memtiles = placement["wts_memtiles"]
     fc_comptiles = placement["compute"]
 
     def _u16(shape):
@@ -111,24 +107,22 @@ def _u16(shape):
         fc1_data = load_wts(data_dir, fc1_f, fc_full_per_tile)
         fc2_data = load_wts(data_dir, fc2_f, fc_full_per_tile)
 
-        # FC weights: ping-pong across two adjacent memtiles. Primary = FC1
-        # (consumed first), ping-pong = FC2. DMA + primary share fc2_memtiles[i]
-        # (odd column) to avoid colliding with post_L1's MM2S on (4,1).
-        fc_pb = StaticWeightStream(
-            obj_type=i8((fc_full_per_tile,)),
-            initial_value=fc1_data,
-            name=f"fc1_wts_{i}",
-            recv_type=i8((fc_recv_per_tile,)),
-            repeat_count=PostOutputSplitL2,
-            memtile_placement=fc2_memtiles[i],
-            compute_placement=fc_comptiles[i],
-            s2mm_channel=1,
-            ping_pong_buf=(i8((fc_full_per_tile,)), fc2_data, f"fc2_wts_{i}"),
-            ping_pong_memtile=fc1_memtiles[i],
-            mem_lock_id=0,
-            comp_lock_id=2,
-            pp_lock_id=0,
+        # FC weights: depth-2 ObjectFifo with both FC1 and FC2 weight sets
+        # as init_values. The consumer loops through buffer 0 (FC1) then
+        # buffer 1 (FC2), each split into PostOutputSplitL2 chunks.
+        fc_wts_of = ObjectFifo(
+            i8((fc_full_per_tile,)),
+            depth=2,
+            name=f"post_L2_wts_{i + 1}",
+            consumer_obj_type=i8((fc_recv_per_tile,)),
+            init_values=[
+                fc1_data.reshape(fc_full_per_tile),
+                fc2_data.reshape(fc_full_per_tile),
+            ],
         )
+        # Pin the producer to a MemTile. Normally a Worker sets its fifo
+        # endpoint implicitly, but this fifo has no producing Worker
+        fc_wts_of.prod().endpoint = ObjectFifoEndpoint(wts_memtiles[i])
 
         def post_l2_fn(
             act_in,
@@ -161,7 +155,7 @@ def post_l2_fn(
             fn_args=[
                 act_in.cons(),
                 act_post_l2_tiles[i].prod(),
-                fc_pb,
+                fc_wts_of.cons(),
                 k_post_l2,
                 post_L1_OutC,
                 post_L2_InC,

programming_examples/ml/mobilenet/lowlevel_dma.py

@@ -20,8 +20,7 @@
     "init": Tile(0, 2),
     "post_l1": {"compute": Tile(6, 4), "memtile": Tile(4, 1)},
     "post_l2": {
-        "fc1_memtiles": [Tile(0, 1), Tile(2, 1), Tile(4, 1), Tile(6, 1)],
-        "fc2_memtiles": [Tile(1, 1), Tile(3, 1), Tile(5, 1), Tile(7, 1)],
+        "wts_memtiles": [Tile(1, 1), Tile(3, 1), Tile(5, 1), Tile(7, 1)],
         "compute": [Tile(6, 3), Tile(7, 4), Tile(7, 3), Tile(7, 2)],
         "join_memtile": Tile(6, 1),
     },