Class PipelineTmaAsync

include/pipeline.h:59–143

// "leader" Currently, it is optional to elect a leader for the Consumers
template <int Stages_, int ClusterX_, int ClusterY_, int ClusterZ_>
class PipelineTmaAsync {
 public:
  using FullBarrier = ClusterTransactionBarrier;
  using EmptyBarrier = ClusterBarrier;
  using ProducerBarrierType = FullBarrier::ValueType;
  using ConsumerBarrierType = EmptyBarrier::ValueType;
  static constexpr uint32_t Stages = Stages_;
  using PipelineState = PipelineState<Stages>;

  struct SharedStorage {
    FullBarrier full_barrier_[Stages];
    EmptyBarrier empty_barrier_[Stages];
  };

  enum class ThreadCategory {
    NonParticipant,
    Producer,
    Consumer,
    ProducerConsumer
  };

  struct Params {
    uint32_t transaction_bytes = 0;
    ThreadCategory role = ThreadCategory::NonParticipant;
    uint32_t is_leader = 0;
    uint32_t num_consumers = 0;
    int active_warps = 0;
  };

  // Constructor
  DEVICE PipelineTmaAsync(SharedStorage& storage, Params params)
      : params_(params),
        full_barrier_ptr_(&storage.full_barrier_[0]),
        empty_barrier_ptr_(&storage.empty_barrier_[0]) {
    int warp_idx = canonical_warp_idx();
    int lane_predicate = elect_one_sync();
    if (warp_idx == params.active_warps && lane_predicate == 1) {
      // Barrier FULL init
      for (int i = 0; i < Stages; ++i) {
        full_barrier_ptr_[i].init(1);
      }
      uint32_t const num_consumer_warpgroups_per_cluster =
          params_.num_consumers / WARP_GROUP_SIZE;
      uint32_t const multicast_consumer_arrival_count =
          (ClusterX_ + ClusterY_ - 1) * num_consumer_warpgroups_per_cluster;
      // Barrier EMPTY init
      for (int i = 0; i < Stages; ++i) {
        empty_barrier_ptr_[i].init(multicast_consumer_arrival_count);
      }
    }
    // Logic to optimally schedule Empty Arrives
    // Goal : To divide SYNCS Empty Arrival duty equally amongst the Warp-Group
    // (128 threads)
    dim3 block_id = block_id_in_cluster();
    static constexpr int cluster_size = ClusterX_ * ClusterY_ * ClusterZ_;
    static_assert(cluster_size <= MAX_CLUSTER_SIZE,
                  "ERROR : Cluster size too large !");