Skip to content

Commit

Permalink
Move duplicated one-liner to TensorUtils (nv-morpheus#682)
Browse files Browse the repository at this point in the history
Simple one-liner for calculating the element count of a tensor based on the shape was copy/pasted into four different places in the code base (Likely I'm to blame for at least one of these).

This just replaces them all with a call to `TensorUtils::get_elem_count`

Authors:
  - David Gardner (https://github.com/dagardner-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: nv-morpheus#682
  • Loading branch information
dagardner-nv authored and jjacobelli committed Mar 7, 2023
1 parent 4b4c57a commit 56cb847
Show file tree
Hide file tree
Showing 5 changed files with 28 additions and 21 deletions.
23 changes: 19 additions & 4 deletions morpheus/_lib/include/morpheus/utilities/tensor_util.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions morpheus/_lib/src/objects/dev_mem_info.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions morpheus/_lib/src/objects/rmm_tensor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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>

Expand Down Expand Up @@ -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
Expand Down
10 changes: 2 additions & 8 deletions morpheus/_lib/src/stages/triton_inference.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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>(
Expand Down
4 changes: 1 addition & 3 deletions morpheus/_lib/src/utilities/tensor_util.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
}

Expand Down

0 comments on commit 56cb847

Please sign in to comment.