Multi-Threaded Simulator (CPU)

This document demonstrates how to simulate a dataflow design using the simulator backend in Allo. The simulator backend provides a fast and flexible environment for verifying the behavior of dataflow kernels before deploying them to hardware. In this example, a simple producer-consumer model is implemented, where data is produced from an input matrix, sent through a pipe, and then consumed with a basic arithmetic operation.

Dataflow Kernel Definition

The design consists of a top-level region that contains two kernels: a producer and a consumer. The producer reads data from an input matrix and sends each element through a pipe, while the consumer receives the data, increments it by one, and writes the result to an output matrix. The following code illustrates the kernel definitions using Allo’s dataflow API:

import allo
from allo.ir.types import float32, Stream
import allo.dataflow as df
import numpy as np

Ty = float32
M, N, K = 16, 16, 16

@df.region()
def top(A: Ty[M, N], B: Ty[M, N]):
    # Create a pipe with a depth of 4
    pipe: Stream[Ty, 4]

    @df.kernel(mapping=[1], args=[A])
    def producer(local_A: Ty[M, N]):
        for i, j in allo.grid(M, N):
            # Load data from the input matrix
            out: Ty = local_A[i, j]
            # Send data to the pipe
            pipe.put(out)

    @df.kernel(mapping=[1], args=[B])
    def consumer(local_B: Ty[M, N]):
        for i, j in allo.grid(M, N):
            # Receive data from the pipe
            data = pipe.get()
            # Perform a simple computation (increment by 1)
            local_B[i, j] = data + 1

Simulation and Testing

To verify the correctness of the dataflow design, a simulation is executed using the simulator backend. The test function initializes an input matrix with random values and an output matrix filled with zeros. The simulation is performed by building the module with the target set to “simulator”. After running the simulation, the output is compared against the expected result using NumPy’s testing utilities.

A = np.random.rand(M, N).astype(np.float32)
B = np.zeros((M, N), dtype=np.float32)
sim_mod = df.build(top, target="simulator")
sim_mod(A, B)
np.testing.assert_allclose(B, A + 1, rtol=1e-5, atol=1e-5)
print("Dataflow Simulator Passed!")

The simulator is implemented using the OMP dialect in MLIR, so it can natively support multi-threaded execution on CPU, which greatly speeds up functional testing at the first place.

Controlling OpenMP Threads in Simulator Tests

When running simulations or tests that use the simulator backend, the number of threads created by OpenMP can be controlled via the OMP_NUM_THREADS environment variable. This variable controls the number of CPU threads OpenMP uses for parallel execution.

Note

OMP_NUM_THREADS is a performance knob, not a correctness requirement. This page used to say that a dataflow design may need a large pool “to simulate properly (otherwise, you might experience hang or stuck)” and to advise raising it to 64 or 128. That requirement is gone: a region whose kernel instances outnumbered the pool used to wedge silently, because a process blocked on a stream held a thread its own producer needed, and since f193c057 the simulator sizes the OpenMP team to the region’s section count rather than to the core count. A region now runs whatever OMP_NUM_THREADS is set to. See item 11 of the limitations register for the measurement this replaces.

For example:

export OMP_NUM_THREADS=8
python3 tests/dataflow/test_1D_systolic.py

The OpenMP runtime determines the size of its thread pool during initialization based on the value of OMP_NUM_THREADS. Once the thread pool has been created, modifying OMP_NUM_THREADS afterward will not affect the existing pool size.

So if you want a particular pool size for a particular test, export it before the process starts, or adjust it individually within each process as needed.

Tip: If you are using pytest to run multiple tests and want each one’s OMP_NUM_THREADS to take effect, be sure to use the --forked option so that each test runs in a separate process and the OMP_NUM_THREADS setting in setup_env takes effect.

Example: To ensure that each test runs in a separate process and the thread pool is recreated according to OMP_NUM_THREADS, you can add the following code to each test file in the target folder:

import os
import pytest

@pytest.fixture(scope="module", autouse=True)
def setup_env():
    os.environ["OMP_NUM_THREADS"] = "8"
    yield
    del os.environ["OMP_NUM_THREADS"]

def test_xxx():
    # your test code here
    pass

Then run the tests with:

python3 -m pytest --forked /path/to/tests/folder -v