updating parallel_for in tests

MeasureTransport · Feb 19, 2024 · 42bb09d · 42bb09d
1 parent c20713a
commit 42bb09d
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 14 deletions.
diff --git a/MParT/Utilities/LinearAlgebra.h b/MParT/Utilities/LinearAlgebra.h
@@ -26,7 +26,7 @@ template<typename ViewType1, typename ViewType2>
 void AddInPlace(ViewType1 x, ViewType2 y) {
     constexpr size_t rank = GetViewRank<ViewType1>::Rank;
     if constexpr(rank == 1) {
-//	Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy(0,x.extent(0));
+	Kokkos::RangePolicy<typename ViewType1::execution_space> policy(0,x.extent(0));
         Kokkos::parallel_for(x.extent(0), KOKKOS_LAMBDA(const int i){x(i) += y(i);});
     } else if (rank == 2) {
         Kokkos::MDRangePolicy<Kokkos::Rank<2>,typename ViewType1::execution_space> policy({0, 0}, {x.extent(0), x.extent(1)});

diff --git a/tests/Test_MapObjective.cpp b/tests/Test_MapObjective.cpp
@@ -46,7 +46,8 @@ TEST_CASE( "Test KLMapObjective", "[KLMapObjective]") {
     const double coeff_def = 1.;
 
     auto map = MapFactory::CreateTriangular<Kokkos::HostSpace>(dim, dim, 2);
-    Kokkos::parallel_for("Fill coeffs", map->numCoeffs, KOKKOS_LAMBDA(const unsigned int i){
+    Kokkos::RangePolicy<typename MemoryToExecution<Kokkos::HostSpace>::Space> policy(0, map->numCoeffs);
+    Kokkos::parallel_for("Fill coeffs", policy, KOKKOS_LAMBDA(const unsigned int i){
         map->Coeffs()(i) = coeff_def;
     });
 

diff --git a/tests/Test_MonotoneComponent.cpp b/tests/Test_MonotoneComponent.cpp
@@ -905,7 +905,8 @@ TEST_CASE( "MonotoneIntegrand1d on device", "[MonotoneIntegrandDevice]") {
     SECTION("Integrand Only") {
         Kokkos::View<double*, DeviceSpace> dres("result", 1);
 
-        Kokkos::parallel_for(1, KOKKOS_LAMBDA(const int i){
+    	Kokkos::RangePolicy<typename MemoryToExecution<DeviceSpace>::Space> policy(0,1);
+        Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const int i){
             MonotoneIntegrand<decltype(expansion), Exp, decltype(dpt), decltype(dcoeffs), DeviceSpace> integrand(dcache.data(), expansion, dpt, dcoeffs, DerivativeFlags::None, 0.0);
             integrand(0.5, &dres(0));
         });
@@ -918,7 +919,8 @@ TEST_CASE( "MonotoneIntegrand1d on device", "[MonotoneIntegrandDevice]") {
     SECTION("Integrand Derivative") {
         Kokkos::View<double*, DeviceSpace> dres("result", 2);
 
-        Kokkos::parallel_for(1, KOKKOS_LAMBDA(const int i){
+    	Kokkos::RangePolicy<typename MemoryToExecution<DeviceSpace>::Space> policy(0,1);
+        Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const int i){
             MonotoneIntegrand<decltype(expansion), Exp, decltype(dpt), decltype(dcoeffs), DeviceSpace> integrand(dcache.data(), expansion, dpt, dcoeffs, DerivativeFlags::Diagonal, 0.0);
             integrand(0.5, dres.data());
         });
@@ -934,7 +936,8 @@ TEST_CASE( "MonotoneIntegrand1d on device", "[MonotoneIntegrandDevice]") {
         Kokkos::View<double*, DeviceSpace> dres_fd("result_fd", hset.Size());
         Kokkos::View<double*, DeviceSpace> testVal("integrand", 1+hset.Size());
 
-        Kokkos::parallel_for(1, KOKKOS_LAMBDA(const int i){
+    	Kokkos::RangePolicy<typename MemoryToExecution<DeviceSpace>::Space> policy(0,1);
+        Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const int i){
 
             MonotoneIntegrand<decltype(expansion), Exp, decltype(dpt), decltype(dcoeffs), DeviceSpace> integrand(dcache.data(), expansion, dpt, dcoeffs, DerivativeFlags::Parameters, 0.0);
             MonotoneIntegrand<decltype(expansion), Exp, decltype(dpt), decltype(dcoeffs), DeviceSpace> integrand2(dcache.data(), expansion, dpt, dcoeffs, DerivativeFlags::None, 0.0);
@@ -968,8 +971,9 @@ TEST_CASE( "MonotoneIntegrand1d on device", "[MonotoneIntegrandDevice]") {
         Kokkos::View<double*, DeviceSpace> dres_fd("result_fd", hset.Size());
         Kokkos::View<double*, DeviceSpace> testVal("integrand", 1+hset.Size());
         Kokkos::View<double*, DeviceSpace> workspace("workspace", hset.Size());
-
-        Kokkos::parallel_for(1, KOKKOS_LAMBDA(const int i){
+
+	Kokkos::RangePolicy<typename MemoryToExecution<DeviceSpace>::Space> policy(0,1);
+        Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const int i){
 
             MonotoneIntegrand<decltype(expansion), Exp, decltype(dpt), decltype(dcoeffs), DeviceSpace> integrand(dcache.data(), expansion, dpt, dcoeffs, DerivativeFlags::Mixed, 0.0, workspace);
             MonotoneIntegrand<decltype(expansion), Exp, decltype(dpt), decltype(dcoeffs), DeviceSpace> integrand2(dcache.data(), expansion, dpt, dcoeffs, DerivativeFlags::Diagonal, 0.0);
@@ -1036,7 +1040,8 @@ TEST_CASE( "Testing MonotoneComponent::EvaluateSingle on Device", "[MonotoneComp
 
     Kokkos::View<double*, DeviceSpace> dres("Device Evaluation", 1);
     // Run the fill cache funciton, using a parallel_for loop to ensure it's run on the device
-    Kokkos::parallel_for(1, KOKKOS_LAMBDA(const int i){
+    Kokkos::RangePolicy<typename MemoryToExecution<DeviceSpace>::Space> policy(0,1);
+    Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const int i){
         dexpansion.FillCache1(dcache.data(), dpt, DerivativeFlags::None);
         dres(0) = MonotoneComponent<decltype(dexpansion),Exp, decltype(quad), DeviceSpace>::EvaluateSingle(dcache.data(), workspace.data(), dpt, dpt(dim-1), dcoeffs, quad, dexpansion);
     });
@@ -1139,4 +1144,4 @@ TEST_CASE( "Testing 1d monotone component evaluation on device", "[MonotoneCompo
 
 }
 
-#endif
+#endif
diff --git a/tests/Test_MultivariateExpansionWorker.cpp b/tests/Test_MultivariateExpansionWorker.cpp
@@ -251,7 +251,8 @@ TEST_CASE( "Testing multivariate expansion on device", "[MultivariateExpansionWo
     hexpansion.FillCache1(hcache.data(), hpt, DerivativeFlags::None);
 
     // Run the fill cache funciton, using a parallel_for loop to ensure it's run on the device
-    Kokkos::parallel_for(1, KOKKOS_LAMBDA(const int i){
+    Kokkos::RangePolicy<typename MemoryToExecution<DeviceSpace>::Space> policy(0,1);
+    Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const int i){
         dexpansion.FillCache1(dcache.data(), dpt, DerivativeFlags::None);
         dexpansion.FillCache2(dcache.data(), dpt, 0.5 * dpt(dim-1), DerivativeFlags::None);
     });
@@ -267,4 +268,4 @@ TEST_CASE( "Testing multivariate expansion on device", "[MultivariateExpansionWo
     }
 }
 
-#endif
+#endif
diff --git a/tests/Test_OrthogonalPolynomials.cpp b/tests/Test_OrthogonalPolynomials.cpp
@@ -251,9 +251,9 @@ TEST_CASE( "Device Hermite polynomial evaluation", "[PhysicistHermiteDevice]" )
     auto xs_device = ToDevice<Kokkos::DefaultExecutionSpace::memory_space>(xs_host);
 
     Kokkos::View<double*,Kokkos::DefaultExecutionSpace::memory_space> ys_device("evals", xs.size());
-
+    Kokkos::RangePolicy<typename Kokkos::DefaultExecutionSpace::execution_space> policy(0,xs.size());
     for(unsigned int p=0; p<10; ++p){
-        Kokkos::parallel_for(xs.size(), KOKKOS_LAMBDA(const size_t ind){
+        Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const size_t ind){
             ys_device(ind) = poly.Evaluate(p, xs_device(ind));
         });
 
@@ -267,4 +267,4 @@ TEST_CASE( "Device Hermite polynomial evaluation", "[PhysicistHermiteDevice]" )
 }
 
 
-#endif 
+#endif