Program Listing for File matx_util.hpp

            

            
#pragma once

#include "morpheus/objects/dev_mem_info.hpp"
#include "morpheus/objects/rmm_tensor.hpp"
#include "morpheus/objects/tensor_object.hpp"
#include "morpheus/utilities/type_util_detail.hpp"

#include <cstddef>
#include <memory>
#include <vector>

namespace morpheus {

struct MatxUtil
{
    static std::shared_ptr<rmm::device_buffer> cast(const DevMemInfo& input, TypeId output_type);

    static std::shared_ptr<rmm::device_buffer> create_seg_ids(size_t row_count, size_t fea_len, TypeId output_type);

    static std::shared_ptr<rmm::device_buffer> logits(const DevMemInfo& input);

    static std::shared_ptr<rmm::device_buffer> transpose(const DevMemInfo& input);

    static std::shared_ptr<rmm::device_buffer> threshold(const DevMemInfo& input, double thresh_val, bool by_row);

    static std::shared_ptr<rmm::device_buffer> reduce_max(const DevMemInfo& input,
                                                          const std::vector<int32_t>& seq_ids,
                                                          size_t seq_id_offset,
                                                          const std::vector<int64_t>& output_shape);
};  // end of group
}  // namespace morpheus