Non-Blocking Stream Operations

Note

Status: implemented on the fork’s main (sunwookim028/allo). This is a fork-local feature: it is not upstream (cornell-zhang/allo) yet and is pending a decision on whether to submit it as an upstream PR. Tracked in fork issue #5 (fork-local preserve list) and fork issue #13 (fork-vs-upstream feature map). The former feature/mesh-accelerator branch that carried the initial implementation was deleted 2026-07-15 after the code landed on main.

Overview

Allo’s dataflow programming model connects kernels via typed FIFO streams (Stream[T, depth]). By default streams are blocking: put() stalls the producer when the FIFO is full, and get() stalls the consumer when it is empty.

Non-blocking operations expose the FIFO’s full/empty status and return a success flag instead of blocking, enabling:

• Valid-ready handshake protocols (AXI-S, NoC credit-based flow control)

• Arbitration — a single kernel polling multiple request channels without deadlocking

• Decoupled memory/compute tiles with explicit message-passing control

New API

Four methods were added to the Stream type (allo/ir/types.py, lines 349–360):

# allo/ir/types.py (fork main)

class Stream:
    def put(self, data): ...          # blocking write (existing)
    def get(self)       -> T: ...     # blocking read  (existing)

    def try_put(self, data) -> int1:  # non-blocking write; returns 1 on success
    def try_get(self)  -> (T, int1):  # non-blocking read; returns (data, success)
    def empty(self)    -> int1:       # True when FIFO has no items
    def full(self)     -> int1:       # True when FIFO is at capacity

Typical usage pattern (spin-until-success):

# Producer: spin until FIFO accepts the token
sent: int1 = 0
while sent == 0:
    sent = req_valid[ct_id].try_put(MSG_WRITE)

# Consumer: poll before committing to a read
if not req_valid[id].empty():
    msg_type, has_req = req_valid[id].try_get()

Why Non-Blocking Is Necessary for Arbitration

Consider a many-to-one topology: two Memory Tiles (MT0, MT1) share one Compute Tile (CT). With blocking get(), if CT blocks waiting for MT0 while MT1 sends first, the system deadlocks permanently.

Non-blocking polling resolves this:

while served < N_MTS:
    if not req_from_mt[0].empty():
        msg, ok = req_from_mt[0].try_get()
        if ok: handle(0)
    if not req_from_mt[1].empty():
        msg, ok = req_from_mt[1].try_get()
        if ok: handle(1)

Working Examples

All examples live in the tests/dataflow/ directory on the fork’s main.

File	Description
tests/dataflow/test_stream_nb_simple.py	Unit tests for try_put / try_get / empty / full with the Allo simulator. Six self-contained test cases, each exercising one primitive. Start here.
tests/dataflow/test_decoupled_mesh.py	End-to-end message-passing tests: 1-CT single-handshake protocol and 2×1 decoupled mesh. Demonstrates the full valid-ready handshake pattern.
tests/dataflow/test_stream_ops_ir.py	IR-level checks — verifies that the four new MLIR ops are emitted correctly.
tests/dataflow/test_stream_ops_sim.py	Simulator-level round-trip tests (producer → consumer via try_put / try_get).
tests/dataflow/test_stream_ops_hls.py	HLS codegen checks — verifies nb_write() / nb_read() / .empty() appear in the emitted Vivado HLS C++. (Its test_tapa_stream_nb case is marked xfail(strict=True): TAPA non-blocking support was removed, so that build raises RuntimeError. See Layer 4b below.)

Running the Simulator Tests

No HLS tool required:

conda activate allo
python -m pytest tests/dataflow/test_stream_nb_simple.py     -v
python -m pytest tests/dataflow/test_decoupled_mesh.py       -v
python -m pytest tests/dataflow/test_stream_ops_sim.py       -v

Running the HLS Codegen Tests

Requires Vitis HLS (brg-zhang-xcel server):

conda activate allo
python -m pytest tests/dataflow/test_stream_ops_hls.py       -v

Implementation Internals

The implementation touches four layers of the compiler stack.

Layer 1 — Python frontend (allo/ir/types.py)

Stream.try_put / try_get / empty / full are stub methods that serve as syntactic targets for the AST builder. They have no runtime Python logic; the builder intercepts them and emits MLIR operations instead.

Location: allo/ir/types.py, lines 349–360 (fork main).

Layer 2 — MLIR operations (mlir/include/allo/Dialect/AlloOps.td)

Two new operations were added to the Allo MLIR dialect:

allo.stream_try_put

Takes (stream, indices, data); returns i1 success flag.

allo.stream_try_get

Takes (stream, indices); returns (data, i1 success_flag) — a two-result op.

These are defined in mlir/include/allo/Dialect/AlloOps.td and the generated Python bindings appear in mlir/build/include/allo/Bindings/dialects/_allo_ops_gen.py (StreamTryPutOp, StreamTryGetOp, lines 2287–2378).

Layer 3 — AST builder (allo/ir/builder.py)

