Move duplicated one-liner to TensorUtils

morpheus/_lib/include/morpheus/utilities/tensor_util.hpp

@@ -19,10 +19,12 @@
 
 #include "morpheus/objects/tensor_object.hpp"
 
-#include <algorithm>  // IWYU pragma: keep
-#include <iosfwd>     // for ostream
-#include <string>     // for string
-#include <vector>     // for vector
+#include <algorithm>   // IWYU pragma: keep
+#include <functional>  // for multiplies
+#include <iosfwd>      // for ostream
+#include <numeric>     // for accumulate
+#include <string>      // for string
+#include <vector>      // for vector
 // <algorithm> is needed for copy, min_element & transform
 
 namespace morpheus {
@@ -106,6 +108,19 @@ struct TensorUtils
 
         return tensor_stride;
     }
+
+    /**
+     * @brief Compute the number of elements in a tensor based on the shape
+     *
+     * @tparam IndexT
+     * @param shape
+     * @return IndexT
+     */
+    template <typename IndexT>
+    static inline IndexT get_elem_count(const std::vector<IndexT>& shape)
+    {
+        return std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>());
+    }
 };
 
 /** @} */  // end of group

morpheus/_lib/src/objects/dev_mem_info.cpp

@@ -17,11 +17,11 @@
 
 #include "morpheus/objects/dev_mem_info.hpp"
 
+#include "morpheus/utilities/tensor_util.hpp"  // for get_elem_count
+
 #include <glog/logging.h>  // for DCHECK
 
 #include <cstdint>     // for uint8_t
-#include <functional>  // for multiplies
-#include <numeric>     // for accumulate
 #include <utility>     // for move
 
 namespace morpheus {
@@ -50,7 +50,7 @@ std::size_t DevMemInfo::bytes() const
 
 std::size_t DevMemInfo::count() const
 {
-    return std::accumulate(m_shape.begin(), m_shape.end(), 1, std::multiplies<>());
+    return TensorUtils::get_elem_count(m_shape);
 }
 
 std::size_t DevMemInfo::offset_bytes() const

morpheus/_lib/src/objects/rmm_tensor.cpp

@@ -21,7 +21,7 @@
 #include "morpheus/objects/dtype.hpp"
 #include "morpheus/objects/tensor_object.hpp"
 #include "morpheus/utilities/matx_util.hpp"
-#include "morpheus/utilities/tensor_util.hpp"  // for get_element_stride
+#include "morpheus/utilities/tensor_util.hpp"  // for get_elem_count & get_element_stride
 
 #include <cuda_runtime.h>            // for cudaMemcpy, cudaMemcpy2D, cudaMemcpyDeviceToDevice
 #include <glog/logging.h>            // for DCHECK_LT, COMPACT_GOOGLE_LOG_FATAL, DCHECK, DCHECK_EQ, LogMessageFatal
@@ -34,7 +34,7 @@
 #include <functional>  // for multiplies, plus, minus
 #include <iterator>    // for back_insert_iterator, back_inserter
 #include <memory>
-#include <numeric>  // for accumulate, transform_reduce
+#include <numeric>  // for transform_reduce
 #include <ostream>  // needed for logging
 #include <vector>
 
@@ -83,7 +83,7 @@ DType RMMTensor::dtype() const
 
 std::size_t RMMTensor::count() const
 {
-    return std::accumulate(m_shape.begin(), m_shape.end(), 1, std::multiplies<>());
+    return TensorUtils::get_elem_count(m_shape);
 }
 
 std::size_t RMMTensor::bytes() const

morpheus/_lib/src/stages/triton_inference.cpp

@@ -29,6 +29,7 @@
 #include "morpheus/utilities/matx_util.hpp"
 #include "morpheus/utilities/stage_util.hpp"
 #include "morpheus/utilities/string_util.hpp"  // for MORPHEUS_CONCAT_STR
+#include "morpheus/utilities/tensor_util.hpp"  // for get_elem_count
 
 #include <cuda_runtime.h>  // for cudaMemcpy, cudaMemcpy2D, cudaMemcpyDeviceToHost, cudaMemcpyHostToDevice
 #include <glog/logging.h>
@@ -43,9 +44,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <exception>
-#include <functional>  // for multiplies
 #include <memory>
-#include <numeric>  // for accumulate
 #include <sstream>
 #include <stdexcept>    // for runtime_error, out_of_range
 #include <type_traits>  // for declval
@@ -71,11 +70,6 @@ void InferenceClientStage__check_triton_errors(triton::client::Error status,
     }
 }
 
-template <typename IndexT>
-inline IndexT get_elem_count(const std::vector<IndexT>& shape)
-{
-    return std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>());
-}
 }  // namespace
 
 namespace morpheus {
@@ -121,7 +115,7 @@ InferenceClientStage::subscribe_fn_t InferenceClientStage::build_operator()
 
                     // First dimension will always end up being the number of rows in the dataframe
                     total_shape[0]  = static_cast<TensorIndex>(x->mess_count);
-                    auto elem_count = get_elem_count(total_shape);
+                    auto elem_count = TensorUtils::get_elem_count(total_shape);
 
                     // Create the output memory
                     auto output_buffer = std::make_shared<rmm::device_buffer>(

morpheus/_lib/src/utilities/tensor_util.cpp

@@ -24,9 +24,7 @@
 // prevent from moving this into the third-party section
 #include <experimental/iterator>  // for make_ostream_joiner
 // clang-format on
-#include <functional>   // for multiplies
 #include <iterator>     // for begin, end
-#include <numeric>      // for accumulate
 #include <ostream>      // for operator<<, ostream, stringstream
 #include <string>       // for char_traits, string
 #include <type_traits>  // for decay_t
@@ -62,7 +60,7 @@ void TensorUtils::set_contiguous_stride(const std::vector<TensorIndex>& shape, s
 bool TensorUtils::has_contiguous_stride(const std::vector<TensorIndex>& shape, const shape_type_t& stride)
 {
     DCHECK_EQ(shape.size(), stride.size());
-    auto count = std::accumulate(std::begin(shape), std::end(shape), 1, std::multiplies<>());
+    auto count = get_elem_count(shape);
     return (shape[0] * stride[0] == count);
 }