MueLu UncoupledAggregationFactory: Add hidden option "aggregation: ba…

…ckend" The option allows to select the aggregation code path: - "default": do whatever matches the input graph, - "host": old, host-only aggregation code path, - "kokkos": Kokkos code path. Signed-off-by: Christian Glusa <[email protected]>
trilinos · Jan 14, 2025 · 1777e78 · 1777e78
1 parent 6247bb7
commit 1777e78
Show file tree

Hide file tree

Showing 13 changed files with 105 additions and 16 deletions.
diff --git a/packages/muelu/doc/UsersGuide/masterList.xml b/packages/muelu/doc/UsersGuide/masterList.xml
@@ -376,6 +376,15 @@
 
   <aggregation>
 
+    <parameter>
+      <name>aggregation: backend</name>
+      <type>string</type>
+      <default>"default"</default>
+      <description>Aggregation scheme. Possible values: "default", "host", "kokkos"</description>
+      <visible>false</visible>
+      <comment-ML>parameter not existing in ML</comment-ML>
+    </parameter>
+
     <parameter>
       <name>aggregation: type</name>
       <type>string</type>

diff --git a/packages/muelu/doc/UsersGuide/paramlist_hidden.tex b/packages/muelu/doc/UsersGuide/paramlist_hidden.tex
@@ -62,6 +62,8 @@
 
 \cbb{coarse: overlap}{int}{0}{Coarse solver subdomain overlap.}
 
+\cbb{aggregation: backend}{string}{"default"}{Aggregation scheme. Possible values: "default", "host", "kokkos"}
+
 \cbb{aggregation: type}{string}{"uncoupled"}{Aggregation scheme. Possible values: see Table~\ref{t:aggregation}.}
 
 \cbb{aggregation: mode}{string}{"uncoupled"}{Controls whether aggregates are allowed to cross processor boundaries. Possible values: "coupled" aggregates can cross processor boundaries, "uncoupled" aggregates cannot cross processor boundaries.}

diff --git a/packages/muelu/src/Graph/Containers/MueLu_LWGraph_decl.hpp b/packages/muelu/src/Graph/Containers/MueLu_LWGraph_decl.hpp
@@ -15,6 +15,7 @@
 
 #include "MueLu_LWGraph_fwd.hpp"
 #include "MueLu_LWGraphBase.hpp"
+#include "MueLu_LWGraph_kokkos_fwd.hpp"
 
 namespace MueLu {
 
@@ -32,6 +33,8 @@ template <class LocalOrdinal  = DefaultLocalOrdinal,
 class LWGraph : public MueLu::LWGraphBase<LocalOrdinal, GlobalOrdinal, Node, true> {
  public:
   using LWGraphBase<LocalOrdinal, GlobalOrdinal, Node, true>::LWGraphBase;
+
+  RCP<MueLu::LWGraph_kokkos<LocalOrdinal, GlobalOrdinal, Node> > copyToDevice();
 };
 
 }  // namespace MueLu

diff --git a/packages/muelu/src/Graph/Containers/MueLu_LWGraph_def.hpp b/packages/muelu/src/Graph/Containers/MueLu_LWGraph_def.hpp
@@ -10,10 +10,35 @@
 #ifndef MUELU_LWGRAPH_DEF_HPP
 #define MUELU_LWGRAPH_DEF_HPP
 
+#include "MueLu_LWGraph_kokkos.hpp"
 #include "MueLu_LWGraph_decl.hpp"
 
 namespace MueLu {
 
+template <class LocalOrdinal, class GlobalOrdinal, class Node>
+RCP<MueLu::LWGraph_kokkos<LocalOrdinal, GlobalOrdinal, Node> > MueLu::LWGraph<LocalOrdinal, GlobalOrdinal, Node>::copyToDevice() {
+  auto graph = this->getGraph();
+
+  auto row_map_d = Kokkos::create_mirror_view(graph.row_map);
+  auto entries_d = Kokkos::create_mirror_view(graph.entries);
+  Kokkos::deep_copy(row_map_d, graph.row_map);
+  Kokkos::deep_copy(entries_d, graph.entries);
+
+  using local_graph_type_device = typename MueLu::LWGraphBase<LocalOrdinal, GlobalOrdinal, Node, false>::local_graph_type;
+  auto graph_d                  = local_graph_type_device(entries_d, row_map_d);
+
+  auto lw_d = rcp(new MueLu::LWGraph_kokkos<LocalOrdinal, GlobalOrdinal, Node>(graph_d, this->GetDomainMap(), this->GetImportMap(), this->getObjectLabel()));
+
+  using bndry_nodes_type = typename MueLu::LWGraphBase<LocalOrdinal, GlobalOrdinal, Node, false>::boundary_nodes_type;
+
+  auto bndry   = this->GetBoundaryNodeMap();
+  auto bndry_d = bndry_nodes_type("boundary_nodes", bndry.extent(0));
+  Kokkos::deep_copy(bndry_d, bndry);
+  lw_d->SetBoundaryNodeMap(bndry_d);
+
+  return lw_d;
+}
+
 }  // namespace MueLu
 
 #endif  // MUELU_LWGRAPH_DEF_HPP
