wehub-resource-sync
375 lines
14 KiB
C++
375 lines
14 KiB
C++
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you under the Apache License, Version 2.0 (the
|
|
* "License"); you may not use this file except in compliance
|
|
* with the License. You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing,
|
|
* software distributed under the License is distributed on an
|
|
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
* KIND, either express or implied. See the License for the
|
|
* specific language governing permissions and limitations
|
|
* under the License.
|
|
*/
|
|
|
|
/*!
|
|
* \file tvm/tirx/buffer.h
|
|
* \brief Symbolic n-dimensional array, to represent a memory buffer.
|
|
*/
|
|
#ifndef TVM_TIRX_BUFFER_H_
|
|
#define TVM_TIRX_BUFFER_H_
|
|
|
|
#include <tvm/ffi/container/array.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
#include <tvm/ffi/string.h>
|
|
#include <tvm/ir/expr.h>
|
|
#include <tvm/tirx/layout.h>
|
|
#include <tvm/tirx/var.h>
|
|
|
|
#include <string>
|
|
|
|
namespace tvm {
|
|
namespace tirx {
|
|
|
|
#ifndef TVM_INDEX_DEFAULT_I64
|
|
#define TVM_INDEX_DEFAULT_I64 1
|
|
#endif
|
|
/*! \brief if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32 */
|
|
inline PrimType DefaultIndexPrimType() {
|
|
#if TVM_INDEX_DEFAULT_I64
|
|
static const PrimType default_index_ty = PrimType::Int(64);
|
|
#else
|
|
static const PrimType default_index_ty = PrimType::Int(32);
|
|
#endif
|
|
return default_index_ty;
|
|
}
|
|
|
|
inline DLDataType DefaultIndexType() {
|
|
#if TVM_INDEX_DEFAULT_I64
|
|
return DLDataType{kDLInt, 64, 1};
|
|
#else
|
|
return DLDataType{kDLInt, 32, 1};
|
|
#endif
|
|
}
|
|
|
|
// forward declare Stmt
|
|
class Stmt;
|
|
|
|
/*! \brief buffer type */
|
|
enum BufferType : int {
|
|
kDefault = 1,
|
|
// Maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 1.
|
|
kAutoBroadcast = 2,
|
|
};
|
|
|
|
/*! \brief Node to represent a buffer */
|
|
class BufferNode : public ffi::Object {
|
|
public:
|
|
// Data fields.
|
|
/*!
|
|
* \brief The pointer to the head of the data
|
|
* \sa data_alignment The alignment of data in bytes.
|
|
*/
|
|
Var data;
|
|
/*! \brief dtype in the content of the tensor */
|
|
PrimType dtype = PrimType::Void();
|
|
/*! \brief The type of the buffer prior to flattening
|
|
*
|
|
* This contains the shape as it is accessed by
|
|
* BufferLoad/BufferStore nodes, and used by the low-level code
|
|
* generators.
|
|
*/
|
|
ffi::Array<PrimExpr> shape;
|
|
/*!
|
|
* \brief Separators between input axes when generating flattened output axes
|
|
*
|
|
* For buffers representing flat 1-d memory (e.g. any buffer in
|
|
* RAM), this should be an empty array. For buffers representing
|
|
* non-flat memory, each entry in axis_separators should be the
|
|
* first input axis that is part of a new flattened axis.
|
|
*/
|
|
ffi::Array<IntImm> axis_separators;
|
|
/*!
|
|
* \brief The strides of each dimension
|
|
* This can be an empty array, indicating array is contiguous
|
|
*/
|
|
ffi::Array<PrimExpr> strides;
|
|
/*! \brief The offset in terms of number of dtype elements (including lanes) */
|
|
PrimExpr elem_offset;
|
|
// Meta data
|
|
/*! \brief optional name of the buffer */
|
|
ffi::String name;
|
|
/*! \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 buffer type */
|
|
BufferType buffer_type;
|
|
/*!
|
|
* \brief Span that points to the original source code.
|
|
* Reserved debug information.
|
|
*/
|
|
mutable Span span;
|
|
|
|
/*! \brief The layout of the buffer */
|
|
ffi::Optional<Layout> layout;
|
|
|
|
/*! \brief The allocated address of the buffer.
|
|
* The address might be multi-dimensional based on its scope.
|
|
* For example, trn.psum takes 2D address, representing (bank, offset).
|
|
*/
|
|
ffi::Array<PrimExpr> allocated_addr;
|
|
|
|
/*! \brief constructor */
|
|
BufferNode() {}
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<BufferNode>()
|
|
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
|
|
.def_ro("data", &BufferNode::data, refl::AttachFieldFlag::SEqHashDefRecursive())
|
|
.def_ro("dtype", &BufferNode::dtype)
|
|
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
|
|
.def_ro("shape", &BufferNode::shape, refl::AttachFieldFlag::SEqHashDefRecursive())
|
|
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
|
|
.def_ro("strides", &BufferNode::strides, refl::AttachFieldFlag::SEqHashDefRecursive())
|
|
.def_ro("axis_separators", &BufferNode::axis_separators,
|
|
refl::AttachFieldFlag::SEqHashDefRecursive())
|
|
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
|
|
.def_ro("elem_offset", &BufferNode::elem_offset,
|
|
refl::AttachFieldFlag::SEqHashDefRecursive())
|
|
.def_ro("name", &BufferNode::name, refl::AttachFieldFlag::SEqHashIgnore())
|
|
.def_ro("data_alignment", &BufferNode::data_alignment)
|
|
.def_ro("offset_factor", &BufferNode::offset_factor)
|
|
.def_ro("buffer_type", &BufferNode::buffer_type)
|
|
.def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore())
|
|
.def_ro("layout", &BufferNode::layout)
|
|
.def_ro("allocated_addr", &BufferNode::allocated_addr);
|
|
}
|
|
|
|
/*! \return preferred index type for this buffer node */
|
|
DLDataType DefaultIndexType() const {
|
|
return shape.size() != 0 ? shape[0].ty()->dtype : tvm::tirx::DefaultIndexType();
|
|
}
|
|
|
|
/*! \return primitive element type for compiler-side uses. */
|
|
PrimType ElementType() const { return dtype; }
|
|
|
|
/*! \brief Determine the offset in the buffer of the given index.
|
|
*
|
|
* Returns the buffer offset, in number of elements of type dtype,
|
|
* without adjusting for number of lanes. (e.g. The number of
|
|
* float16x4 elements in a buffer of type float16x4.)
|
|
*
|
|
* \param index The index to be accessed.
|
|
* \param inner Ignore the elem_offset, return inner offset only
|
|
*/
|
|
ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index, bool inner = false) const;
|
|
|
|
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
|
|
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Buffer", BufferNode, ffi::Object);
|
|
};
|
|
|
|
/*!
|
|
* \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 ffi::ObjectRef {
|
|
public:
|
|
// User can specify data_alignment and offset_factor to be 0
|
|
// A default value will be picked.
|
|
TVM_DLL Buffer(Var data, PrimType dtype, ffi::Array<PrimExpr> shape, ffi::Array<PrimExpr> strides,
|
|
PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor,
|
|
BufferType buffer_type, ffi::Array<IntImm> axis_separators = {},
|
|
Span span = Span(), ffi::Optional<Layout> layout = std::nullopt,
|
|
ffi::Array<PrimExpr> allocated_addr = {});
|
|
|
|
/*!
|
|
* \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(ffi::Array<PrimExpr> begins, ffi::Array<PrimExpr> 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.
|
|
* \param input_extent The extent of ptr.
|
|
*/
|
|
TVM_DLL Expr access_ptr(int access_mask, PointerType ptr_type = PointerType::VoidPointerTy(),
|
|
int content_lanes = 1, PrimExpr offset = IntImm::Int32(0),
|
|
ffi::Optional<PrimExpr> input_extent = std::nullopt) 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.
|
|
* \param predicate A vector mask of boolean values indicating which lanes of a vector are to be
|
|
* loaded. The number lanes of the mask must be equal to the number of lanes in being loaded.
|
|
*/
|
|
TVM_DLL PrimExpr vload(ffi::Array<PrimExpr> begin, PrimType dtype,
|
|
ffi::Optional<PrimExpr> predicate = std::nullopt) 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.
|
|
* \param predicate A vector mask of boolean values indicating which lanes of a vector are to be
|
|
* stored. The number lanes of the mask must be equal to the number of lanes in value.
|
|
*/
|
|
TVM_DLL Stmt vstore(ffi::Array<PrimExpr> begin, PrimExpr value,
|
|
ffi::Optional<PrimExpr> predicate = std::nullopt) const;
|
|
|
|
/*!
|
|
* \brief Get a flattened version of the buffer
|
|
*/
|
|
Buffer GetFlattenedBuffer() const;
|
|
|
|
/*! \brief Determine the offset in the buffer of the given index.
|
|
*
|
|
* Returns the buffer offset, in number of elements of type dtype,
|
|
* without adjusting for number of lanes. (e.g. The number of
|
|
* float16x4 elements in a buffer of type float16x4.)
|
|
*/
|
|
ffi::Array<PrimExpr> OffsetOf(ffi::Array<PrimExpr> index) const;
|
|
|
|
/*!
|
|
* \brief Get the buffer_offset op for the given index.
|
|
* \param index The index to be accessed.
|
|
* \return The buffer_offset op.
|
|
*/
|
|
PrimExpr OffsetOf_p(const ffi::Array<PrimExpr>& indices) const;
|
|
|
|
/*!
|
|
* \brief Return the storage scope associated with this buffer.
|
|
*/
|
|
TVM_DLL ffi::String scope() const;
|
|
|
|
/*!
|
|
* \brief Return a new buffer with the allocated address.
|
|
*/
|
|
TVM_DLL Buffer with_allocated_addr(ffi::Array<PrimExpr> allocated_addr) const;
|
|
|
|
/*!
|
|
* \brief Return true if the buffer is a scalar.
|
|
* \param alloc_or_decl Whether to consider alloc_scalar and decl_scalar as scalar. True for
|
|
* alloc_scalar, False for decl_scalar.
|
|
*/
|
|
TVM_DLL bool IsScalar(bool alloc_or_decl = true) const;
|
|
|
|
/*!
|
|
* \brief Return a new buffer with the dtype.
|
|
*/
|
|
TVM_DLL Buffer with_dtype(PrimType dtype) const;
|
|
|
|
/*! \return primitive element type for compiler-side uses. */
|
|
PrimType ElementType() const { return (*this)->ElementType(); }
|
|
|
|
/*!
|
|
* \brief Return a new buffer with the data.
|
|
*/
|
|
TVM_DLL Buffer with_data(Var data) const;
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Buffer, ffi::ObjectRef, BufferNode);
|
|
TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferNode);
|
|
};
|
|
|
|
/*!
|
|
* \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
|
|
* \param storage_scope The storage scope associated with this buffer
|
|
* \param axis_separators Divisions defining the groups of axes that will be flattened together.
|
|
* \param span The location of this object in the source code.
|
|
* \return The created buffer.
|
|
* \sa Buffer for complete constructor.
|
|
*/
|
|
TVM_DLL Buffer decl_buffer(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
|
|
ffi::String name = "buffer", ffi::String storage_scope = "",
|
|
ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
|
|
Span span = Span());
|
|
|
|
/*!
|
|
* \brief Base node for data producers.
|
|
*
|
|
* A DataProducer stores necessary information(e.g. a tensor expression) to produce
|
|
* a multi-dimensional array. The stored information is opaque to the TIR.
|
|
* DataProducer can appear in high-level DSLs that are built on top of the TIR.
|
|
*
|
|
* A valid TIR PrimFunc should not contain any DataProducer, high level DSLs should lower
|
|
* all DataProducers to Buffers before TIR transformations.
|
|
*
|
|
* \sa tvm::te::Tensor
|
|
*/
|
|
class DataProducerNode : public PrimExprConvertibleNode {
|
|
public:
|
|
/*! \brief destructor. */
|
|
virtual ~DataProducerNode() {}
|
|
/*!
|
|
* \brief Get the shape of the result.
|
|
* \return The shape.
|
|
*/
|
|
virtual ffi::Array<PrimExpr> GetShape() const = 0;
|
|
/*!
|
|
* \brief Get the raw element dtype of the result.
|
|
* \return The raw dtype.
|
|
*/
|
|
virtual PrimType GetDataType() const = 0;
|
|
/*!
|
|
* \brief Get the name hint of the data producer.
|
|
* \return The data type.
|
|
*/
|
|
virtual ffi::String GetNameHint() const = 0;
|
|
TVM_FFI_DECLARE_OBJECT_INFO("tirx.DataProducer", DataProducerNode, PrimExprConvertibleNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to DataProducerNode.
|
|
* \sa DataProducerNode
|
|
*/
|
|
class DataProducer : public PrimExprConvertible {
|
|
public:
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataProducer, PrimExprConvertible, DataProducerNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Creates TIR Buffer for provided parameters
|
|
* \param shape shape of the buffer
|
|
* \param dtype data type
|
|
* \param name buffer name
|
|
* \param data_alignment alignment requirement of data pointer in bytes
|
|
* \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
|
|
* multiple of offset_factor
|
|
User can specify data_alignment and offset_factor to be 0
|
|
* A default value will be picked.
|
|
* \param compact If the statement has already bound to a compact buffer.
|
|
* \param memory_scope memory scope of the buffer
|
|
*/
|
|
TVM_DLL tirx::Buffer BufferWithOffsetAlignment(ffi::Array<PrimExpr> shape, PrimType dtype,
|
|
std::string name, int data_alignment,
|
|
int offset_factor, bool compact,
|
|
std::string memory_scope = "");
|
|
} // namespace tirx
|
|
} // namespace tvm
|
|
#endif // TVM_TIR_BUFFER_H_
|