tvm_tirx.backend.cuda.lang.Pipeline

Contents

tvm_tirx.backend.cuda.lang.Pipeline#

class tvm_tirx.backend.cuda.lang.Pipeline(*args, **kwargs)#

Bases: object

A full/empty mbarrier pair for a software-pipelined data flow.

Pass barrier-type tags and Pipeline constructs and inits the barriers itself. Tags: "tma" (TMABar), "tcgen05" (TCGen05Bar), "mbar" (MBarrier). The barrier type and arrival count of each event stay explicit at the call site – e.g. Pipeline(pool, n, full="tma", empty="tcgen05", init_empty=NUM_CONSUMER).

Both signals are required: a Pipeline is a pair. For a one-way event (a pure “X happened” signal with no slot to recycle) use a bare barrier (TMABar/TCGen05Bar/MBarrier) directly – it has no empty side.

Parameters:

  • pool (SMEMPool) – Shared memory pool allocator.

  • stages (int) – Number of pipeline stages (barrier slots).

  • full (str) – Barrier-type tag for the full / empty signal (see above).

  • empty (str) – Barrier-type tag for the full / empty signal (see above).

  • init_full (int) – Expected arrival count for the full / empty barrier.

  • init_empty (int) – Expected arrival count for the full / empty barrier.

  • empty_phase_offset (int) – XORed into the empty barrier’s phase bit on every wait / arrive.

  • leader (Expr, optional) – Propagated to both barriers; defaults to thread 0 of the CTA.

Contents