Learn » Tutorial » Learning about collection reductions

// Forward declaration for message
struct ColRedMsg;

//                  VT Base Class for a collection
//            \------------------------------------------/
//             \                                        /
//              \                              Index   /
//               \                          \---------/
//                \                          \       /
struct ReduceCol : ::vt::Collection<ReduceCol,Index1D> {

  void reduceHandler();
  void reduceTarget(int value);

};

//                  VT Base Message for Collections
//               \-----------------------------------/
//                \                                 /
struct ColRedMsg : ::vt::CollectionMessage<ReduceCol> { };

void ReduceCol::reduceHandler() {
  auto cur_node = theContext()->getNode();
  (void)cur_node;  // don't warn about unused variable
  auto idx = this->getIndex();
  (void)idx;  // don't warn about unused variable

  //::fmt::print("MyCol::reduceHandler index={}, node={}\n", idx.x(), cur_node);

  auto proxy = getCollectionProxy();
  int val = 100;

  // Invoke the reduce!
  proxy.allreduce<&ReduceCol::reduceTarget, collective::PlusOp>(val);
}

int number_of_elements = 32;

void ReduceCol::reduceTarget(int value) {
  fmt::print("collection reduce value={}\n", value);
  assert(number_of_elements * 100 == value);

}

// Tutorial code to demonstrate reducing a collection
static inline void collectionReduce() {
  NodeType const this_node = ::vt::theContext()->getNode();
  NodeType const num_nodes = ::vt::theContext()->getNumNodes();
  (void)num_nodes;  // don't warn about unused variable

  /*
   * This is an example of reducing over a virtual context collection
   */

  if (this_node == 0) {
    // Range of `number_of_elements` elements for the collection
    auto range = vt::Index1D(number_of_elements);

    // Construct the collection in a rooted manner. By default, the elements
    // will be block mapped to the nodes
    auto proxy = vt::makeCollectionRooted<ReduceCol>("tutorial_collection_reduce")
      .bounds(range)     // Set the bounds for the collection
      .bulkInsert()      // Bulk insert all the elements within the bounds
      .wait();           // Wait for construction and get the proxy back

    // Broadcast a message to the entire collection. The reduceHandler will be
    // invoked on every element to the collection
    proxy.broadcast<&ReduceCol::reduceHandler>();
  }
}