From 0ddaef6f4cc7efa1555cb08774e6089e505bbcca Mon Sep 17 00:00:00 2001 From: Joseph Schuchart Date: Fri, 16 Aug 2024 12:26:13 -0400 Subject: [PATCH] POTRF: add option to disable explicit device mapping hint Signed-off-by: Joseph Schuchart --- examples/potrf/potrf.h | 19 ++++++++++++++----- examples/potrf/testing_dpotrf.cc | 5 ++++- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/examples/potrf/potrf.h b/examples/potrf/potrf.h index d3a92cb2f..689a932e2 100644 --- a/examples/potrf/potrf.h +++ b/examples/potrf/potrf.h @@ -669,7 +669,8 @@ namespace potrf { template auto make_potrf_ttg(MatrixT& A, ttg::Edge>& input, - ttg::Edge>& output, bool defer_write) { + ttg::Edge>& output, bool defer_write, + bool enable_device_map = true) { using T = typename MatrixT::element_type; auto keymap1 = [&](const Key1& key) { return A.rank_of(key[0], key[0]); }; @@ -705,28 +706,36 @@ namespace potrf { tt_potrf->set_keymap(keymap1); tt_potrf->set_defer_writer(defer_write); #ifdef ENABLE_DEVICE_KERNEL - tt_potrf->set_devicemap(devmap1); + if (enable_device_map) { + 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_devicemap(devmap2a); + if (enable_device_map) { + 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_devicemap(devmap2b); + if (enable_device_map) { + 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_devicemap(devmap3); + if (enable_device_map) { + tt_gemm->set_devicemap(devmap3); + } #endif // 0 /* Priorities taken from DPLASMA */ diff --git a/examples/potrf/testing_dpotrf.cc b/examples/potrf/testing_dpotrf.cc index 2a5f6dacb..781cc417c 100644 --- a/examples/potrf/testing_dpotrf.cc +++ b/examples/potrf/testing_dpotrf.cc @@ -61,6 +61,9 @@ int main(int argc, char **argv) bool check = !cmdOptionExists(argv+1, argv+argc, "-x"); bool cow_hint = !cmdOptionExists(argv+1, argv+argc, "-w"); + /* whether we set a device mapping */ + bool enable_device_map = !cmdOptionExists(argv, argv+argc, "--default-device-map"); + // TODO: need to filter out our arguments to make parsec happy ttg::initialize(1, argv, nthreads); @@ -130,7 +133,7 @@ int main(int argc, char **argv) init_tt->set_keymap([&]() {return world.rank();}); auto plgsy_ttg = make_plgsy_ttg(A, N, random_seed, startup, topotrf, cow_hint); - auto potrf_ttg = potrf::make_potrf_ttg(A, topotrf, result, cow_hint); + auto potrf_ttg = potrf::make_potrf_ttg(A, topotrf, result, cow_hint, enable_device_map); auto result_ttg = make_result_ttg(A, result, cow_hint); auto connected = make_graph_executable(init_tt.get());