Fixes to device hint implementation

Signed-off-by: Joseph Schuchart <[email protected]>
devreal · May 15, 2024 · 8164c44 · 8164c44
1 parent 4f86f7b
commit 8164c44
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 13 deletions.
diff --git a/examples/potrf/potrf.h b/examples/potrf/potrf.h
@@ -683,12 +683,12 @@ namespace potrf {
      * data movement between devices. This provides hints for load-balancing up front
      * and avoids movement of the TRSM result to GEMM tasks.
      */
-    auto devmap1 = [&](const key1& key) { return (key[0] / A.P()) % ttg::device::num_devices(); }
+    auto devmap1 = [&](const Key1& key) { return (key[0] / A.P()) % ttg::device::num_devices(); };
 
-    auto devmap2a = [&](const key2& key) { return (key[0] / A.P()) % ttg::device::num_devices(); }
-    auto devmap2b = [&](const key2& key) { return (key[1] / A.P()) % ttg::device::num_devices(); }
+    auto devmap2a = [&](const Key2& key) { return (key[0] / A.P()) % ttg::device::num_devices(); };
+    auto devmap2b = [&](const Key2& key) { return (key[1] / A.P()) % ttg::device::num_devices(); };
 
-    auto devmap3 = [&](const key3& key) { return (key[0] / A.P()) % ttg::device::num_devices(); }
+    auto devmap3 = [&](const Key3& key) { return (key[0] / A.P()) % ttg::device::num_devices(); };
 
     ttg::Edge<Key1, MatrixTile<T>> syrk_potrf("syrk_potrf"), disp_potrf("disp_potrf");
 
@@ -705,28 +705,28 @@ namespace potrf {
     tt_potrf->set_keymap(keymap1);
     tt_potrf->set_defer_writer(defer_write);
 #ifdef ENABLE_DEVICE_KERNEL
-    tt_potrf->set_devmap(devmap1);
+    tt_potrf->set_devicemap(devmap1);
 #endif // 0
 
     auto tt_trsm = make_trsm(A, disp_trsm, potrf_trsm, gemm_trsm, trsm_syrk, trsm_gemm_row, trsm_gemm_col, output);
     tt_trsm->set_keymap(keymap2a);
     tt_trsm->set_defer_writer(defer_write);
 #ifdef ENABLE_DEVICE_KERNEL
-    tt_trsm->set_devmap(devmap2a);
+    tt_trsm->set_devicemap(devmap2a);
 #endif // 0
 
     auto tt_syrk = make_syrk(A, disp_syrk, trsm_syrk, syrk_syrk, syrk_potrf, syrk_syrk);
     tt_syrk->set_keymap(keymap2b);
     tt_syrk->set_defer_writer(defer_write);
 #ifdef ENABLE_DEVICE_KERNEL
-    tt_syrk->set_devmap(devmap2b);
+    tt_syrk->set_devicemap(devmap2b);
 #endif // 0
 
     auto tt_gemm = make_gemm(A, disp_gemm, trsm_gemm_row, trsm_gemm_col, gemm_gemm, gemm_trsm, gemm_gemm);
     tt_gemm->set_keymap(keymap3);
     tt_gemm->set_defer_writer(defer_write);
 #ifdef ENABLE_DEVICE_KERNEL
-    tt_gemm->set_devmap(devmap3);
+    tt_gemm->set_devicemap(devmap3);
 #endif // 0
 
     /* Priorities taken from DPLASMA */

diff --git a/ttg/ttg/parsec/ttg.h b/ttg/ttg/parsec/ttg.h
@@ -1537,13 +1537,13 @@ namespace ttg_parsec {
         if (tt->devicemap) {
           int parsec_dev;
           if constexpr (std::is_void_v<keyT>) {
-            parsec_dev = ttg::device::ttg_device_to_parsec_device(tt->devicemap());
+            parsec_dev = detail::ttg_device_to_parsec_device(tt->devicemap());
           } else {
-            parsec_dev = ttg::device::ttg_device_to_parsec_device(tt->devicemap(tt->key));
+            parsec_dev = detail::ttg_device_to_parsec_device(tt->devicemap(task->key));
           }
           for (int i = 0; i < MAX_PARAM_COUNT; ++i) {
             /* only set on mutable data since we have exclusive access */
-            if (tc.in[i].flow_flags & PARSEC_FLOW_ACCESS_WRITE) {
+            if (tc.in[i]->flow_flags & PARSEC_FLOW_ACCESS_WRITE) {
               parsec_data_t *data = parsec_task->data[i].data_in->original;
               /* only set the preferred device if the host has the latest copy
                * as otherwise we may end up with the wrong data if there is a newer

diff --git a/ttg/ttg/util/meta.h b/ttg/ttg/util/meta.h
@@ -851,11 +851,11 @@ namespace ttg {
       template <typename Key, typename Return, typename Enabler = void>
       struct keymap;
       template <typename Key, typename Return>
-      struct keymap<Key, std::enable_if_t<!is_void_v<Key>>> {
+      struct keymap<Key, Return, std::enable_if_t<!is_void_v<Key>>> {
         using type = std::function<Return(const Key &)>;
       };
       template <typename Key, typename Return>
-      struct keymap<Key, std::enable_if_t<is_void_v<Key>>> {
+      struct keymap<Key, Return, std::enable_if_t<is_void_v<Key>>> {
         using type = std::function<Return()>;
       };
       template <typename Key, typename Return = int>