/*! * Copyright (c) 2016 by Contributors * \file buffer.h * \brief Symbolic n-dimensional array, to represent a memory buffer. */ #ifndef TVM_BUFFER_H_ #define TVM_BUFFER_H_ #include <tvm/container.h> #include <string> #include "./base.h" #include "./expr.h" namespace tvm { // Internal node container Buffer class BufferNode; /*! * \brief Buffer is a symbolic n-darray structure. * It is a composition of primitive symbolic types, * used to specify input/output strcuture of the program. */ class Buffer : public NodeRef { public: Buffer() {} explicit Buffer(std::shared_ptr<Node> n) : NodeRef(n) {} /*! * \brief construct a new buffer based on shape and strides. */ explicit Buffer(Array<Expr> shape, Type dtype = Float(32), std::string name = "buffer"); /*! * \brief Generate a load expression loading the index location of buffer. * \param index The index to the buffer. * \return The load expression. */ Expr MakeLoad(Array<Expr> index) const; /*! * \brief Generate a store statement. * \param index The index to the buffer. * \param value The value to be stored. * \return The load expression. */ Stmt MakeStore(Array<Expr> index, Expr value) const; /*! * \brief access the internal node container * \return the pointer to the internal node container */ inline const BufferNode* operator->() const; /*! \brief specify container node */ using ContainerType = BufferNode; }; /*! \brief Node to represent a buffer */ class BufferNode : public Node { public: /*! \brief optional name of the buffer */ std::string name; /*! \brief The pointer to the head of the data */ Var data; /*! \brief The shape of the buffer */ Array<Expr> shape; /*! * \brief The strides of each dimension * This can be an empty array, indicating array is contiguous */ Array<Expr> strides; /*! \brief data type in the content of the tensor */ Type dtype; /*! * \brief The offset in bytes to the beginning pointer to data * Can be undefined, indicating this must be zero. */ Expr byte_offset; /*! \brief Alignment bytes size of byte_offset */ int offset_alignment; /*! \brief constructor */ BufferNode() {} void VisitAttrs(AttrVisitor* v) final { v->Visit("name", &name); v->Visit("data", &data); v->Visit("shape", &shape); v->Visit("strides", &strides); v->Visit("dtype", &dtype); v->Visit("byte_offset", &byte_offset); v->Visit("offset_alignment", &offset_alignment); } static Buffer make(std::string name, Var ptr, Array<Expr> shape, Array<Expr> strides, Type dtype, Expr byte_offset, int offset_alignment); static constexpr const char* _type_key = "Buffer"; TVM_DECLARE_NODE_TYPE_INFO(BufferNode, Node); }; inline const BufferNode* Buffer::operator->() const { return static_cast<const BufferNode*>(node_.get()); } } // namespace tvm #endif // TVM_BUFFER_H_