From d222c40dafdfe6c8f05b4d9d2f8b1e50923fd4ad Mon Sep 17 00:00:00 2001
From: Oscar Antepara <oantepara@lbl.gov>
Date: Thu, 15 May 2025 17:43:06 -0400
Subject: [PATCH] Adding AlbanyLaunchBounds to control Kokkos Launch Bounds
 configurations for GPU kernel optimization

---
 src/landIce/evaluators/LandIce_StokesFOResid.hpp | 5 +++--
 src/utility/Albany_KokkosUtils.hpp               | 6 ++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/landIce/evaluators/LandIce_StokesFOResid.hpp b/src/landIce/evaluators/LandIce_StokesFOResid.hpp
index 71c2283c82..b59a53866a 100644
--- a/src/landIce/evaluators/LandIce_StokesFOResid.hpp
+++ b/src/landIce/evaluators/LandIce_StokesFOResid.hpp
@@ -14,6 +14,7 @@
 
 #include "Albany_Layouts.hpp"
 #include "PHAL_Dimension.hpp"
+#include "Albany_KokkosUtils.hpp"
 
 namespace LandIce {
 /** \brief Finite Element Interpolation Evaluator
@@ -77,8 +78,8 @@ class StokesFOResid : public PHX::EvaluatorWithBaseImpl<Traits>,
   struct POISSON_2D_Tag{};
 
   typedef Kokkos::RangePolicy<ExecutionSpace,LandIce_3D_Tag> LandIce_3D_Policy;
-  typedef Kokkos::RangePolicy<ExecutionSpace,LandIce_3D_Opt_Tag<8>> LandIce_3D_Opt_Hex_Policy;
-  typedef Kokkos::RangePolicy<ExecutionSpace,LandIce_3D_Opt_Tag<6>> LandIce_3D_Opt_Wedge_Policy;
+  typedef Kokkos::RangePolicy<ExecutionSpace,KU::AlbanyLaunchBounds,LandIce_3D_Opt_Tag<8>> LandIce_3D_Opt_Hex_Policy;
+  typedef Kokkos::RangePolicy<ExecutionSpace,KU::AlbanyLaunchBounds,LandIce_3D_Opt_Tag<6>> LandIce_3D_Opt_Wedge_Policy;
   typedef Kokkos::RangePolicy<ExecutionSpace,POISSON_3D_Tag> POISSON_3D_Policy;
   typedef Kokkos::RangePolicy<ExecutionSpace,LandIce_2D_Tag> LandIce_2D_Policy;
   typedef Kokkos::RangePolicy<ExecutionSpace,LandIce_XZ_2D_Tag> LandIce_XZ_2D_Policy;
diff --git a/src/utility/Albany_KokkosUtils.hpp b/src/utility/Albany_KokkosUtils.hpp
index 68a9ae6c4d..86b858c43d 100644
--- a/src/utility/Albany_KokkosUtils.hpp
+++ b/src/utility/Albany_KokkosUtils.hpp
@@ -97,6 +97,12 @@ atomic_add(D dst, const V& val)
   KU::AtomicAddImpl<ExeSpace, D, V>::atomic_add(dst, val);
 }
 
+#ifdef KOKKOS_ENABLE_CUDA
+typedef Kokkos::LaunchBounds<> AlbanyLaunchBounds;
+#else
+typedef Kokkos::LaunchBounds<128,2> AlbanyLaunchBounds;
+#endif
+
 }
 
 #endif // ALBANY_KOKKOS_UTILS_HPP