tvm
ir.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 #ifndef TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
20 #define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
21 
24 #include <tvm/tir/op.h>
25 
26 namespace tvm {
27 namespace script {
28 namespace ir_builder {
29 namespace tir {
30 
32 using tvm::tir::Buffer;
33 using tvm::tir::Var;
34 
50 Buffer BufferDecl(Array<PrimExpr> shape, DataType dtype, String buffer_name, Optional<Var> data,
51  Optional<Array<PrimExpr>> strides, Optional<PrimExpr> elem_offset,
52  String storage_scope, int align, int offset_factor, String buffer_type,
53  Optional<Array<IntImm>> axis_separators);
54 
59 PrimFuncFrame PrimFunc(bool is_private);
60 
67 Var Arg(String name, Var var);
68 
75 Buffer Arg(String name, Buffer buffer);
76 
81 void FuncName(String name);
82 
87 void FuncAttrs(Map<String, ffi::Any> attrs);
88 
94 Type FuncRet(Type ret_type);
95 
111 Buffer MatchBuffer(ObjectRef param, Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
112  Optional<Var> data = std::nullopt, Array<PrimExpr> strides = {},
113  PrimExpr elem_offset = PrimExpr(), String storage_scope = "global",
114  int align = -1, int offset_factor = 0, String buffer_type = "default",
115  Optional<Array<IntImm>> axis_separators = std::nullopt);
116 
123 BlockFrame Block(String name, bool no_realize = false);
124 
130 
135 void Where(PrimExpr predicate);
136 
141 void Reads(Array<ObjectRef> buffer_slices);
142 
147 void Writes(Array<ObjectRef> buffer_slices);
148 
153 void BlockAttrs(Map<String, ffi::Any> attrs);
154 
169 Buffer AllocBuffer(Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
170  Optional<Var> data = std::nullopt, Array<PrimExpr> strides = {},
171  PrimExpr elem_offset = PrimExpr(), String storage_scope = "", int align = -1,
172  int offset_factor = 0, String buffer_type = "default",
173  Optional<Array<IntImm>> axis_separators = std::nullopt);
174 namespace axis {
175 
183 Var Spatial(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
184 
192 Var Reduce(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
193 
201 Var Scan(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
202 
210 Var Opaque(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
211 
219 Array<Var> Remap(String kinds, Array<PrimExpr> bindings, DataType dtype = DataType::Int(32));
220 
221 } // namespace axis
222 
231  Optional<Map<String, Any>> annotations = std::nullopt);
240  Optional<Map<String, Any>> annotations = std::nullopt);
249  Optional<Map<String, Any>> annotations = std::nullopt);
258  Optional<Map<String, Any>> annotations = std::nullopt);
267 ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, String thread,
268  Optional<Map<String, Any>> annotations = std::nullopt);
274 ForFrame Grid(Array<PrimExpr> extents);
275 
282 AssertFrame Assert(PrimExpr condition, String message);
283 
293 LetFrame LetStmt(PrimExpr value, Optional<Type> type_annotation = std::nullopt,
294  Optional<Var> var = std::nullopt);
295 
303 RealizeFrame Realize(tvm::tir::BufferRegion buffer_slice, String storage_scope, PrimExpr condition);
304 
314 AllocateFrame Allocate(Array<PrimExpr> extents, DataType dtype, String storage_scope = "",
315  Optional<PrimExpr> condition = std::nullopt,
316  Optional<Map<String, Any>> annotations = std::nullopt);
317 
326 AllocateConstFrame AllocateConst(NDArray data, DataType dtype, Array<PrimExpr> extents,
327  Optional<Map<String, Any>> annotations = std::nullopt);
328 
336 AttrFrame Attr(ffi::Any node, String attr_key, PrimExpr value);
337 
344 
350 IfFrame If(PrimExpr condition);
351 
357 
363 
379 DeclBufferFrame DeclBuffer(Array<PrimExpr> shape, DataType dtype, String buffer_name,
380  Optional<Var> data, Optional<Array<PrimExpr>> strides,
381  Optional<PrimExpr> elem_offset, String storage_scope, int align,
382  int offset_factor, String buffer_type,
383  Optional<Array<IntImm>> axis_separators);
384 
392 
399 LaunchThreadFrame LaunchThread(String thread_tag, PrimExpr extent);
400 
407 Var EnvThread(String thread_tag, DataType dtype = DataType::Int(32));
408 
417 void BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr> indices,
418  Optional<PrimExpr> predicate);
419 
424 void Evaluate(PrimExpr value);
425 
444  String storage_scope = "global", bool is_size_var = false,
445  bool is_unknown_type = false) {
446  Type type_annotation{nullptr};
447  if (is_unknown_type && storage_scope == "global") {
448  type_annotation = PrimType(runtime::DataType::Handle());
449  } else {
450  type_annotation = PointerType(PrimType(dtype), storage_scope);
451  }
452  return is_size_var ? tvm::tir::SizeVar("", type_annotation) : tvm::tir::Var("", type_annotation);
453 }
454 
456 
457 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType) \
458  inline PrimExpr FuncName(Optional<PrimExpr> expr = std::nullopt, bool is_size_var = false) { \
459  DataType dtype = DType; \
460  return expr.defined() \
461  ? tvm::cast(dtype, expr.value()) \
462  : (is_size_var ? tvm::tir::SizeVar("", dtype) : tvm::tir::Var("", dtype)); \
463  }
464 
465 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType) \
466  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##8, FDType(8)); \
467  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##16, FDType(16)); \
468  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##32, FDType(32)); \
469  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##64, FDType(64));
470 
475 
476 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, FDType, Size) \
477  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, FDType(Size, 4)); \
478  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, FDType(Size, 8)); \
479  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, FDType(Size, 16)); \
480  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, FDType(Size, 32)); \
481  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, FDType(Size, 64));
482 
483 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType) \
484  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, FDType, 8); \
485  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, FDType, 16); \
486  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, FDType, 32); \
487  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, FDType, 64);
488 
493 
494 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
495  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1)); \
496  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4)); \
497  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8)); \
498  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16)); \
499  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, FDType(32)); \
500  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, FDType(64));
501 
510 
513 
515 
518 
519 #undef TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST
520 
521 } // namespace tir
522 } // namespace ir_builder
523 } // namespace script
524 } // namespace tvm
525 
526 #endif // TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
Definition: type.h:182
Reference to PrimExprNode.
Definition: expr.h:129
Definition: type.h:134
Range container
Definition: expr.h:698
Managed reference to TensorMapTypeNode.
Definition: type.h:316
Managed reference to TypeNode.
Definition: type.h:101
Runtime primitive data type.
Definition: data_type.h:47
static DataType Float8E4M3FNUZ(int lanes=1)
Construct float8 e4m3fnuz datatype.
Definition: data_type.h:334
static DataType Float4E2M1FN(int lanes=1)
Construct float4 e2m1fn datatype.
Definition: data_type.h:376
static DataType Float(int bits, int lanes=1)
Construct an float type.
Definition: data_type.h:291
static DataType Float8E8M0FNU(int lanes=1)
Construct float8 e8m0fnu datatype.
Definition: data_type.h:355
static DataType Float8E4M3FN(int lanes=1)
Construct float8 e4m3fn datatype.
Definition: data_type.h:327
static DataType Float8E5M2FNUZ(int lanes=1)
Construct float8 e5m2fnuz datatype.
Definition: data_type.h:348
static DataType Float8E4M3B11FNUZ(int lanes=1)
Construct float8 e4m3b11fnuz datatype.
Definition: data_type.h:318
static DataType Float8E5M2(int lanes=1)
Construct float8 e5m2 datatype.
Definition: data_type.h:341
static DataType Float6E2M3FN(int lanes=1)
Construct float6 e2m3fn datatype.
Definition: data_type.h:362
static DataType Bool(int lanes=1, bool is_scalable=false)
Construct a bool type.
Definition: data_type.h:383
static DataType Float6E3M2FN(int lanes=1)
Construct float6 e3m2fn datatype.
Definition: data_type.h:369
static DataType BFloat(int bits, int lanes=1)
Construct an bfloat type.
Definition: data_type.h:298
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
static DataType Void()
Construct a Void type.
Definition: data_type.h:397
static DataType Handle(int bits=64, int lanes=1)
Construct a handle type.
Definition: data_type.h:392
static DataType UInt(int bits, int lanes=1, bool is_scalable=false)
Construct an uint type.
Definition: data_type.h:282
static DataType Float8E4M3(int lanes=1)
Construct float8 e4m3 datatype.
Definition: data_type.h:311
static DataType Float8E3M4(int lanes=1)
Construct float8 e3m4 datatype.
Definition: data_type.h:304
Managed NDArray. The array is backed by reference counted blocks.
Definition: ndarray.h:53
Managed reference to AllocateConstFrameNode.
Definition: frame.h:545
Managed reference to AllocateFrameNode.
Definition: frame.h:496
Managed reference to AssertFrameNode.
Definition: frame.h:318
Managed reference to AttrFrameNode.
Definition: frame.h:588
Managed reference to BlockFrameNode.
Definition: frame.h:187
Managed reference to BlockInitFrameNode.
Definition: frame.h:225
Managed reference to ElseFrameNode.
Definition: frame.h:743
Managed reference to ForFrameNode.
Definition: frame.h:277
Managed reference to IfFrameNode.
Definition: frame.h:667
Managed reference to LaunchThreadFrameNode.
Definition: frame.h:400
Managed reference to LetFrameNode.
Definition: frame.h:358
Managed reference to PrimFuncFrameNode.
Definition: frame.h:116
Managed reference to RealizeFrameNode.
Definition: frame.h:444
Managed reference to ThenFrameNode.
Definition: frame.h:705
Managed reference to WhileFrameNode.
Definition: frame.h:624
Managed reference to BufferRegionNode.
Definition: stmt.h:871
Buffer is a symbolic n-darray structure. It is a composition of primitive symbolic types,...
Definition: buffer.h:157
a named variable represents a tensor index size
Definition: var.h:144
a named variable in TIR
Definition: var.h:78
Var Opaque(Range dom, PrimExpr binding, DataType dtype=DataType::Int(32))
The opaque block axis defining function.
Var Reduce(Range dom, PrimExpr binding, DataType dtype=DataType::Int(32))
The reduced block axis defining function.
Var Scan(Range dom, PrimExpr binding, DataType dtype=DataType::Int(32))
The scanning block axis defining function.
Var Spatial(Range dom, PrimExpr binding, DataType dtype=DataType::Int(32))
The spatial block axis defining function.
Array< Var > Remap(String kinds, Array< PrimExpr > bindings, DataType dtype=DataType::Int(32))
The block axis remapping function.
PrimExpr Float8E5M2FNUZ(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:508
ForFrame Grid(Array< PrimExpr > extents)
The grid For statement.
void Evaluate(PrimExpr value)
Evaluate the input expression.
PrimExpr Float8E3M4(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:502
PrimExpr Void(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:517
Buffer MatchBuffer(ObjectRef param, Array< PrimExpr > shape, DataType dtype=DataType::Float(32), Optional< Var > data=std::nullopt, Array< PrimExpr > strides={}, PrimExpr elem_offset=PrimExpr(), String storage_scope="global", int align=-1, int offset_factor=0, String buffer_type="default", Optional< Array< IntImm >> axis_separators=std::nullopt)
The buffer match statement.
LaunchThreadFrame LaunchThread(Var var, PrimExpr extent)
Launch a thread.
ForFrame Vectorized(PrimExpr start, PrimExpr stop, Optional< Map< String, Any >> annotations=std::nullopt)
The vectorized For statement.
PrimFuncFrame PrimFunc(bool is_private)
The primitive function statement.
PrimExpr Float8E4M3(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:503
Buffer AllocBuffer(Array< PrimExpr > shape, DataType dtype=DataType::Float(32), Optional< Var > data=std::nullopt, Array< PrimExpr > strides={}, PrimExpr elem_offset=PrimExpr(), String storage_scope="", int align=-1, int offset_factor=0, String buffer_type="default", Optional< Array< IntImm >> axis_separators=std::nullopt)
The buffer allocation function.
Var Handle(runtime::DataType dtype=runtime::DataType::Void(), String storage_scope="global", bool is_size_var=false, bool is_unknown_type=false)
Create a TIR var that represents a pointer.
Definition: ir.h:443
ThenFrame Then()
Create a then.
RealizeFrame Realize(tvm::tir::BufferRegion buffer_slice, String storage_scope, PrimExpr condition)
The realization.
IfFrame If(PrimExpr condition)
Create an if statement.
Buffer BufferDecl(Array< PrimExpr > shape, DataType dtype, String buffer_name, Optional< Var > data, Optional< Array< PrimExpr >> strides, Optional< PrimExpr > elem_offset, String storage_scope, int align, int offset_factor, String buffer_type, Optional< Array< IntImm >> axis_separators)
The buffer declaration function.
ElseFrame Else()
Create an else.
PrimExpr Float8E4M3FNUZ(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:506
void Where(PrimExpr predicate)
The block predicate statement.
LetFrame LetStmt(PrimExpr value, Optional< Type > type_annotation=std::nullopt, Optional< Var > var=std::nullopt)
The let binding.
Type FuncRet(Type ret_type)
The PrimFunc return type statement.
AllocateConstFrame AllocateConst(NDArray data, DataType dtype, Array< PrimExpr > extents, Optional< Map< String, Any >> annotations=std::nullopt)
The allocate constant node.
void FuncAttrs(Map< String, ffi::Any > attrs)
The PrimFunc annotation statement.
ForFrame Serial(PrimExpr start, PrimExpr stop, Optional< Map< String, Any >> annotations=std::nullopt)
The serial For statement.
PrimExpr Float8E8M0FNU(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:509
void FuncName(String name)
The PrimFunc naming statement.
WhileFrame While(PrimExpr condition)
Create a while loop.
DeclBufferFrame DeclBuffer(Array< PrimExpr > shape, DataType dtype, String buffer_name, Optional< Var > data, Optional< Array< PrimExpr >> strides, Optional< PrimExpr > elem_offset, String storage_scope, int align, int offset_factor, String buffer_type, Optional< Array< IntImm >> axis_separators)
The buffer declaration frame.
Var EnvThread(String thread_tag, DataType dtype=DataType::Int(32))
Bind a var to thread env.
PrimExpr Float4E2M1FN(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:514
ForFrame Parallel(PrimExpr start, PrimExpr stop, Optional< Map< String, Any >> annotations=std::nullopt)
The parallel For statement.
PrimExpr Float6E3M2FN(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:512
ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, String thread, Optional< Map< String, Any >> annotations=std::nullopt)
The thread-binding For statement.
ForFrame Unroll(PrimExpr start, PrimExpr stop, Optional< Map< String, Any >> annotations=std::nullopt)
The unrolled For statement.
void BufferStore(Buffer buffer, PrimExpr value, Array< PrimExpr > indices, Optional< PrimExpr > predicate)
Store data in a buffer.
void BlockAttrs(Map< String, ffi::Any > attrs)
The block annotation statement.
PrimExpr Float8E4M3FN(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:505
Var TensormapHandle()
Definition: ir.h:455
PrimExpr Float6E2M3FN(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:511
PrimExpr Float8E4M3B11FNUZ(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:504
Var Arg(String name, Var var)
The PrimFunc variable arguments adding function.
AllocateFrame Allocate(Array< PrimExpr > extents, DataType dtype, String storage_scope="", Optional< PrimExpr > condition=std::nullopt, Optional< Map< String, Any >> annotations=std::nullopt)
The allocate node.
void Reads(Array< ObjectRef > buffer_slices)
The block buffer region reading statement.
BlockInitFrame Init()
The block initialization statement.
AttrFrame Attr(ffi::Any node, String attr_key, PrimExpr value)
Create an attribute.
void Writes(Array< ObjectRef > buffer_slices)
The block buffer region writing statement.
BlockFrame Block(String name, bool no_realize=false)
The block declaration statement.
AssertFrame Assert(PrimExpr condition, String message)
The assertion statement.
PrimExpr Float8E5M2(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:507
PrimExpr Boolean(Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:516
Var var(std::string name_hint, DataType t=DataType::Int(32))
Construct a new Var expression.
constexpr const char * axis_separators
Marks the physical axis separators.
Definition: stmt.h:1123
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
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType)
Definition: ir.h:457
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType)
Definition: ir.h:465
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType)
Definition: ir.h:483
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType)
Definition: ir.h:494
Common operators defined for Expr.