| /*! |
| * Copyright (c) 2016 by Contributors |
| * \file tvm/buffer.h |
| * \brief Symbolic n-dimensional array, to represent a memory buffer. |
| */ |
| #ifndef TVM_BUFFER_H_ |
| #define TVM_BUFFER_H_ |
| |
| #include <string> |
| |
| #include "base.h" |
| #include "expr.h" |
| #include "ir_operator.h" |
| #include "tvm/node/container.h" |
| |
| namespace tvm { |
| |
| // Internal node container Buffer |
| class BufferNode; |
| |
| /*! \brief memory access kind */ |
| enum class AccessMask : int { |
| kRead = 1, |
| kWrite = 2 |
| }; |
| |
| /*! |
| * \brief Buffer is a symbolic n-darray structure. |
| * It is a composition of primitive symbolic types, |
| * used to specify the memory layout of the Tensor used in program input. |
| */ |
| class Buffer : public NodeRef { |
| public: |
| Buffer() {} |
| explicit Buffer(NodePtr<Node> n) : NodeRef(n) {} |
| /*! |
| * \brief Return a new buffer that is equivalent with current one |
| * but always add stride field. |
| * \return The strided version of the buffer. |
| */ |
| TVM_DLL Buffer MakeStrideView() const; |
| /*! |
| * \brief Make a new symbolic buffer representing a slice of the buffer. |
| * \param begins The beginning position of each dimension. |
| * \param extents The extent of each dimension. |
| * \note This function will make target buffer as compact as possible. |
| * If stride is not needed in the slice, it won't be presented |
| * \return the result buffer. |
| */ |
| TVM_DLL Buffer MakeSlice(Array<Expr> begins, Array<Expr> extents) const; |
| /*! |
| * \brief Get access ptr to the entire buffer. |
| * \param access_mask The access mask |
| * \param ptr_type The type of the pointer. |
| * \param content_lanes The number of lanes for the (data) type. |
| * \param offset The offset of ptr. |
| */ |
| TVM_DLL Expr access_ptr(int access_mask, Type ptr_type = Handle(), |
| int content_lanes = 1, Expr offset = make_const(Int(32), 0)) const; |
| /*! |
| * \brief Create an Expr that does a vector load at begin index. |
| * \param begin The beginning index |
| * \param dtype The data type to be loaded. |
| */ |
| TVM_DLL Expr vload(Array<Expr> begin, Type dtype) const; |
| /*! |
| * \brief Create a Stmt that does a vector store at begin index. |
| * \param begin The beginning index |
| * \param value The value to be stored. |
| */ |
| TVM_DLL Stmt vstore(Array<Expr> begin, 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: |
| // Data fields. |
| /*! |
| * \brief The pointer to the head of the data |
| * \sa data_alignment The alignment of data in bytes. |
| */ |
| Var data; |
| /*! \brief data type in the content of the tensor */ |
| Type dtype; |
| /*! \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 The offset in terms of number of dtype elements (including lanes) */ |
| Expr elem_offset; |
| // Meta data |
| /*! \brief optional name of the buffer */ |
| std::string name; |
| /*! \brief storage scope of the buffer, if other than global */ |
| std::string scope; |
| /*! \brief Alignment requirement of data pointer in bytes. */ |
| int data_alignment; |
| /*! |
| * \brief Factor of elem_offset field, |
| * elem_offset is guaranteed to be multiple of offset_factor. |
| */ |
| int offset_factor; |
| /*! \brief constructor */ |
| BufferNode() {} |
| |
| void VisitAttrs(AttrVisitor* v) final { |
| v->Visit("data", &data); |
| v->Visit("dtype", &dtype); |
| v->Visit("shape", &shape); |
| v->Visit("strides", &strides); |
| v->Visit("elem_offset", &elem_offset); |
| v->Visit("name", &name); |
| v->Visit("scope", &scope); |
| v->Visit("data_alignment", &data_alignment); |
| v->Visit("offset_factor", &offset_factor); |
| } |
| |
| /*! \return preferred index type for this buffer node */ |
| Type DefaultIndexType() const { |
| return shape.size() != 0 ? shape[0].type() : Int(32); |
| } |
| |
| // User can specify data_alignment and offset_factor to be 0 |
| // A default value will be picked. |
| TVM_DLL static Buffer make(Var ptr, |
| Type dtype, |
| Array<Expr> shape, |
| Array<Expr> strides, |
| Expr elem_offset, |
| std::string name, |
| std::string scope, |
| int data_alignment, |
| int offset_factor); |
| |
| 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()); |
| } |
| |
| /*! |
| * \brief Construct a new buffer given shape, and dtype. |
| * \param shape The shape of the buffer, |
| * \param dtype The content data type. |
| * \param name The name of the buffer |
| * \return The created buffer. |
| * \sa BufferNode::make for complete constructor. |
| */ |
| TVM_DLL Buffer decl_buffer(Array<Expr> shape, |
| Type dtype = Float(32), |
| std::string name = "buffer"); |
| } // namespace tvm |
| #endif // TVM_BUFFER_H_ |