From e9fcd8429884a19262f828c728b6cbc4e2d71331 Mon Sep 17 00:00:00 2001
From: M Shehtab Zaman <szaman5@binghamton.edu>
Date: Mon, 17 Jun 2024 20:13:02 -0700
Subject: [PATCH] Updated instantiation code

---
 .../lbann/layers/misc/channelwise_softmax.hpp |  2 +-
 .../layers/misc/channelwise_softmax_impl.hpp  |  2 +-
 .../distconv/distconv_channelwise_softmax.hpp | 55 +++++++++----------
 .../misc/channelwise_softmax_kernels.cuh      |  1 +
 .../distconv/distconv_channelwise_softmax.cu  | 17 ++----
 5 files changed, 36 insertions(+), 41 deletions(-)

diff --git a/include/lbann/layers/misc/channelwise_softmax.hpp b/include/lbann/layers/misc/channelwise_softmax.hpp
index 9168bb60fac..086ace20b73 100644
--- a/include/lbann/layers/misc/channelwise_softmax.hpp
+++ b/include/lbann/layers/misc/channelwise_softmax.hpp
@@ -1,5 +1,5 @@
 ////////////////////////////////////////////////////////////////////////////////
-// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
+// Copyright (c) 2014-2024, Lawrence Livermore National Security, LLC.
 // Produced at the Lawrence Livermore National Laboratory.
 // Written by the LBANN Research Team (B. Van Essen, et al.) listed in
 // the CONTRIBUTORS file. <lbann-dev@llnl.gov>
diff --git a/include/lbann/layers/misc/channelwise_softmax_impl.hpp b/include/lbann/layers/misc/channelwise_softmax_impl.hpp
index fd1e7b1702d..6879eb294b5 100644
--- a/include/lbann/layers/misc/channelwise_softmax_impl.hpp
+++ b/include/lbann/layers/misc/channelwise_softmax_impl.hpp
@@ -1,5 +1,5 @@
 ////////////////////////////////////////////////////////////////////////////////
-// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
+// Copyright (c) 2014-2024, Lawrence Livermore National Security, LLC.
 // Produced at the Lawrence Livermore National Laboratory.
 // Written by the LBANN Research Team (B. Van Essen, et al.) listed in
 // the CONTRIBUTORS file. <lbann-dev@llnl.gov>
diff --git a/include/lbann/layers/misc/distconv/distconv_channelwise_softmax.hpp b/include/lbann/layers/misc/distconv/distconv_channelwise_softmax.hpp
index b039bf09738..4ca40035f84 100644
--- a/include/lbann/layers/misc/distconv/distconv_channelwise_softmax.hpp
+++ b/include/lbann/layers/misc/distconv/distconv_channelwise_softmax.hpp
@@ -1,5 +1,5 @@
 ////////////////////////////////////////////////////////////////////////////////
-// Copyright (c) 2014-2022, Lawrence Livermore National Security, LLC.
+// Copyright (c) 2014-2024, Lawrence Livermore National Security, LLC.
 // Produced at the Lawrence Livermore National Laboratory.
 // Written by the LBANN Research Team (B. Van Essen, et al.) listed in
 // the CONTRIBUTORS file. <lbann-dev@llnl.gov>
@@ -29,33 +29,32 @@
 #include "lbann/utils/distconv.hpp"
 
 #ifdef LBANN_HAS_DISTCONV