diff --git a/packages/muelu/src/Graph/UncoupledAggregation/MueLu_UncoupledAggregationFactory_def.hpp b/packages/muelu/src/Graph/UncoupledAggregation/MueLu_UncoupledAggregationFactory_def.hpp
@@ -76,6 +76,8 @@ RCP<const ParameterList> UncoupledAggregationFactory<LocalOrdinal, GlobalOrdinal
   SET_VALID_ENTRY("aggregation: error on nodes with no on-rank neighbors");
   SET_VALID_ENTRY("aggregation: phase3 avoid singletons");
   SET_VALID_ENTRY("aggregation: phase 1 algorithm");
+  SET_VALID_ENTRY("aggregation: backend");
+  validParamList->getEntry("aggregation: backend").setValidator(rcp(new Teuchos::StringValidator(Teuchos::tuple<std::string>("default", "host", "kokkos"))));
 #undef SET_VALID_ENTRY
 
   // general variables needed in AggregationFactory
@@ -175,20 +177,53 @@ void UncoupledAggregationFactory<LocalOrdinal, GlobalOrdinal, Node>::Build(Level
   RCP<Aggregates> aggregates;
   RCP<const Teuchos::Comm<int>> comm;
   LO numRows;
+
+  const std::string aggregationBackend = pL.get<std::string>("aggregation: backend");
+
+  // "Graph" can have type "LWGraph" or "LWGraph_kokkos".
+  // The aggregation phases can call either "BuildAggregatesNonKokkos" or "BuildAggregates".
+
+  // "aggregation: backend" can take values "default", "non-Kokkos" or "Kokkos".
+  // "default": run depending on the type of "Graph"
+  // "non-Kokkos": run the non-Kokkos aggregation, moving "Graph" to host if necessary
+  // "Kokkos": run the Kokkos aggregation, potentially move "Graph", moving "Graph" to device if necessary
+
   bool runOnHost;
   if (IsType<RCP<LWGraph>>(currentLevel, "Graph")) {
-    graph      = Get<RCP<LWGraph>>(currentLevel, "Graph");
-    aggregates = rcp(new Aggregates(*graph));
-    comm       = graph->GetComm();
-    numRows    = graph->GetNodeNumVertices();
-    runOnHost  = true;
+    if ((aggregationBackend == "default") || (aggregationBackend == "non-Kokkos")) {
+      graph      = Get<RCP<LWGraph>>(currentLevel, "Graph");
+      aggregates = rcp(new Aggregates(*graph));
+      comm       = graph->GetComm();
+      numRows    = graph->GetNodeNumVertices();
+      runOnHost  = true;
+    } else {
+      RCP<LWGraph> tmp_graph = Get<RCP<LWGraph>>(currentLevel, "Graph");
+      graph_kokkos           = tmp_graph->copyToDevice();
+      aggregates             = rcp(new Aggregates(*graph_kokkos));
+      comm                   = graph_kokkos->GetComm();
+      numRows                = graph_kokkos->GetNodeNumVertices();
+      runOnHost              = false;
+    }
+  } else if (IsType<RCP<LWGraph_kokkos>>(currentLevel, "Graph")) {
+    if ((aggregationBackend == "default") || (aggregationBackend == "Kokkos")) {
+      graph_kokkos = Get<RCP<LWGraph_kokkos>>(currentLevel, "Graph");
+      aggregates   = rcp(new Aggregates(*graph_kokkos));
+      comm         = graph_kokkos->GetComm();
+      numRows      = graph_kokkos->GetNodeNumVertices();
+      runOnHost    = false;
+    } else {
+      RCP<LWGraph_kokkos> tmp_graph_kokkos = Get<RCP<LWGraph_kokkos>>(currentLevel, "Graph");
+      graph                                = tmp_graph_kokkos->copyToHost();
+      aggregates                           = rcp(new Aggregates(*graph));
+      comm                                 = graph->GetComm();
+      numRows                              = graph->GetNodeNumVertices();
+      runOnHost                            = true;
+    }
   } else {
-    graph_kokkos = Get<RCP<LWGraph_kokkos>>(currentLevel, "Graph");
-    aggregates   = rcp(new Aggregates(*graph_kokkos));
-    comm         = graph_kokkos->GetComm();
-    numRows      = graph_kokkos->GetNodeNumVertices();
-    runOnHost    = false;
+    TEUCHOS_TEST_FOR_EXCEPTION(true, std::invalid_argument, "Graph has bad type.");
+  }
 
+  if (!runOnHost) {
     TEUCHOS_TEST_FOR_EXCEPTION(pL.get<bool>("aggregation: use interface aggregation"), std::invalid_argument, "Option: 'aggregation: use interface aggregation' is not supported in the Kokkos version of uncoupled aggregation");
     // Sanity Checking: match ML behavior is not supported in UncoupledAggregation_Kokkos in Phase 1 , but it is in 2a and 2b
     TEUCHOS_TEST_FOR_EXCEPTION(pL.get<bool>("aggregation: match ML phase1"), std::invalid_argument, "Option: 'aggregation: match ML phase1' is not supported in the Kokkos version of uncoupled aggregation");

diff --git a/packages/muelu/src/Interface/MueLu_ParameterListInterpreter_def.hpp b/packages/muelu/src/Interface/MueLu_ParameterListInterpreter_def.hpp
@@ -1099,12 +1099,10 @@ void ParameterListInterpreter<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
     MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: min agg size", int, aggParams);
     MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: max agg size", int, aggParams);
     MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: max selected neighbors", int, aggParams);
-    if (useKokkos_) {
-      // if not using kokkos refactor Uncoupled, there is no algorithm option (always Serial)
-      MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: phase 1 algorithm", std::string, aggParams);
-      MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: deterministic", bool, aggParams);
-      MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: coloring algorithm", std::string, aggParams);
-    }
+    MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: backend", std::string, aggParams);
+    MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: phase 1 algorithm", std::string, aggParams);
+    MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: deterministic", bool, aggParams);
+    MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: coloring algorithm", std::string, aggParams);
     MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: enable phase 1", bool, aggParams);
     MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: enable phase 2a", bool, aggParams);
     MUELU_TEST_AND_SET_PARAM_2LIST(paramList, defaultList, "aggregation: enable phase 2b", bool, aggParams);

diff --git a/packages/muelu/src/MueCentral/MueLu_MasterList.cpp b/packages/muelu/src/MueCentral/MueLu_MasterList.cpp
@@ -164,6 +164,7 @@ namespace MueLu {
   "<Parameter name=\"coarse: type\" type=\"string\" value=\"KLU\"/>"
   "<ParameterList name=\"coarse: params\"/>"
   "<Parameter name=\"coarse: overlap\" type=\"int\" value=\"0\"/>"
+  "<Parameter name=\"aggregation: backend\" type=\"string\" value=\"default\"/>"
   "<Parameter name=\"aggregation: type\" type=\"string\" value=\"uncoupled\"/>"
   "<Parameter name=\"aggregation: mode\" type=\"string\" value=\"uncoupled\"/>"
   "<Parameter name=\"aggregation: ordering\" type=\"string\" value=\"natural\"/>"
@@ -597,6 +598,8 @@ namespace MueLu {
 
          ("coarse: overlap","coarse: overlap")
 
+         ("aggregation: backend","aggregation: backend")
+
          ("aggregation: type","aggregation: type")
 
          ("aggregation: mode","aggregation: mode")

diff --git a/packages/muelu/test/interface/default/Output/operator_solve_1_np1_tpetra.gold b/packages/muelu/test/interface/default/Output/operator_solve_1_np1_tpetra.gold
@@ -24,6 +24,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)
@@ -60,6 +61,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)
@@ -96,6 +98,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)

diff --git a/packages/muelu/test/interface/default/Output/operator_solve_1_np4_tpetra.gold b/packages/muelu/test/interface/default/Output/operator_solve_1_np4_tpetra.gold
@@ -24,6 +24,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)
@@ -60,6 +61,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)
@@ -96,6 +98,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)

