tensor_buffer.h

Go to the documentation of this file.

#ifndef ATEN_CORE_TENSOR_BUFFER_H_
#define ATEN_CORE_TENSOR_BUFFER_H_
 
#include <base/core/refcount.h>
#include <base/core/allocator.h>
#include <ATen/core/tensor_types.h>
 
namespace container {
 
class TensorBuffer : public base::core::counted_base {
   public:
    TensorBuffer(base::core::Allocator* alloc, void* data_ptr);
 
    TensorBuffer(base::core::Allocator* alloc, size_t size);
 
    explicit TensorBuffer(void* data_ptr);
 
    TensorBuffer(TensorBuffer&& other) noexcept;
 
    ~TensorBuffer();
 
    void* data() const;
 
    void resize(size_t size);
 
    size_t GetAllocatedBytes() const;
 
    TensorBuffer* root_buffer();
 
    base::core::Allocator* allocator() const;
 
    template <typename T>
    T* base() const {
        return reinterpret_cast<T*>(data());
    }
 
    virtual bool OwnsMemory() const;
 
    DeviceType GetDeviceType() const;
 
    TensorBuffer& operator=(const TensorBuffer& other);
 
    TensorBuffer& operator=(TensorBuffer&& other) noexcept;
 
 
  private:
    base::core::Allocator* alloc_ = nullptr; 
    void *data_ = nullptr;       
    bool owns_memory_ = false;    
    int64_t allocated_bytes_ = 0; 
};
 
}  // namespace container
 
#endif  // ATEN_CORE_TENSOR_BUFFER_H_