-namespace distconv{
-  template <typename Backend, typename DataType>
-  class ChannelwiseSoftmax{
-    using LocaleMPI = tensor::LocaleMPI;
-
-    public:
-      ChannelwiseSoftmax(Backend &backend):m_be(backend){};
-    
-    template <typename Allocator>
-    int forward(
-      const tensor::Tensor<DataType, LocaleMPI, Allocator> &input_0,
-      tensor::Tensor<DataType, LocaleMPI, Allocator> &output);
-
-    template <typename Allocator>
-    int backward(
-      const tensor::Tensor<DataType, LocaleMPI, Allocator> &input_0,
-      const tensor::Tensor<DataType, LocaleMPI, Allocator> &output_grad,
-      tensor::Tensor<DataType, LocaleMPI, Allocator> &input_grad_0);
-
-    protected:
-      Backend &m_be;
-
-  };
-
-  extern template class ChannelwiseSoftmax<::distconv::BackendDNNLib, float>;
-  extern template class ChannelwiseSoftmax<::distconv::BackendDNNLib, double>;
-}
+namespace distconv {
+template <typename Backend, typename DataType>
+class ChannelwiseSoftmax
+{
+  using LocaleMPI = tensor::LocaleMPI;
+
+public:
+  ChannelwiseSoftmax(Backend& backend) : m_be(backend){};
+
+  template <typename Allocator>
+  int forward(const tensor::Tensor<DataType, LocaleMPI, Allocator>& input_0,
+              tensor::Tensor<DataType, LocaleMPI, Allocator>& output);
+
+  template <typename Allocator>
+  int backward(
+    const tensor::Tensor<DataType, LocaleMPI, Allocator>& input_0,
+    const tensor::Tensor<DataType, LocaleMPI, Allocator>& output_grad,
+    tensor::Tensor<DataType, LocaleMPI, Allocator>& input_grad_0);
+
+protected:
+  Backend& m_be;
+};
+
+extern template class ChannelwiseSoftmax<::distconv::BackendDNNLib, float>;
+extern template class ChannelwiseSoftmax<::distconv::BackendDNNLib, double>;
+} // namespace distconv
 
 #endif // LBANN_HAS_DISTCONV
 #endif // LBANN_LAYERS_MISC_DISTCONV_CHANNELWISE_SOFTMAX
\ No newline at end of file
diff --git a/src/layers/misc/channelwise_softmax_kernels.cuh b/src/layers/misc/channelwise_softmax_kernels.cuh
index 95a9c2423d0..9c1ae86bc1b 100644
--- a/src/layers/misc/channelwise_softmax_kernels.cuh
+++ b/src/layers/misc/channelwise_softmax_kernels.cuh
@@ -25,6 +25,7 @@
 ////////////////////////////////////////////////////////////////////////////////
 #ifndef LBANN_LAYERS_MISC_CHANNELWISE_SOFTMAX_KERNELS
 #define LBANN_LAYERS_MISC_CHANNELWISE_SOFTMAX_KERNELS
+#include "lbann/utils/gpu/helpers.hpp"
 namespace lbann{
 namespace{
 using Size3 = gpu_lib::array<size_t,3>;
diff --git a/src/layers/misc/distconv/distconv_channelwise_softmax.cu b/src/layers/misc/distconv/distconv_channelwise_softmax.cu
index b99c6cd00a2..0d3db2bad8d 100644
--- a/src/layers/misc/distconv/distconv_channelwise_softmax.cu
+++ b/src/layers/misc/distconv/distconv_channelwise_softmax.cu
@@ -123,26 +123,21 @@ int ChannelwiseSoftmax<Backend, DataType>::backward(
 // Explicit template instantiation
 // =========================================================
 
-#define ETI(T, Backend)                                                        \
-  template class ChannelwiseSoftmax<Backend, T>;                               \
-  template int ChannelwiseSoftmax<Backend, T>::forward<tensor::CUDAAllocator>( \
+#define PROTO(T)                                                               \
+  template class ChannelwiseSoftmax<BackendDNNLib, T>;                         \
+  template int                                                                 \
+  ChannelwiseSoftmax<BackendDNNLib, T>::forward<tensor::CUDAAllocator>(        \
     const tensor::Tensor<T, tensor::LocaleMPI, tensor::CUDAAllocator>&         \
       input_0,                                                                 \
     tensor::Tensor<T, tensor::LocaleMPI, tensor::CUDAAllocator>& output_0);    \
   template int                                                                 \
-  ChannelwiseSoftmax<Backend, T>::backward<tensor::CUDAAllocator>(             \
+  ChannelwiseSoftmax<BackendDNNLib, T>::backward<tensor::CUDAAllocator>(       \
     const tensor::Tensor<T, tensor::LocaleMPI, tensor::CUDAAllocator>&         \
       input_0,                                                                 \
     const tensor::Tensor<T, tensor::LocaleMPI, tensor::CUDAAllocator>&         \
       input_1,                                                                 \
     tensor::Tensor<T, tensor::LocaleMPI, tensor::CUDAAllocator>& output_grad);
 
-/// @todo: fp16
-ETI(float, BackendDNNLib)
-#ifdef LBANN_HAS_DOUBLE
-ETI(double, BackendDNNLib)
-#endif // LBANN_HAS_DOUBLE
-
-#undef ETI
+#include "lbann/macros/instantiate.hpp"
 } // namespace distconv
 #endif // LBANN_HAS_DISTCONV