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  Array<PrimExpr> shape;
87  Array<IntImm> axis_separators;
92  Array<PrimExpr> strides;
95  // Meta data
97  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  Array<PrimExpr> ElemOffset(Array<PrimExpr> index) const;
144 
145  static constexpr const char* _type_key = "tir.Buffer";
146  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
147 
150 };
151 
157 class Buffer : public ObjectRef {
158  public:
159  // User can specify data_alignment and offset_factor to be 0
160  // A default value will be picked.
161  TVM_DLL Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr> strides,
162  PrimExpr elem_offset, String name, int data_alignment, int offset_factor,
163  BufferType buffer_type, Array<IntImm> axis_separators = {}, Span span = Span());
164 
170  TVM_DLL Buffer MakeStrideView() const;
179  TVM_DLL Buffer MakeSlice(Array<PrimExpr> begins, 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  Optional<PrimExpr> input_extent = std::nullopt) const;
198  TVM_DLL PrimExpr vload(Array<PrimExpr> begin, DataType dtype,
199  Optional<PrimExpr> predicate = std::nullopt) const;
207  TVM_DLL Stmt vstore(Array<PrimExpr> begin, PrimExpr value,
208  Optional<PrimExpr> predicate = std::nullopt) const;
209 
214 
221  Array<PrimExpr> OffsetOf(Array<PrimExpr> index) const;
222 
226  TVM_DLL String scope() const;
227 
230 };
231 
243 TVM_DLL Buffer decl_buffer(Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
244  String name = "buffer", String storage_scope = "",
245  Optional<Array<IntImm>> axis_separators = std::nullopt,
246  Span span = Span());
247 
261  public:
263  virtual ~DataProducerNode() {}
268  virtual Array<PrimExpr> GetShape() const = 0;
273  virtual DataType GetDataType() const = 0;
278  virtual String GetNameHint() const = 0;
279 
280  static constexpr const char* _type_key = "tir.DataProducer";
282 };
283 
289  public:
291 };
292 
306 TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
307  std::string name, int data_alignment,
308  int offset_factor, bool compact,
309  std::string memory_scope = "");
310 } // namespace tir
311 } // namespace tvm
312 #endif // TVM_TIR_BUFFER_H_
Managed reference class to IntImmNode.
Definition: expr.h:520
Base class for other IR constructs that can be converted to PrimExpr. This is useful for the FFI to c...
Definition: expr.h:159
Managed reference to PrimExprConvertibleNode.
Definition: expr.h:172
Reference to PrimExprNode.
Definition: expr.h:129
Definition: source_map.h:113
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
String name
optional name of the buffer
Definition: buffer.h:97
static constexpr const char * _type_key
Definition: buffer.h:145
Span span
Span that points to the original source code. Reserved debug information.
Definition: buffer.h:111
BufferNode()
constructor
Definition: buffer.h:113
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
Array< PrimExpr > shape
The type of the buffer prior to flattening.
Definition: buffer.h:78
PrimExpr elem_offset
The offset in terms of number of dtype elements (including lanes)
Definition: buffer.h:94
TVM_DECLARE_FINAL_OBJECT_INFO(BufferNode, Object)
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: buffer.h:146
int offset_factor
Factor of elem_offset field, elem_offset is guaranteed to be multiple of offset_factor.
Definition: buffer.h:104
Array< IntImm > axis_separators
Separators between input axes when generating flattened output axes.
Definition: buffer.h:87
int data_alignment
Alignment requirement of data pointer in bytes.
Definition: buffer.h:99
Array< PrimExpr > strides
The strides of each dimension This can be an empty array, indicating array is contiguous.
Definition: buffer.h:92
DataType DefaultIndexType() const
Definition: buffer.h:133
DataType dtype
data type in the content of the tensor
Definition: buffer.h:71
Array< PrimExpr > ElemOffset(Array< PrimExpr > index) const
Determine the offset in the buffer of the given index.
Buffer is a symbolic n-darray structure. It is a composition of primitive symbolic types,...
Definition: buffer.h:157
Array< PrimExpr > OffsetOf(Array< PrimExpr > index) const
Determine the offset in the buffer of the given index.
String scope() const
Return the storage scope associated with this buffer.
Stmt vstore(Array< PrimExpr > begin, PrimExpr value, Optional< PrimExpr > predicate=std::nullopt) const
Create a Stmt that does a vector store at begin index.
PrimExpr vload(Array< PrimExpr > begin, DataType dtype, Optional< PrimExpr > predicate=std::nullopt) const
Create an Expr that does a vector load at begin index.
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.
PrimExpr access_ptr(int access_mask, DataType ptr_type=DataType::Handle(), int content_lanes=1, PrimExpr offset=IntImm(DataType::Int(32), 0), Optional< PrimExpr > input_extent=std::nullopt) const
Get access ptr to the entire buffer.
Buffer(Var data, DataType dtype, Array< PrimExpr > shape, Array< PrimExpr > strides, PrimExpr elem_offset, String name, int data_alignment, int offset_factor, BufferType buffer_type, Array< IntImm > axis_separators={}, Span span=Span())
Buffer MakeSlice(Array< PrimExpr > begins, Array< PrimExpr > extents) const
Make a new symbolic buffer representing a slice of the buffer.
TVM_DEFINE_OBJECT_REF_METHODS(Buffer, ObjectRef, BufferNode)
Buffer GetFlattenedBuffer() const
Get a flattened version of the buffer.
Base node for data producers.
Definition: buffer.h:260
static constexpr const char * _type_key
Definition: buffer.h:280
virtual DataType GetDataType() const =0
Get the data type of the result.
virtual Array< PrimExpr > GetShape() const =0
Get the shape of the result.
virtual String GetNameHint() const =0
Get the name hint of the data producer.
TVM_DECLARE_BASE_OBJECT_INFO(DataProducerNode, PrimExprConvertibleNode)
virtual ~DataProducerNode()
destructor.
Definition: buffer.h:263
Managed reference to DataProducerNode.
Definition: buffer.h:288
TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, PrimExprConvertible, DataProducerNode)
Container of all statements.
Definition: stmt.h:64
a named variable in TIR
Definition: var.h:78
Base expr nodes in TVM.
Definition: repr_printer.h:91
constexpr const char * axis_separators
Marks the physical axis separators.
Definition: stmt.h:1123
DataType DefaultIndexType()
if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32
Definition: buffer.h:43
Buffer decl_buffer(Array< PrimExpr > shape, DataType dtype=DataType::Float(32), String name="buffer", String storage_scope="", Optional< Array< IntImm >> axis_separators=std::nullopt, Span span=Span())
Construct a new buffer given shape, and dtype.
BufferType
buffer type
Definition: buffer.h:55
@ kAutoBroadcast
Definition: buffer.h:58
@ kDefault
Definition: buffer.h:56
tir::Buffer BufferWithOffsetAlignment(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.
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:1945
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Variables in the TIR.