From 5ba125a344c4c2f0fc7309260c0249fa8d9f2508 Mon Sep 17 00:00:00 2001 From: Joseph Schuchart Date: Fri, 21 Jun 2024 09:57:26 -0400 Subject: [PATCH 1/2] Add devicemap to spmm benchmark Signed-off-by: Joseph Schuchart --- examples/spmm/spmm_cuda.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/spmm/spmm_cuda.cc b/examples/spmm/spmm_cuda.cc index 9dcf5928e..6519f21e0 100644 --- a/examples/spmm/spmm_cuda.cc +++ b/examples/spmm/spmm_cuda.cc @@ -787,7 +787,11 @@ class SpMM25D { , a_rowidx_to_colidx_(a_rowidx_to_colidx) , b_colidx_to_rowidx_(b_colidx_to_rowidx) { this->set_priomap([this](const Key<3> &ijk) { return this->prio(ijk); }); // map a key to an integral priority value - + auto num_devices = ttg::device::num_devices(); + this->set_devicemap( + [num_devices](const Key<3> &ijk){ + return ((((uint64_t)ijk[0]) << 32) + ijk[1]) % num_devices; + }); // for each {i,j} determine first k that contributes AND belongs to this node, // initialize input {i,j,first_k} flow to 0 for (auto i = 0ul; i != a_rowidx_to_colidx_.size(); ++i) { From 5569bdb44df591f1197efe7e71d91442e8027502 Mon Sep 17 00:00:00 2001 From: Joseph Schuchart Date: Thu, 27 Jun 2024 09:46:01 -0400 Subject: [PATCH 2/2] Remove bogus output Thanks to Nilesh for catching this. Signed-off-by: Joseph Schuchart --- examples/spmm/spmm_cuda.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/spmm/spmm_cuda.cc b/examples/spmm/spmm_cuda.cc index 6519f21e0..469aef8f7 100644 --- a/examples/spmm/spmm_cuda.cc +++ b/examples/spmm/spmm_cuda.cc @@ -1348,7 +1348,6 @@ static void initBlSpHardCoded(const std::function &)> &keymap, a_colidx_to_rowidx[3].emplace_back(0); // A[0][3] A.setFromTriplets(A_elements.begin(), A_elements.end()); - std::cout << "A_elements.begin()" << A_elements.begin() << "A_elements.end()" << A_elements.end() << "\n"; if (buildRefs && 0 == rank) { Aref.setFromTriplets(Aref_elements.begin(), Aref_elements.end());