Updated instantiation code

LLNL · Jun 18, 2024 · e9fcd84 · e9fcd84
1 parent 8eef5a0
commit e9fcd84
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 41 deletions.
diff --git 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. <[email protected]>

diff --git 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. <[email protected]>

diff --git 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. <[email protected]>
@@ -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
diff --git 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
@@ -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