Merge branch 'master' into spmm_constraints

devreal · Jul 26, 2024 · d8d2597 · d8d2597
2 parents 2dd5905 + 7fb4027
commit d8d2597
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/examples/spmm/spmm_cuda.cc b/examples/spmm/spmm_cuda.cc
@@ -779,7 +779,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) {