diff --git a/packages/muelu/test/interface/default/Output/operator_solve_5_np1_tpetra.gold b/packages/muelu/test/interface/default/Output/operator_solve_5_np1_tpetra.gold
@@ -35,6 +35,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Build (MueLu::CoarseMapFactory)
 matrixmatrix: kernel params -> 
  [empty list]
@@ -69,6 +70,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)

diff --git a/packages/muelu/test/interface/default/Output/operator_solve_5_np4_tpetra.gold b/packages/muelu/test/interface/default/Output/operator_solve_5_np4_tpetra.gold
@@ -35,6 +35,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Build (MueLu::CoarseMapFactory)
 matrixmatrix: kernel params -> 
  [empty list]
@@ -69,6 +70,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)

diff --git a/packages/muelu/test/interface/default/Output/operator_solve_6_np1_tpetra.gold b/packages/muelu/test/interface/default/Output/operator_solve_6_np1_tpetra.gold
@@ -39,6 +39,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Build (MueLu::CoarseMapFactory)
 matrixmatrix: kernel params -> 
  [empty list]
@@ -74,6 +75,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)

diff --git a/packages/muelu/test/interface/default/Output/operator_solve_6_np4_tpetra.gold b/packages/muelu/test/interface/default/Output/operator_solve_6_np4_tpetra.gold
@@ -39,6 +39,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Build (MueLu::CoarseMapFactory)
 matrixmatrix: kernel params -> 
  [empty list]
@@ -74,6 +75,7 @@ BuildAggregatesNonKokkos (Phase 1 (main))
 BuildAggregatesNonKokkos (Phase 2a (secondary))
 BuildAggregatesNonKokkos (Phase 2b (expansion))
 BuildAggregatesNonKokkos (Phase 3 (cleanup))
+aggregation: deterministic = 0   [unused]
 Nullspace factory (MueLu::NullspaceFactory)
 Fine level nullspace = Nullspace
 Build (MueLu::CoarseMapFactory)