Program Listing for File tensor.hpp

            

            
#pragma once

#include "morpheus/objects/tensor_object.hpp"
#include "morpheus/utilities/type_util.hpp"

#include <rmm/device_buffer.hpp>

#include <cstddef>// for size_t
#include <cstdint>
#include <memory>
#include <string>
#include <vector>

namespace morpheus {
/****** Component public implementations *******************/
/****** Tensor****************************************/

#pragma GCC visibility push(default)
class Tensor
{
  public:
    Tensor(std::shared_ptr<rmm::device_buffer> buffer,
           std::string init_typestr,
           std::vector<int32_t> init_shape,
           std::vector<int32_t> init_strides,
           size_t init_offset = 0);

    std::vector<int32_t> shape;
    std::vector<int32_t> strides;
    std::string typestr;

    void *data() const;

    size_t bytes_count() const;

    std::vector<uint8_t> get_host_data() const;

    auto get_stream() const;

    static TensorObject create(std::shared_ptr<rmm::device_buffer> buffer,
                               DType dtype,
                               std::vector<TensorIndex> shape,
                               std::vector<TensorIndex> strides,
                               size_t offset = 0);

  private:
    size_t m_offset;
    std::shared_ptr<rmm::device_buffer> m_device_buffer;
};

#pragma GCC visibility pop// end of group
}  // namespace morpheus