The builder’s visit_Call handler was extended to recognise the four new method names on Stream objects and emit the corresponding MLIR ops:

Location: allo/ir/builder.py, lines 2622–2676 (fork main).

Key lowering details:

• try_put emits allo.stream_try_put; the i1 result is unpacked as the return value.

• try_get emits allo.stream_try_get; the two results (data, success) are tuple-bound and unpacked at the call site via Python destructuring (val, ok = stream.try_get()).

• empty / full are lowered to allo.stream_try_get / allo.stream_try_put with a probe-only mode (no data consumed/produced).

Layer 4a — Simulator backend (allo/backend/simulator.py)

The OpenMP-threaded simulator was extended to:

1. Interpret StreamTryPutOp and StreamTryGetOp by inspecting the head/tail pointers of the shared FIFO without blocking.

2. Flush shared pointers with openmp_d.FlushOp before and after every non-blocking access to prevent stale cache lines across OpenMP threads.

3. Recurse _process_function_streams into nested kernel calls (fixing a pre-existing deadlock for hierarchical dataflow regions).

Layer 4b — HLS backends

Vivado HLS (mlir/lib/Translation/EmitVivadoHLS.cpp):

// try_put → hls::stream::write_nb()
os << streamName << ".write_nb(" << dataArg << ")";

// try_get → hls::stream::read_nb()
os << streamName << ".read_nb(" << dataArg << ")";

// empty / full → hls::stream::empty() / full()
os << streamName << ".empty()";

TAPA: not supported. The TAPA overrides (try_write / try_read) were removed in 2026-07 — TAPA is not used in the mesh flow and the dead codepath was a maintenance burden (see Catapult SystemC Flow and the Wire). The hooks in EmitBaseHLS.h have empty default bodies, but they are never reached: the op falls through to visitUnhandledOp and the emitter reports “can’t be correctly emitted”, so a non-blocking op built for target="tapa" fails the build with RuntimeError: Failed to emit HLS code. Note the message it prints blames wrap_io, which is unrelated to the actual cause. Use the Vitis HLS target for non-blocking streams.

Catapult HLS (mlir/lib/Translation/EmitCatapultHLS.cpp):

// try_put → blocking ac_channel::write(); try_get → blocking ac_channel::read()
// nb_write / nb_read inside a spin-while loop segfault Catapult's go compile
// (LOOP-19), so the non-blocking ops are emitted as blocking calls and the
// success flag is hard-coded to true.
os << streamName << ".write(" << dataArg << ")";

// empty → !ac_channel::available(1)  (ac_channel has no .empty() in the
// synthesizable subset, EDG CIN-59)
os << "!" << streamName << ".available(1)";

HLS Synthesis Cost (Vitis HLS, U280, estimated)

Based on synthesis experiments in tests/dataflow/hls_synth_streams.py:

Design	LUT	FF	II
Blocking stream (FIFO depth 2)	1,417	248	7
Non-blocking stream (try_put/try_get)	1,457 (+2.8%)	260 (+4.8%)	7

Note

The FF and II figures in this table are labelled estimated and disagree with the measured record. The primary record, dev/records/vitis_nb_streams.rst (Vitis HLS 2023.2, U280, 2026-03-08), measured the same two designs at LUT 1417 / 1457 but FF 1325 / 1369, and reported pipeline II 7 for the blocking design and undef (statically unbounded spin-wait) for the non-blocking one. Where the two disagree, use the record.

Non-blocking operations add only ~O(1) LUTs per stream (one comparator + branch). The success bit propagation is the main scheduling concern in HLS — keep it in control logic, not on the datapath critical path.

Files Modified (PR Scope)

File	Change
allo/ir/types.py	Add try_put, try_get, empty, full stubs to Stream
allo/ir/builder.py	Lower the four new methods to MLIR ops in visit_Call
allo/ir/infer.py	Type-inference rules for StreamTryPutOp / StreamTryGetOp
mlir/include/allo/Dialect/AlloOps.td	Define allo.stream_try_put and allo.stream_try_get ops
mlir/include/allo/Translation/EmitBaseHLS.h	Declare virtual emitter hooks for non-blocking ops
mlir/lib/Translation/EmitVivadoHLS.cpp	Vivado HLS: emit write_nb / read_nb / empty / full
mlir/lib/Translation/EmitTapaHLS.cpp	TAPA: try_write / try_read were added here and later removed (2026-07)
allo/backend/simulator.py	Interpret non-blocking ops; OMP flush; recursive kernel injection
tests/dataflow/test_stream_nb_simple.py	Simulator unit tests (3 test cases)
tests/dataflow/test_decoupled_mesh.py	End-to-end 1-CT and 2×1 mesh tests
tests/dataflow/test_stream_ops_ir.py	MLIR op emission tests
tests/dataflow/test_stream_ops_sim.py	Simulator round-trip tests
tests/dataflow/test_stream_ops_hls.py	HLS codegen string-match tests

See Also

• Dataflow Programming — Dataflow programming model overview