Skip to content

Commit e036b85

Browse files
author
Matt Davis
committed
[IRON] PacketFifo: add end-to-end lowering tests against an aie.device context
Address review issue 5. The previous test suite covered the construction surface, registry integration, and host-side behavioural simulation, but didn't exercise PacketFifo.resolve() against a real MLIR context — so a mismatch with the dialect's packetflow() signature would not have been caught. Add two new tests in test/iron/test_packet_fifo.py: - test_resolve_emits_packetflow_op_in_module: builds a 2-producer / 1-consumer PacketFifo inside an aie.device body, calls resolve(), and asserts the resulting MLIR contains 2 aie.packetflow ops (one per producer). - test_resolve_idempotent_does_not_emit_twice: verifies the resolve() reentrancy guard. Pattern mirrors the existing test_resolve_emits_cascade_flow_op_in_module in test_cascade_fifo.py. Verified the call-site signature matches dialects.aie.packetflow's __init__ (pkt_id, source, source_port, source_channel, dests, keep_pkt_header).
1 parent c14c5e2 commit e036b85

1 file changed

Lines changed: 70 additions & 0 deletions

File tree

test/iron/test_packet_fifo.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -575,3 +575,73 @@ def test_packet_fifo_module_is_aie_iron_packet():
575575
from aie.iron import PacketFifo
576576

577577
assert PacketFifo.__module__.endswith(".packet")
578+
579+
# ---------------------------------------------------------------------------
580+
# Lowering tests — exercise resolve() against a real MLIR context to prove
581+
# the call into ``dialects.aie.packetflow`` matches the dialect's signature.
582+
# ---------------------------------------------------------------------------
583+
584+
def test_resolve_emits_packetflow_op_in_module():
585+
"""``resolve()`` inside an ``aie.device`` body emits one
586+
``aie.packetflow`` op per producer connecting to the consumer's tile.
587+
"""
588+
from aie.dialects.aie import AIEDevice, device, tile
589+
from aie.extras.context import mlir_mod_ctx
590+
from aie.iron import PacketFifo
591+
from aie.iron.device import Tile
592+
593+
with mlir_mod_ctx() as ctx:
594+
595+
@device(AIEDevice.npu2)
596+
def device_body():
597+
t_p0_op = tile(0, 2)
598+
t_p1_op = tile(0, 3)
599+
t_c_op = tile(0, 5)
600+
t_p0 = Tile(0, 2)
601+
t_p0.op = t_p0_op
602+
t_p1 = Tile(0, 3)
603+
t_p1.op = t_p1_op
604+
t_c = Tile(0, 5)
605+
t_c.op = t_c_op
606+
pf = PacketFifo(producers=[t_p0, t_p1], consumers=[t_c])
607+
pf.resolve()
608+
609+
module_str = str(ctx.module)
610+
611+
assert "aie.packetflow" in module_str, (
612+
f"expected aie.packetflow op in module after resolve(); "
613+
f"got:\n{module_str}"
614+
)
615+
# One packetflow per producer (2 producers -> 2 packetflow ops).
616+
assert module_str.count("aie.packetflow") == 2, (
617+
f"expected exactly 2 packetflow ops (one per producer); "
618+
f"got {module_str.count('aie.packetflow')} in:\n{module_str}"
619+
)
620+
621+
def test_resolve_idempotent_does_not_emit_twice():
622+
"""Calling ``resolve()`` twice on the same fifo does not double-emit."""
623+
from aie.dialects.aie import AIEDevice, device, tile
624+
from aie.extras.context import mlir_mod_ctx
625+
from aie.iron import PacketFifo
626+
from aie.iron.device import Tile
627+
628+
with mlir_mod_ctx() as ctx:
629+
630+
@device(AIEDevice.npu2)
631+
def device_body():
632+
t_p_op = tile(0, 2)
633+
t_c_op = tile(0, 5)
634+
t_p = Tile(0, 2)
635+
t_p.op = t_p_op
636+
t_c = Tile(0, 5)
637+
t_c.op = t_c_op
638+
pf = PacketFifo(producers=[t_p], consumers=[t_c])
639+
pf.resolve()
640+
pf.resolve() # second call must be a no-op.
641+
642+
module_str = str(ctx.module)
643+
644+
assert module_str.count("aie.packetflow") == 1, (
645+
f"expected exactly 1 packetflow op (resolve must be idempotent); "
646+
f"got {module_str.count('aie.packetflow')} in:\n{module_str}"
647+
)

0 commit comments

Comments
 (0)