Example Jacobi 1D Program
The full code for this example can be found here: examples/collection/jacobi1d_vt.cc.
// // This code applies a few steps of the Jacobi iteration to // the linear system A x = 0 // where A is a tridiagonal matrix with pattern [-1 2 -1] // The initial guess for x is a made-up non-zero vector. // The exact solution is the vector 0. // // The matrix A is square and invertible. // The number of rows is ((number of objects) * (number of rows per object)) // // Such a matrix A is obtained when using 2nd-order finite difference // for discretizing (-d^2 u /dx^2 = f) on [0, 1] with homogeneous // Dirichlet condition (u(0) = u(1) = 0) using a uniform grid // with grid size 1 / ((number of objects) * (number of rows per object) + 1) // #include <vt/transport.h> #include <cstdlib> static constexpr std::size_t const default_nrow_object = 8; static constexpr std::size_t const default_num_objs = 4; static constexpr double const default_tol = 1.0e-02; struct NodeObj { bool is_finished_ = false; void workFinishedHandler() { is_finished_ = true; } bool isWorkFinished() { return is_finished_; } }; using NodeObjProxy = vt::objgroup::proxy::Proxy<NodeObj>; struct LinearPb1DJacobi : vt::Collection<LinearPb1DJacobi,vt::Index1D> { private: std::vector<double> tcur_, told_; std::vector<double> rhs_; size_t iter_ = 0; size_t msgReceived_ = 0, totalReceive_ = 0; size_t numObjs_ = 1; size_t numRowsPerObject_ = 1; size_t maxIter_ = 8; NodeObjProxy objProxy_; public: explicit LinearPb1DJacobi() : tcur_(), told_(), rhs_(), iter_(0), msgReceived_(0), totalReceive_(0), numObjs_(1), numRowsPerObject_(1), maxIter_(8) { } using BlankMsg = vt::CollectionMessage<LinearPb1DJacobi>; struct LPMsg : vt::CollectionMessage<LinearPb1DJacobi> { size_t numObjects = 0; size_t nRowPerObject = 0; size_t iterMax = 0; NodeObjProxy objProxy; LPMsg() = default; LPMsg(const size_t nobjs, const size_t nrow, const size_t itMax, NodeObjProxy proxy) : numObjects(nobjs), nRowPerObject(nrow), iterMax(itMax), objProxy(proxy) { } }; void checkCompleteCB(double normRes) { // // Only one object for the reduction will visit // this function // auto const iter_max_reached = iter_ > maxIter_; auto const norm_res_done = normRes < default_tol; if (iter_max_reached or norm_res_done) { auto const to_print = iter_max_reached ? "\n Maximum Number of Iterations Reached. \n\n" : fmt::format("\n Max-Norm Residual Reduced by {} \n\n", default_tol); fmt::print(to_print); // Notify all nodes that computation is finished objProxy_.broadcast<&NodeObj::workFinishedHandler>(); } else { fmt::print(" ## ITER {} >> Residual Norm = {} \n", iter_, normRes); } } void doIteration() { iter_ += 1; // //--- Copy extremal values // tcur_[0] = told_[0]; tcur_[numRowsPerObject_+1] = told_[numRowsPerObject_+1]; // //---- Jacobi iteration step //---- A tridiagonal matrix = "tridiag" ( [-1.0 2.0 -1.0] ) //---- rhs_ right hand side vector // for (size_t ii = 1; ii <= numRowsPerObject_; ++ii) { tcur_[ii] = 0.5*(rhs_[ii] + told_[ii-1] + told_[ii+1]); } std::copy(tcur_.begin(), tcur_.end(), told_.begin()); // // Compute the maximum entries among the rows on this object // We do not take into account the "ghost" entries // as they may be "out of date". // double maxNorm = 0.0; for (size_t ii = 1; ii < tcur_.size()-1; ++ii) { double val = tcur_[ii]; maxNorm = (maxNorm > std::fabs(val)) ? maxNorm : std::fabs(val); } auto proxy = this->getCollectionProxy(); proxy.reduce<&LinearPb1DJacobi::checkCompleteCB, vt::collective::MaxOp>( proxy[0], maxNorm ); } struct VecMsg : vt::CollectionMessage<LinearPb1DJacobi> { using MessageParentType = vt::CollectionMessage<LinearPb1DJacobi>; vt_msg_serialize_if_needed_by_parent_or_type1(vt::IdxBase); VecMsg() = default; VecMsg(vt::IdxBase const& in_index, double const& ref) : vt::CollectionMessage<LinearPb1DJacobi>(), from_index(in_index), val(ref) { } template <typename Serializer> void serialize(Serializer& s) { MessageParentType::serialize(s); s | from_index; s | val; } vt::IdxBase from_index = 0; double val = 0.0; }; void exchange(VecMsg *msg) { // Receive and treat the message from a neighboring object. const vt::IdxBase myIdx = getIndex().x(); if (myIdx > msg->from_index) { this->told_[0] = msg->val; msgReceived_ += 1; } if (myIdx < msg->from_index) { this->told_[numRowsPerObject_ + 1] = msg->val; msgReceived_ += 1; } // Check whether this 'object' has received all the expected messages. if (msgReceived_ == totalReceive_) { msgReceived_ = 0; doIteration(); } } void doIter([[maybe_unused]] BlankMsg *msg) { // // Treat the particular case of 1 object // where no communication is needed. // Without this treatment, the code would not iterate. // if (numObjs_ == 1) { doIteration(); return; } //--------------------------------------- // // Routine to send information to a different object // vt::IdxBase const myIdx = getIndex().x(); //--- Send the values to the left auto proxy = this->getCollectionProxy(); if (myIdx > 0) { proxy[myIdx - 1].send<VecMsg, &LinearPb1DJacobi::exchange>( myIdx, told_[1] ); } //--- Send values to the right if (size_t(myIdx) < numObjs_ - 1) { proxy[myIdx + 1].send<VecMsg, &LinearPb1DJacobi::exchange>( myIdx, told_[numRowsPerObject_] ); } } void init() { tcur_.assign(numRowsPerObject_ + 2, 0.0); told_.assign(numRowsPerObject_ + 2, 0.0); rhs_.assign(numRowsPerObject_ + 2, 0.0); double h = 1.0 / (numRowsPerObject_ * numObjs_ + 1.0); int nf = 3 * int(numRowsPerObject_ * numObjs_ + 1) / 4; size_t const myIdx = getIndex().x(); for (size_t ii = 0; ii < tcur_.size(); ++ii) { double x0 = ( numRowsPerObject_ * myIdx + ii) * h; tcur_[ii] = sin(nf * M_PI * x0 * x0); } totalReceive_ = 2; if (myIdx == 0) { tcur_[0] = 0.0; totalReceive_ -= 1; } if (myIdx == numObjs_ - 1) { tcur_[numRowsPerObject_+1] = 0.0; totalReceive_ -= 1; } std::copy(tcur_.begin(), tcur_.end(), told_.begin()); } void init(LPMsg* msg) { numObjs_ = msg->numObjects; numRowsPerObject_ = msg->nRowPerObject; maxIter_ = msg->iterMax; objProxy_ = msg->objProxy; // Initialize the starting vector init(); } }; bool isWorkDone( vt::objgroup::proxy::Proxy<NodeObj> const& proxy){ auto const this_node = vt::theContext()->getNode(); return proxy[this_node].invoke<&NodeObj::isWorkFinished>(); } int main(int argc, char** argv) { size_t num_objs = default_num_objs; size_t numRowsPerObject = default_nrow_object; size_t maxIter = 8; std::string name(argv[0]); vt::initialize(argc, argv); vt::NodeType this_node = vt::theContext()->getNode(); vt::NodeType num_nodes = vt::theContext()->getNumNodes(); if (argc == 1) { if (this_node == 0) { fmt::print( stderr, "{}: using default arguments since none provided\n", name ); } num_objs = default_num_objs * num_nodes; } else if (argc == 2) { num_objs = static_cast<size_t>(strtol(argv[1], nullptr, 10)); } else if (argc == 3) { num_objs = static_cast<size_t>(strtol(argv[1], nullptr, 10)); numRowsPerObject = static_cast<size_t>(strtol(argv[2], nullptr, 10)); } else if (argc == 4) { num_objs = static_cast<size_t>(strtol(argv[1], nullptr, 10)); numRowsPerObject = static_cast<size_t>(strtol(argv[2], nullptr, 10)); maxIter = static_cast<size_t>(strtol(argv[3], nullptr, 10)); } else { fmt::print( stderr, "usage: {} <num-objects> <num-rows-per-object> <maxiter>\n", name ); return 1; } // Object group of all nodes that take part in computation // Used to determine whether the computation is finished auto grp_proxy = vt::theObjGroup()->makeCollective<NodeObj>("examples_jacobi1d"); // Create the decomposition into objects using BaseIndexType = typename vt::Index1D::DenseIndexType; auto range = vt::Index1D(static_cast<BaseIndexType>(num_objs)); auto col_proxy = vt::makeCollection<LinearPb1DJacobi>("examples_jacobi1d") .bounds(range) .bulkInsert() .wait(); vt::runInEpochCollective([col_proxy, grp_proxy, num_objs, numRowsPerObject, maxIter]{ col_proxy.broadcastCollective<LinearPb1DJacobi::LPMsg, &LinearPb1DJacobi::init>( num_objs, numRowsPerObject, maxIter, grp_proxy ); }); while(!isWorkDone(grp_proxy)){ vt::runInEpochCollective([col_proxy]{ col_proxy.broadcastCollective< LinearPb1DJacobi::BlankMsg, &LinearPb1DJacobi::doIter >(); }); vt::thePhase()->nextPhaseCollective(); } vt::finalize(); return 0; }