tvm
buffer.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
24 #ifndef TVM_TIR_BUFFER_H_
25 #define TVM_TIR_BUFFER_H_
26 
27 #include <tvm/ffi/container/array.h>
28 #include <tvm/ffi/reflection/registry.h>
29 #include <tvm/ffi/string.h>
30 #include <tvm/ir/expr.h>
32 #include <tvm/tir/var.h>
33 
34 #include <string>
35 
36 namespace tvm {
37 namespace tir {
38 
39 #ifndef TVM_INDEX_DEFAULT_I64
40 #define TVM_INDEX_DEFAULT_I64 1
41 #endif
44 #if TVM_INDEX_DEFAULT_I64
45  return DataType::Int(64);
46 #else
47  return DataType::Int(32);
48 #endif
49 }
50 
51 // forward declare Stmt
52 class Stmt;
53 
55 enum BufferType : int {
56  kDefault = 1,
57  // Maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 1.
59 };
60 
62 class BufferNode : public Object {
63  public:
64  // Data fields.
78  ffi::Array<PrimExpr> shape;
87  ffi::Array<IntImm> axis_separators;
92  ffi::Array<PrimExpr> strides;
95  // Meta data
97  ffi::String name;
111  mutable Span span;
114 
115  static void RegisterReflection() {
116  namespace refl = tvm::ffi::reflection;
117  refl::ObjectDef<BufferNode>()
118  .def_ro("data", &BufferNode::data, refl::AttachFieldFlag::SEqHashDef())
119  .def_ro("dtype", &BufferNode::dtype)
120  .def_ro("shape", &BufferNode::shape, refl::AttachFieldFlag::SEqHashDef())
121  .def_ro("strides", &BufferNode::strides, refl::AttachFieldFlag::SEqHashDef())
122  .def_ro("axis_separators", &BufferNode::axis_separators,
123  refl::AttachFieldFlag::SEqHashDef())
124  .def_ro("elem_offset", &BufferNode::elem_offset, refl::AttachFieldFlag::SEqHashDef())
125  .def_ro("name", &BufferNode::name, refl::AttachFieldFlag::SEqHashIgnore())
126  .def_ro("data_alignment", &BufferNode::data_alignment)
127  .def_ro("offset_factor", &BufferNode::offset_factor)
128  .def_ro("buffer_type", &BufferNode::buffer_type)
129  .def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore());
130  }
131 
134  return shape.size() != 0 ? shape[0].dtype() : tvm::tir::DefaultIndexType();
135  }
136 
143  ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index) const;
144 
145  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
146 
149 };
150 
156 class Buffer : public ObjectRef {
157  public:
158  // User can specify data_alignment and offset_factor to be 0
159  // A default value will be picked.
160  TVM_DLL Buffer(Var data, DataType dtype, ffi::Array<PrimExpr> shape, ffi::Array<PrimExpr> strides,
161  PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor,
162  BufferType buffer_type, ffi::Array<IntImm> axis_separators = {},
163  Span span = Span());
164 
170  TVM_DLL Buffer MakeStrideView() const;
179  TVM_DLL Buffer MakeSlice(ffi::Array<PrimExpr> begins, ffi::Array<PrimExpr> extents) const;
188  TVM_DLL PrimExpr access_ptr(int access_mask, DataType ptr_type = DataType::Handle(),
189  int content_lanes = 1, PrimExpr offset = IntImm(DataType::Int(32), 0),
190  ffi::Optional<PrimExpr> input_extent = std::nullopt) const;
198  TVM_DLL PrimExpr vload(ffi::Array<PrimExpr> begin, DataType dtype,
199  ffi::Optional<PrimExpr> predicate = std::nullopt) const;
207  TVM_DLL Stmt vstore(ffi::Array<PrimExpr> begin, PrimExpr value,
208  ffi::Optional<PrimExpr> predicate = std::nullopt) const;
209 
214 
221  ffi::Array<PrimExpr> OffsetOf(ffi::Array<PrimExpr> index) const;
222 
226  TVM_DLL ffi::String scope() const;
227 
230 };
231 
243 TVM_DLL Buffer decl_buffer(ffi::Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
244  ffi::String name = "buffer", ffi::String storage_scope = "",
245  ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
246  Span span = Span());
247 
261  public:
263  virtual ~DataProducerNode() {}
268  virtual ffi::Array<PrimExpr> GetShape() const = 0;
273  virtual DataType GetDataType() const = 0;
278  virtual ffi::String GetNameHint() const = 0;
280 };
281 
287  public:
289 };
290 
304 TVM_DLL tir::Buffer BufferWithOffsetAlignment(ffi::Array<PrimExpr> shape, DataType dtype,
305  std::string name, int data_alignment,
306  int offset_factor, bool compact,
307  std::string memory_scope = "");
308 } // namespace tir
309 } // namespace tvm
310 #endif // TVM_TIR_BUFFER_H_
Managed reference class to IntImmNode.
Definition: expr.h:510
Base class for other IR constructs that can be converted to PrimExpr. This is useful for the FFI to c...
Definition: expr.h:154
Managed reference to PrimExprConvertibleNode.
Definition: expr.h:165
Reference to PrimExprNode.
Definition: expr.h:124
Definition: source_map.h:111
Runtime primitive data type.
Definition: data_type.h:47
static DataType Float(int bits, int lanes=1)
Construct an float type.
Definition: data_type.h:291
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
static DataType Handle(int bits=64, int lanes=1)
Construct a handle type.
Definition: data_type.h:392
Node to represent a buffer.
Definition: buffer.h:62
Span span
Span that points to the original source code. Reserved debug information.
Definition: buffer.h:111
BufferNode()
constructor
Definition: buffer.h:113
ffi::String name
optional name of the buffer
Definition: buffer.h:97
BufferType buffer_type
buffer type
Definition: buffer.h:106
Var data
The pointer to the head of the data.
Definition: buffer.h:69
static void RegisterReflection()
Definition: buffer.h:115
ffi::Array< IntImm > axis_separators
Separators between input axes when generating flattened output axes.
Definition: buffer.h:87
PrimExpr elem_offset
The offset in terms of number of dtype elements (including lanes)
Definition: buffer.h:94
ffi::Array< PrimExpr > shape
The type of the buffer prior to flattening.
Definition: buffer.h:78
ffi::Array< PrimExpr > ElemOffset(ffi::Array< PrimExpr > index) const
Determine the offset in the buffer of the given index.
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: buffer.h:145
ffi::Array< PrimExpr > strides
The strides of each dimension This can be an empty array, indicating array is contiguous.
Definition: buffer.h:92
int offset_factor
Factor of elem_offset field, elem_offset is guaranteed to be multiple of offset_factor.
Definition: buffer.h:104
int data_alignment
Alignment requirement of data pointer in bytes.
Definition: buffer.h:99
DataType DefaultIndexType() const
Definition: buffer.h:133
DataType dtype
data type in the content of the tensor
Definition: buffer.h:71
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Buffer", BufferNode, Object)
Buffer is a symbolic n-darray structure. It is a composition of primitive symbolic types,...
Definition: buffer.h:156
ffi::String scope() const
Return the storage scope associated with this buffer.
PrimExpr vload(ffi::Array< PrimExpr > begin, DataType dtype, ffi::Optional< PrimExpr > predicate=std::nullopt) const
Create an Expr that does a vector load at begin index.
Buffer(Var data, DataType 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())
TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferNode)
Buffer MakeStrideView() const
Return a new buffer that is equivalent with current one but always add stride field.
ffi::Array< PrimExpr > OffsetOf(ffi::Array< PrimExpr > index) const
Determine the offset in the buffer of the given index.
PrimExpr access_ptr(int access_mask, DataType ptr_type=DataType::Handle(), int content_lanes=1, PrimExpr offset=IntImm(DataType::Int(32), 0), ffi::Optional< PrimExpr > input_extent=std::nullopt) const
Get access ptr to the entire buffer.
Buffer MakeSlice(ffi::Array< PrimExpr > begins, ffi::Array< PrimExpr > extents) const
Make a new symbolic buffer representing a slice of the buffer.
Stmt vstore(ffi::Array< PrimExpr > begin, PrimExpr value, ffi::Optional< PrimExpr > predicate=std::nullopt) const
Create a Stmt that does a vector store at begin index.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Buffer, ObjectRef, BufferNode)
Buffer GetFlattenedBuffer() const
Get a flattened version of the buffer.
Base node for data producers.
Definition: buffer.h:260
virtual ffi::String GetNameHint() const =0
Get the name hint of the data producer.
virtual DataType GetDataType() const =0
Get the data type of the result.
TVM_FFI_DECLARE_OBJECT_INFO("tir.DataProducer", DataProducerNode, PrimExprConvertibleNode)
virtual ffi::Array< PrimExpr > GetShape() const =0
Get the shape of the result.
virtual ~DataProducerNode()
destructor.
Definition: buffer.h:263
Managed reference to DataProducerNode.
Definition: buffer.h:286
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataProducer, PrimExprConvertible, DataProducerNode)
Container of all statements.
Definition: stmt.h:63
a named variable in TIR
Definition: var.h:77
Base expr nodes in TVM.
Definition: repr_printer.h:91
constexpr const char * axis_separators
Marks the physical axis separators.
Definition: stmt.h:1094
DataType DefaultIndexType()
if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32
Definition: buffer.h:43
BufferType
buffer type
Definition: buffer.h:55
@ kAutoBroadcast
Definition: buffer.h:58
@ kDefault
Definition: buffer.h:56
tir::Buffer BufferWithOffsetAlignment(ffi::Array< PrimExpr > shape, DataType dtype, std::string name, int data_alignment, int offset_factor, bool compact, std::string memory_scope="")
Creates TIR Buffer for provided parameters.
Buffer decl_buffer(ffi::Array< PrimExpr > shape, DataType dtype=DataType::Float(32), ffi::String name="buffer", ffi::String storage_scope="", ffi::Optional< ffi::Array< IntImm >> axis_separators=std::nullopt, Span span=Span())
Construct a new buffer given shape, and dtype.
Tensor shape(const Tensor &src, DataType dtype, const std::string name="T_shape", const std::string tag=kInjective)
Get the shape of input tensor.
Definition: transform.h:1960
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Variables in the TIR.