19 #ifndef TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
20 #define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
28 namespace ir_builder {
51 Optional<Array<PrimExpr>> strides, Optional<PrimExpr> elem_offset,
52 String storage_scope,
int align,
int offset_factor, String buffer_type,
112 Optional<Var> data = std::nullopt, Array<PrimExpr> strides = {},
114 int align = -1,
int offset_factor = 0, String buffer_type =
"default",
141 void Reads(Array<ObjectRef> buffer_slices);
147 void Writes(Array<ObjectRef> buffer_slices);
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",
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);
268 Optional<Map<String, Any>> annotations = std::nullopt);
294 Optional<Var>
var = std::nullopt);
315 Optional<PrimExpr> condition = std::nullopt,
316 Optional<Map<String, Any>> annotations = std::nullopt);
327 Optional<Map<String, Any>> annotations = std::nullopt);
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,
418 Optional<PrimExpr> predicate);
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") {
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)); \
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));
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));
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);
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));
519 #undef TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST
Reference to PrimExprNode.
Definition: expr.h:129
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.