19 #ifndef TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
20 #define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
28 namespace ir_builder {
51 ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
52 ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope,
int align,
53 int offset_factor, ffi::String buffer_type,
88 void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
115 ffi::String storage_scope =
"global",
int align = -1,
int offset_factor = 0,
116 ffi::String buffer_type =
"default",
143 void Reads(ffi::Array<ObjectRef> buffer_slices);
149 void Writes(ffi::Array<ObjectRef> buffer_slices);
172 ffi::Optional<Var> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
174 int align = -1,
int offset_factor = 0, ffi::String buffer_type =
"default",
221 ffi::Array<Var>
Remap(ffi::String kinds, ffi::Array<PrimExpr> bindings,
234 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
243 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
252 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
261 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
271 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
297 ffi::Optional<Var>
var = std::nullopt);
319 ffi::Optional<PrimExpr> condition = std::nullopt,
320 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
332 ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
385 ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
386 ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope,
387 int align,
int offset_factor, ffi::String buffer_type,
423 ffi::Optional<PrimExpr> predicate);
449 ffi::String storage_scope =
"global",
bool is_size_var =
false,
450 bool is_unknown_type =
false) {
451 Type type_annotation{
nullptr};
452 if (is_unknown_type && storage_scope ==
"global") {
462 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType) \
463 inline PrimExpr FuncName(ffi::Optional<PrimExpr> expr = std::nullopt, \
464 bool is_size_var = false) { \
465 DataType dtype = DType; \
466 return expr.defined() \
467 ? tvm::cast(dtype, expr.value()) \
468 : (is_size_var ? tvm::tir::SizeVar("", dtype) : tvm::tir::Var("", dtype)); \
471 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType) \
472 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##8, FDType(8)); \
473 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##16, FDType(16)); \
474 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##32, FDType(32)); \
475 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##64, FDType(64));
482 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, FDType, Size) \
483 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, FDType(Size, 4)); \
484 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, FDType(Size, 8)); \
485 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, FDType(Size, 16)); \
486 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, FDType(Size, 32)); \
487 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, FDType(Size, 64));
489 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType) \
490 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, FDType, 8); \
491 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, FDType, 16); \
492 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, FDType, 32); \
493 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, FDType, 64);
500 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
501 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1)); \
502 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4)); \
503 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8)); \
504 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16)); \
505 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, FDType(32)); \
506 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, FDType(64));
525 #undef TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST
Reference to PrimExprNode.
Definition: expr.h:124
Range container
Definition: expr.h:689
Managed reference to TensorMapTypeNode.
Definition: type.h:305
Managed reference to TypeNode.
Definition: type.h:100
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 Tensor. The array is backed by reference counted blocks.
Definition: tensor.h:53
Managed reference to AllocateConstFrameNode.
Definition: frame.h:570
Managed reference to AllocateFrameNode.
Definition: frame.h:518
Managed reference to AssertFrameNode.
Definition: frame.h:330
Managed reference to AttrFrameNode.
Definition: frame.h:616
Managed reference to BlockFrameNode.
Definition: frame.h:190
Managed reference to BlockInitFrameNode.
Definition: frame.h:231
Managed reference to ElseFrameNode.
Definition: frame.h:778
Managed reference to ForFrameNode.
Definition: frame.h:286
Managed reference to IfFrameNode.
Definition: frame.h:700
Managed reference to LaunchThreadFrameNode.
Definition: frame.h:417
Managed reference to LetFrameNode.
Definition: frame.h:372
Managed reference to PrimFuncFrameNode.
Definition: frame.h:116
Managed reference to RealizeFrameNode.
Definition: frame.h:463
Managed reference to ThenFrameNode.
Definition: frame.h:739
Managed reference to WhileFrameNode.
Definition: frame.h:655
Managed reference to BufferRegionNode.
Definition: stmt.h:846
Buffer is a symbolic n-darray structure. It is a composition of primitive symbolic types,...
Definition: buffer.h:156
a named variable represents a tensor index size
Definition: var.h:143
a named variable in TIR
Definition: var.h:77
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.
ffi::Array< Var > Remap(ffi::String kinds, ffi::Array< PrimExpr > bindings, DataType dtype=DataType::Int(32))
The block axis remapping function.
LetFrame LetStmt(PrimExpr value, ffi::Optional< Type > type_annotation=std::nullopt, ffi::Optional< Var > var=std::nullopt)
The let binding.
PrimExpr Float8E4M3FNUZ(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:512
void Evaluate(PrimExpr value)
Evaluate the input expression.
PrimExpr Boolean(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:522
ForFrame Grid(ffi::Array< PrimExpr > extents)
The grid For statement.
PrimExpr Float8E4M3FN(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:511
void FuncAttrs(ffi::Map< ffi::String, ffi::Any > attrs)
The PrimFunc annotation statement.
PrimExpr Float8E3M4(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:508
LaunchThreadFrame LaunchThread(Var var, PrimExpr extent)
Launch a thread.
PrimFuncFrame PrimFunc(bool is_private)
The primitive function statement.
PrimExpr Float8E4M3(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:509
ThenFrame Then()
Create a then.
void Reads(ffi::Array< ObjectRef > buffer_slices)
The block buffer region reading statement.
RealizeFrame Realize(tvm::tir::BufferRegion buffer_slice, ffi::String storage_scope, PrimExpr condition)
The realization.
IfFrame If(PrimExpr condition)
Create an if statement.
ElseFrame Else()
Create an else.
void Where(PrimExpr predicate)
The block predicate statement.
ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The thread-binding For statement.
Type FuncRet(Type ret_type)
The PrimFunc return type statement.
PrimExpr Float4E2M1FN(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:520
Buffer BufferDecl(ffi::Array< PrimExpr > shape, DataType dtype, ffi::String buffer_name, ffi::Optional< Var > data, ffi::Optional< ffi::Array< PrimExpr >> strides, ffi::Optional< PrimExpr > elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::String buffer_type, ffi::Optional< ffi::Array< IntImm >> axis_separators)
The buffer declaration function.
WhileFrame While(PrimExpr condition)
Create a while loop.
void BlockAttrs(ffi::Map< ffi::String, ffi::Any > attrs)
The block annotation statement.
ForFrame Parallel(PrimExpr start, PrimExpr stop, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The parallel For statement.
void FuncName(ffi::String name)
The PrimFunc naming statement.
Buffer MatchBuffer(ObjectRef param, ffi::Array< PrimExpr > shape, DataType dtype=DataType::Float(32), ffi::Optional< Var > data=std::nullopt, ffi::Array< PrimExpr > strides={}, PrimExpr elem_offset=PrimExpr(), ffi::String storage_scope="global", int align=-1, int offset_factor=0, ffi::String buffer_type="default", ffi::Optional< ffi::Array< IntImm >> axis_separators=std::nullopt)
The buffer match statement.
AllocateConstFrame AllocateConst(Tensor data, DataType dtype, ffi::Array< PrimExpr > extents, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The allocate constant node.
BlockFrame Block(ffi::String name, bool no_realize=false)
The block declaration statement.
PrimExpr Float8E5M2(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:513
Var Handle(runtime::DataType dtype=runtime::DataType::Void(), ffi::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:448
void Writes(ffi::Array< ObjectRef > buffer_slices)
The block buffer region writing statement.
PrimExpr Void(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:523
PrimExpr Float6E2M3FN(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:517
Var TensormapHandle()
Definition: ir.h:460
ForFrame Unroll(PrimExpr start, PrimExpr stop, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The unrolled For statement.
AssertFrame Assert(PrimExpr condition, ffi::String message)
The assertion statement.
PrimExpr Float8E5M2FNUZ(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:514
BlockInitFrame Init()
The block initialization statement.
Var Arg(ffi::String name, Var var)
The PrimFunc variable arguments adding function.
AttrFrame Attr(ffi::Any node, ffi::String attr_key, PrimExpr value)
Create an attribute.
PrimExpr Float6E3M2FN(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:518
DeclBufferFrame DeclBuffer(ffi::Array< PrimExpr > shape, DataType dtype, ffi::String buffer_name, ffi::Optional< Var > data, ffi::Optional< ffi::Array< PrimExpr >> strides, ffi::Optional< PrimExpr > elem_offset, ffi::String storage_scope, int align, int offset_factor, ffi::String buffer_type, ffi::Optional< ffi::Array< IntImm >> axis_separators)
The buffer declaration frame.
Buffer AllocBuffer(ffi::Array< PrimExpr > shape, DataType dtype=DataType::Float(32), ffi::Optional< Var > data=std::nullopt, ffi::Array< PrimExpr > strides={}, PrimExpr elem_offset=PrimExpr(), ffi::String storage_scope="", int align=-1, int offset_factor=0, ffi::String buffer_type="default", ffi::Optional< ffi::Array< IntImm >> axis_separators=std::nullopt)
The buffer allocation function.
PrimExpr Float8E4M3B11FNUZ(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:510
PrimExpr Float8E8M0FNU(ffi::Optional< PrimExpr > expr=std::nullopt, bool is_size_var=false)
Definition: ir.h:515
AllocateFrame Allocate(ffi::Array< PrimExpr > extents, DataType dtype, ffi::String storage_scope="", ffi::Optional< PrimExpr > condition=std::nullopt, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The allocate node.
ForFrame Serial(PrimExpr start, PrimExpr stop, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The serial For statement.
Var EnvThread(ffi::String thread_tag, DataType dtype=DataType::Int(32))
Bind a var to thread env.
ForFrame Vectorized(PrimExpr start, PrimExpr stop, ffi::Optional< ffi::Map< ffi::String, Any >> annotations=std::nullopt)
The vectorized For statement.
void BufferStore(Buffer buffer, PrimExpr value, ffi::Array< PrimExpr > indices, ffi::Optional< PrimExpr > predicate)
Store data in a buffer.
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:1094
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
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType)
Definition: ir.h:462
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType)
Definition: ir.h:471
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType)
Definition: ir.h:489
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType)
Definition: ir.h:500
Common operators defined for Expr.