wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/tirx/script/builder/frame.h
+740
View File
@@ -0,0 +1,740 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
#ifndef TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/stmt.h>
#include <utility>
namespace tvm {
namespace script {
namespace ir_builder {
namespace tirx {
/*!
* \brief A base frame that represents the TIR fame with body of statements.
*
* \sa TIRFrame
*/
class TIRFrameNode : public IRBuilderFrameNode {
public:
/*! \brief The Stmt within in this frame. */
ffi::Array<tvm::tirx::Stmt> stmts;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TIRFrameNode>().def_ro("stmts", &TIRFrameNode::stmts);
}
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.tirx.TIRFrame", TIRFrameNode, IRBuilderFrameNode);
};
/*!
* \brief Managed reference to TIRFrameNode.
*
* \sa TIRFrameNode
*/
class TIRFrame : public IRBuilderFrame {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TIRFrame, IRBuilderFrame, TIRFrameNode);
protected:
TIRFrame() = default;
explicit TIRFrame(ffi::ObjectPtr<TIRFrameNode> data) : IRBuilderFrame(data) {}
};
/*!
* \brief A frame that represents the PrimFunc containing TIR statements.
*
* \sa PrimFuncFrame
*/
class PrimFuncFrameNode : public TIRFrameNode {
public:
/*! \brief The name of the block. */
ffi::Optional<ffi::String> name;
/*! \brief Function parameters. */
ffi::Array<tvm::tirx::Var> args;
/*! \brief Whether the PrimFunc is annotated as private. */
bool is_private;
/*! \brief The return type of the function. */
ffi::Optional<Type> ret_type;
/*! \brief Maps some parameters to specific Buffer data structures. */
ffi::Map<tvm::tirx::Var, tvm::tirx::Buffer> buffer_map;
/*! \brief Additional attributes storing the meta-data */
ffi::Map<ffi::String, Any> attrs;
/*! \brief The variable map bound to thread env. */
ffi::Map<tvm::tirx::Var, tvm::tirx::IterVar> env_threads;
/*! \brief The buffer allocated in root block. */
ffi::Array<tvm::tirx::Buffer> root_alloc_buffers;
// TIR utils
/*! \brief Whether this PrimFunc uses s_tir semantics (root SBlock wrap,
* parser layout default = None). Default (false) = tirx semantics. */
bool s_tir;
/*! \brief Whether it is a persistent kernel. */
bool persistent;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PrimFuncFrameNode>()
.def_ro("name", &PrimFuncFrameNode::name)
.def_ro("args", &PrimFuncFrameNode::args)
.def_ro("is_private", &PrimFuncFrameNode::is_private)
.def_ro("ret_type", &PrimFuncFrameNode::ret_type)
.def_ro("buffer_map", &PrimFuncFrameNode::buffer_map)
.def_ro("attrs", &PrimFuncFrameNode::attrs)
.def_ro("env_threads", &PrimFuncFrameNode::env_threads)
.def_ro("root_alloc_buffers", &PrimFuncFrameNode::root_alloc_buffers)
.def_ro("s_tir", &PrimFuncFrameNode::s_tir)
.def_ro("persistent", &PrimFuncFrameNode::persistent);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.PrimFuncFrame", PrimFuncFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to PrimFuncFrameNode.
*
* \sa PrimFuncFrameNode
*/
class PrimFuncFrame : public TIRFrame {
public:
explicit PrimFuncFrame(ffi::ObjectPtr<PrimFuncFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimFuncFrame, TIRFrame, PrimFuncFrameNode);
};
/*!
* \brief A frame that represents the block.
*
* \sa SBlockFrame
*/
class SBlockFrameNode : public TIRFrameNode {
public:
/*! \brief The name of the block. */
ffi::String name;
/*! \brief The variables of the block. */
ffi::Array<tvm::tirx::IterVar> iter_vars;
/*! \brief The read buffer regions of the block. */
ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> reads;
/*! \brief The write buffer regions of the block. */
ffi::Optional<ffi::Array<tvm::tirx::BufferRegion>> writes;
/*! \brief The init statement of the bolck. */
ffi::Optional<tvm::tirx::Stmt> init;
/*! \brief The buffer allocated in the block. */
ffi::Array<tvm::tirx::Buffer> alloc_buffers;
/*! \brief The match buffer regions. */
ffi::Array<tvm::tirx::MatchBufferRegion> match_buffers;
/*! \brief The annotation of the block. */
ffi::Optional<ffi::Map<ffi::String, Any>> annotations;
/*! \brief The corresponding values of the iter vars. */
ffi::Array<PrimExpr> iter_values;
/*!
* \brief The predicate of the block realization, the block will only be executed when the
* predicate is true.
*/
ffi::Optional<PrimExpr> predicate;
/*! \brief The flag whether to construct BlockRealize or Block. */
bool no_realize;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SBlockFrameNode>()
.def_ro("name", &SBlockFrameNode::name)
.def_ro("iter_vars", &SBlockFrameNode::iter_vars)
.def_ro("reads", &SBlockFrameNode::reads)
.def_ro("writes", &SBlockFrameNode::writes)
.def_ro("init", &SBlockFrameNode::init)
.def_ro("alloc_buffers", &SBlockFrameNode::alloc_buffers)
.def_ro("match_buffers", &SBlockFrameNode::match_buffers)
.def_ro("annotations", &SBlockFrameNode::annotations)
.def_ro("iter_values", &SBlockFrameNode::iter_values)
.def_ro("predicate", &SBlockFrameNode::predicate)
.def_ro("no_realize", &SBlockFrameNode::no_realize);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.SSBlockFrame", SBlockFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to SBlockFrameNode.
*
* \sa SBlockFrameNode
*/
class SBlockFrame : public TIRFrame {
public:
explicit SBlockFrame(ffi::ObjectPtr<SBlockFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockFrame, TIRFrame, SBlockFrameNode);
};
/*!
* \brief A frame that represents the block initialization statment.
*
* \sa BlockInitFrame
*/
class BlockInitFrameNode : public TIRFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BlockInitFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.SBlockInitFrame", BlockInitFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to BlockInitFrameNode.
*
* \sa BlockInitFrameNode
*/
class BlockInitFrame : public TIRFrame {
public:
explicit BlockInitFrame(ffi::ObjectPtr<BlockInitFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockInitFrame, TIRFrame, BlockInitFrameNode);
};
/*!
* \brief A frame that represents the for loop.
*
* \sa ForFrame
*/
class ForFrameNode : public TIRFrameNode {
public:
/*!
* \brief Functions that generate loop nests.
* \param loop_vars The loop variables, from outer to inner
* \param loop_extents The loop extents that correspond to loop variables
* \param loop_body The loop body
* \return A stmt, the loop nest
*/
using FMakeForLoop = ffi::TypedFunction<tvm::tirx::Stmt(
ffi::Array<tvm::tirx::Var> loop_vars, ffi::Array<Range> loop_extents,
ffi::Array<ffi::Optional<PrimExpr>> loop_steps, tvm::tirx::Stmt loop_body)>;
/*! \brief The loop variable. */
ffi::Array<tvm::tirx::Var> vars;
/*! \brief The domains of iteration. */
ffi::Array<Range> doms;
/*! \brief The optional steps of iteration. */
ffi::Array<ffi::Optional<PrimExpr>> steps;
/*! \brief The for loop generating function. */
FMakeForLoop f_make_for_loop;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ForFrameNode>()
.def_ro("vars", &ForFrameNode::vars)
.def_ro("doms", &ForFrameNode::doms);
// `f_make_for_loop` is not registered as it's not visited.
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ForFrame", ForFrameNode, TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ForFrameNode.
*
* \sa ForFrameNode
*/
class ForFrame : public TIRFrame {
public:
explicit ForFrame(ffi::ObjectPtr<ForFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ForFrame, TIRFrame, ForFrameNode);
};
/*!
* \brief A frame that represents the assert statement. Proceeds if the condition is true,
* otherwise aborts with the message.
*
* \sa AssertFrame
*/
class AssertFrameNode : public TIRFrameNode {
public:
/*! \brief The PrimExpr to test. */
PrimExpr condition;
/*! \brief The error kind, e.g. "RuntimeError", "TypeError", "ValueError". */
tvm::tirx::StringImm error_kind;
/*! \brief Error message fragments, concatenated at runtime when assertion fails. */
ffi::Array<tvm::tirx::StringImm> message_parts;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AssertFrameNode>()
.def_ro("condition", &AssertFrameNode::condition)
.def_ro("error_kind", &AssertFrameNode::error_kind)
.def_ro("message_parts", &AssertFrameNode::message_parts);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AssertFrame", AssertFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to AssertFrameNode.
*
* \sa AssertFrameNode
*/
class AssertFrame : public TIRFrame {
public:
explicit AssertFrame(ffi::ObjectPtr<AssertFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssertFrame, TIRFrame, AssertFrameNode);
};
/*!
* \brief The LaunchThreadFrameNode.
* \note It is used only inside a PrimFunc.
*/
class LaunchThreadFrameNode : public TIRFrameNode {
public:
/*! \brief The extent of environment thread. */
PrimExpr extent;
/*! \brief The attribute key, could be either virtual_thread or thread_extent. */
ffi::String attr_key;
/*! \brief The iteration variable. */
tvm::tirx::IterVar iter_var;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<LaunchThreadFrameNode>()
.def_ro("extent", &LaunchThreadFrameNode::extent)
.def_ro("attr_key", &LaunchThreadFrameNode::attr_key)
.def_ro("iter_var", &LaunchThreadFrameNode::iter_var);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.LaunchThreadFrame",
LaunchThreadFrameNode, TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to LaunchThreadFrameNode.
*
* \sa LaunchThreadFrameNode
*/
class LaunchThreadFrame : public TIRFrame {
public:
explicit LaunchThreadFrame(ffi::ObjectPtr<LaunchThreadFrameNode> data)
: TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LaunchThreadFrame, TIRFrame, LaunchThreadFrameNode);
};
/*!
* \brief A frame that represents attribute node.
*
* \sa AttrFrame
*/
class AttrFrameNode : public TIRFrameNode {
public:
/*! \brief The node to annotate the attribute. */
Any node;
/*! \brief Attribute type key. */
ffi::String attr_key;
/*! \brief The value of the attribute. */
PrimExpr value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AttrFrameNode>()
.def_ro("node", &AttrFrameNode::node)
.def_ro("attr_key", &AttrFrameNode::attr_key)
.def_ro("value", &AttrFrameNode::value);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AttrFrame", AttrFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to AttrFrameNode.
*
* \sa AttrFrameNode
*/
class AttrFrame : public TIRFrame {
public:
explicit AttrFrame(ffi::ObjectPtr<AttrFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrFrame, TIRFrame, AttrFrameNode);
};
/*!
* \brief A frame that represents while loop.
*
* \sa WhileFrame
*/
class WhileFrameNode : public TIRFrameNode {
public:
/*! \brief The termination condition of while. */
PrimExpr condition;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<WhileFrameNode>().def_ro("condition", &WhileFrameNode::condition);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.WhileFrame", WhileFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to WhileFrameNode.
*
* \sa WhileFrameNode
*/
class WhileFrame : public TIRFrame {
public:
explicit WhileFrame(ffi::ObjectPtr<WhileFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WhileFrame, TIRFrame, WhileFrameNode);
};
/*!
* \brief A frame that represents if statement.
*
* \sa IfFrame
*/
class IfFrameNode : public TIRFrameNode {
public:
/*! \brief The condition of the if statement. */
PrimExpr condition;
/*! \brief The statements in the true branch. */
ffi::Optional<ffi::Array<tvm::tirx::Stmt>> then_stmts;
/*! \brief The stetements in the false branch. */
ffi::Optional<ffi::Array<tvm::tirx::Stmt>> else_stmts;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IfFrameNode>()
.def_ro("condition", &IfFrameNode::condition)
.def_ro("then_stmts", &IfFrameNode::then_stmts)
.def_ro("else_stmts", &IfFrameNode::else_stmts);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.IfFrame", IfFrameNode, TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to IfFrameNode.
*
* \sa IfFrameNode
*/
class IfFrame : public TIRFrame {
public:
explicit IfFrame(ffi::ObjectPtr<IfFrameNode> data) : TIRFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfFrame, TIRFrame, IfFrameNode);
};
/*!
* \brief A frame that represents then.
*
* \sa ThenFrame
*/
class ThenFrameNode : public TIRFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ThenFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ThenFrame", ThenFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ThenFrameNode.
*
* \sa ThenFrameNode
*/
class ThenFrame : public TIRFrame {
public:
explicit ThenFrame(ffi::ObjectPtr<ThenFrameNode> data) : TIRFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThenFrame, TIRFrame, ThenFrameNode);
};
/*!
* \brief A frame that represents else.
*
* \sa ElseFrame
*/
class ElseFrameNode : public TIRFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ElseFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ElseFrame", ElseFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ElseFrameNode.
*
* \sa ElseFrameNode
*/
class ElseFrame : public TIRFrame {
public:
explicit ElseFrame(ffi::ObjectPtr<ElseFrameNode> data) : TIRFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, TIRFrame, ElseFrameNode);
};
class DeclBufferFrameNode : public TIRFrameNode {
public:
/*! \brief The declared buffer. */
tvm::tirx::Buffer buffer;
/*! \brief The buffer allocated or not. */
bool allocated;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DeclBufferFrameNode>()
.def_ro("buffer", &DeclBufferFrameNode::buffer)
.def_ro("allocated", &DeclBufferFrameNode::allocated);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.DeclBufferFrame", DeclBufferFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
class DeclBufferFrame : public TIRFrame {
public:
explicit DeclBufferFrame(ffi::ObjectPtr<DeclBufferFrameNode> data) : TIRFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeclBufferFrame, TIRFrame, DeclBufferFrameNode);
};
class ComposeOpFrameNode : public TIRFrameNode {
public:
/*! \brief The workspace of the compose op. */
ffi::Map<ffi::String, tvm::tirx::Buffer> workspace;
/*! \brief The config of the compose op. */
ffi::Map<ffi::String, ffi::Any> config;
/*! \brief The optional dispatch variant name of the compose op. */
ffi::Optional<ffi::String> dispatch{std::nullopt};
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ComposeOpFrameNode>()
.def_ro("workspace", &ComposeOpFrameNode::workspace)
.def_ro("config", &ComposeOpFrameNode::config)
.def_ro("dispatch", &ComposeOpFrameNode::dispatch);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ComposeOpFrame", ComposeOpFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
class ComposeOpFrame : public TIRFrame {
public:
explicit ComposeOpFrame(ffi::ObjectPtr<ComposeOpFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComposeOpFrame, TIRFrame, ComposeOpFrameNode);
};
class AllocBufferFrameNode : public TIRFrameNode {
public:
/*! \brief The allocated buffer. */
tvm::tirx::Buffer buffer;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AllocBufferFrameNode>().def_ro("buffer", &AllocBufferFrameNode::buffer);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AllocBufferFrame", AllocBufferFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
class AllocBufferFrame : public TIRFrame {
public:
explicit AllocBufferFrame(ffi::ObjectPtr<AllocBufferFrameNode> data)
: TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AllocBufferFrame, TIRFrame, AllocBufferFrameNode);
};
/*!
* \brief A frame that represents a hint directive for the sketch language.
*
* \sa HintFrame
*/
class HintFrameNode : public TIRFrameNode {
public:
/*! \brief The free-form hint message string. */
ffi::String message;
/*! \brief Optional structured key-value attributes. */
ffi::Map<ffi::String, ffi::Any> attrs;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<HintFrameNode>()
.def_ro("message", &HintFrameNode::message)
.def_ro("attrs", &HintFrameNode::attrs);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.HintFrame", HintFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
/*!
* \brief Managed reference to HintFrameNode.
*
* \sa HintFrameNode
*/
class HintFrame : public TIRFrame {
public:
explicit HintFrame(ffi::ObjectPtr<HintFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(HintFrame, TIRFrame, HintFrameNode);
};
} // namespace tirx
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_TIRX_SCRIPT_BUILDER_FRAME_H_
include/tvm/tirx/script/builder/ir.h
+576
View File
@@ -0,0 +1,576 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
#ifndef TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/container/variant.h>
#include <tvm/runtime/tensor.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/script/builder/frame.h>
#include <tvm/tirx/tirx_stmt.h>
namespace tvm {
namespace script {
namespace ir_builder {
namespace tirx {
using tvm::ffi::Tuple;
using tvm::ffi::Variant;
using tvm::runtime::Tensor;
using tvm::tirx::Buffer;
using tvm::tirx::ExecScope;
using tvm::tirx::Layout;
using tvm::tirx::Var;
/*!
* \brief The buffer declaration function.
* \param shape The type of the buffer prior to flattening.
* \param dtype The data type in the content of the buffer.
* \param buffer_name The name of the buffer.
* \param data The pointer to the head of the data.
* \param strides The strides of each dimension.
* \param elem_offset The offset in terms of number of dtype elements (including lanes).
* \param storage_scope The optional storage scope of buffer data pointer.
* \param align The alignment requirement of data pointer in bytes.
* \param offset_factor The factor of elem_offset field.
* \param buffer_type The buffer type.
* \param axis_separators The separators between input axes when generating flattened output axes.
* \return The declared buffer.
*/
Buffer BufferDecl(ffi::Array<PrimExpr> shape, PrimType 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,
ffi::Optional<Layout> layout = std::nullopt,
ffi::Array<PrimExpr> allocated_addr = {});
/*!
* \brief The primitive function statement.
* \return The PrimFuncFrame.
*/
PrimFuncFrame PrimFunc(bool is_private, bool s_tir = false, bool persistent = false);
/*!
* \brief The PrimFunc variable arguments adding function.
* \param name The name of the variable.
* \param var The variable argument.
* \return The variable.
*/
Var Arg(ffi::String name, Var var);
/*!
* \brief The PrimFunc buffer arguments adding function.
* \param name The name of the buffer.
* \param buffer The buffer argument.
* \return The buffer.
*/
Buffer Arg(ffi::String name, Buffer buffer);
/*!
* \brief The PrimFunc naming statement.
* \param name The name of the PrimFunc.
*/
void FuncName(ffi::String name);
/*!
* \brief The PrimFunc annotation statement.
* \param attrs The annotations of the PrimFunc.
*/
void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
/*!
* \brief The PrimFunc return type statement.
* \param ret_type The return type of the PrimFunc.
* \return The return type.
*/
Type FuncRet(Type ret_type);
/*!
* \brief The buffer match statement.
* \param param The parameter of the PrimFunc to match.
* \param shape The type of the buffer prior to flattening.
* \param dtype The data type in the content of the buffer.
* \param data The pointer to the head of the data.
* \param strides The strides of each dimension.
* \param elem_offset The offset in terms of number of dtype elements (including lanes).
* \param storage_scope The optional storage scope of buffer data pointer.
* \param align The alignment requirement of data pointer in bytes.
* \param offset_factor The factor of elem_offset field.
* \param buffer_type The buffer type.
* \param axis_separators The separators between input axes when generating flattened output axes.
* \return The matched buffer.
*/
Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
PrimType dtype = PrimType::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,
ffi::Optional<Layout> layout = std::nullopt);
/*!
* \brief The block declaration statement.
* \param name The name of the block.
* \param no_realize The flag whether to construct SBlockRealize or SBlock.
* \return The SBlockFrame.
*/
SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_scope = "");
void TilePrimitiveCall(tvm::tirx::TilePrimitiveCall op_call);
ffi::Array<tvm::tirx::Var> KernelId(ffi::Array<PrimExpr> extents, ffi::String parent);
ffi::Array<tvm::tirx::Var> CtaId(ffi::Array<PrimExpr> extents, ffi::String parent);
ffi::Array<tvm::tirx::Var> CtaIdInPair();
ffi::Array<tvm::tirx::Var> WarpId(ffi::Array<PrimExpr> extents, ffi::String parent);
ffi::Array<tvm::tirx::Var> ThreadId(ffi::Array<PrimExpr> extents, ffi::String parent);
/*!
* \brief The block initialization statement.
* \return The BlockInitFrame.
*/
BlockInitFrame Init();
/*!
* \brief The block predicate statement.
* \param predicate The predicate condition.
*/
void Where(PrimExpr predicate);
/*!
* \brief The block buffer region reading statement.
* \param buffer_slices The array of buffer regions to read.
*/
void Reads(ffi::Array<ffi::ObjectRef> buffer_slices);
/*!
* \brief The block buffer region writing statement.
* \param buffer_slices The array of buffer regions to write.
*/
void Writes(ffi::Array<ffi::ObjectRef> buffer_slices);
/*!
* \brief The block annotation statement.
* \param attrs The annotation of the block.
*/
void BlockAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
/*!
* \brief The buffer allocation function.
* \param shape The type of the buffer prior to flattening.
* \param dtype The data type in the content of the buffer.
* \param data The pointer to the head of the data.
* \param strides The strides of each dimension.
* \param elem_offset The offset in terms of number of dtype elements (including lanes).
* \param storage_scope The optional storage scope of buffer data pointer.
* \param align The alignment requirement of data pointer in bytes.
* \param offset_factor The factor of elem_offset field.
* \param buffer_type The buffer type.
* \param axis_separators The separators between input axes when generating flattened output axes.
* \param layout The layout of the buffer.
* \param allocated_addr The allocated address of the buffer. Might be multi-dimensional.
* \return The allocated buffer or the AllocBufferFrame if the function is called under
* T.prim_func(tirx=True).
*/
ffi::Variant<Buffer, AllocBufferFrame> SBlockAllocBuffer(
ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::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,
ffi::Optional<Layout> layout = std::nullopt, ffi::Array<PrimExpr> allocated_addr = {});
namespace axis {
/*!
* \brief The spatial block axis defining function.
* \param dom The domain of the iteration variable.
* \param binding The binding value of the iteration variable.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
Var Spatial(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
/*!
* \brief The reduced block axis defining function.
* \param dom The domain of the iteration variable.
* \param binding The binding value of the iteration variable.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
Var Reduce(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
/*!
* \brief The scanning block axis defining function.
* \param dom The domain of the iteration variable.
* \param binding The binding value of the iteration variable.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
Var Scan(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
/*!
* \brief The opaque block axis defining function.
* \param dom The domain of the iteration variable.
* \param binding The binding value of the iteration variable.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
Var Opaque(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
/*!
* \brief The block axis remapping function.
* \param kinds The types of the iteration variables.
* \param bindings The binding values of the iteration variables.
* \param dtype The data types of the iteration variables.
* \return The iteration variables.
*/
ffi::Array<Var> Remap(ffi::String kinds, ffi::Array<PrimExpr> bindings,
PrimType dtype = PrimType::Int(32));
} // namespace axis
/*!
* \brief The serial For statement.
* \param start The minimum value of iteration.
* \param stop The maximum value of iteration.
* \param annotations The optional annotations of the For statement.
* \param step The optional step value of iteration.
* \return The ForFrame.
*/
ForFrame Serial(PrimExpr start, PrimExpr stop,
ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
ffi::Optional<PrimExpr> step = std::nullopt);
/*!
* \brief The parallel For statement.
* \param start The minimum value of iteration.
* \param stop The maximum value of iteration.
* \param annotations The optional annotations of the For statement.
* \param step The optional step value of iteration.
* \return The ForFrame.
*/
ForFrame Parallel(PrimExpr start, PrimExpr stop,
ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
ffi::Optional<PrimExpr> step = std::nullopt);
/*!
* \brief The vectorized For statement.
* \param start The minimum value of iteration.
* \param stop The maximum value of iteration.
* \param annotations The optional annotations of the For statement.
* \param step The optional step value of iteration.
* \return The ForFrame.
*/
ForFrame Vectorized(PrimExpr start, PrimExpr stop,
ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
ffi::Optional<PrimExpr> step = std::nullopt);
/*!
* \brief The unrolled For statement.
* \param start The minimum value of iteration.
* \param stop The maximum value of iteration.
* \param annotations The optional annotations of the For statement.
* \param step The optional step value of iteration.
* \return The ForFrame.
*/
ForFrame Unroll(PrimExpr start, PrimExpr stop,
ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
ffi::Optional<PrimExpr> step = std::nullopt);
/*!
* \brief The thread-binding For statement.
* \param start The minimum value of iteration.
* \param stop The maximum value of iteration.
* \param thread The thread for loop variable to bind.
* \param annotations The optional annotations of the For statement.
* \return The ForFrame.
*/
ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread,
ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
/*!
* \brief The grid For statement.
* \param extents The extents of the iteration.
* \return The ForFrame.
*/
ForFrame Grid(ffi::Array<Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>> extents);
/*!
* \brief The assertion statement.
* \param condition The assertion condition.
* \param error_kind The error kind (e.g. "RuntimeError", "TypeError", "ValueError").
* \param message_parts The error message parts (stored as separate fragments in the IR).
* \return The AssertFrame.
*/
AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
ffi::Array<ffi::String> message_parts);
/*!
* \brief Create a Bind (variable binding).
*
* Emits a flat Bind statement to the current frame and returns the bound variable.
*
* \param value The value to be bound.
* \param type_annotation The type annotation of the binding.
* Usually it is used for fine-grained var typing,
* particularly, PointerType.
* \param var The variable to be bound. If not specified, a new variable will be created.
* \return The bound Var.
*/
Var Bind(Expr value, ffi::Optional<Type> type_annotation = std::nullopt,
ffi::Optional<Var> var = std::nullopt);
/*!
* \brief Create an attribute.
* \param node The node to annotate the attribute.
* \param attr_key Attribute type key.
* \param value The value of the attribute.
* \return The result AttrFrame.
*/
AttrFrame Attr(ffi::Any node, ffi::String attr_key, PrimExpr value);
/*!
* \brief Mark the device-region entry within the enclosing PrimFunc body.
* Returns an AttrFrame keyed ``tirx.device_entry`` (value ``Bool(true)``).
* Subsequent stmts accumulate into the frame's body; the frame is closed
* by ``PrimFuncFrameNode::ExitWithScope`` which drains leftover frames.
*
* Python sugar: ``Tx.device_entry()`` is a flat-call (no ``with``), which
* auto-enters the frame.
*/
AttrFrame DeviceEntry();
/*!
* \brief Create a while loop.
* \param condition The termination condition of the loop.
* \return The result WhileFrame.
*/
WhileFrame While(PrimExpr condition);
/*!
* \brief Create a break statement.
*/
void Break();
/*!
* \brief Create a continue statement.
*/
void Continue();
/*!
* \brief Create an if statement.
* \param condition The condition of if statement.
* \return The result IfFrame.
*/
IfFrame If(PrimExpr condition);
/*!
* \brief Create a then.
* \return The result ThenFrame.
*/
ThenFrame Then();
/*!
* \brief Create an else.
* \return The result ElseFrame.
*/
ElseFrame Else();
/*!
* \brief The buffer declaration frame.
* \param shape The type of the buffer prior to flattening.
* \param dtype The data type in the content of the buffer.
* \param buffer_name The name of the buffer.
* \param data The pointer to the head of the data.
* \param strides The strides of each dimension.
* \param elem_offset The offset in terms of number of dtype elements (including lanes).
* \param storage_scope The optional storage scope of buffer data pointer.
* \param align The alignment requirement of data pointer in bytes.
* \param offset_factor The factor of elem_offset field.
* \param buffer_type The buffer type.
* \param axis_separators The separators between input axes when generating flattened output axes.
* \param layout The layout of the buffer.
* \return The declaration frame.
*/
DeclBufferFrame DeclBuffer(ffi::Array<PrimExpr> shape, PrimType 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,
ffi::Optional<Layout> layout = std::nullopt,
ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
/*!
* \brief Statement-level buffer allocation (creates an AllocBuffer IR node).
* \param shape The shape of the buffer to allocate.
* \param dtype The data type of buffer elements.
* \param storage_scope The storage scope (e.g., "global", "shared").
* \param annotations Optional annotations for the allocation.
* \return The allocated buffer.
*/
Buffer AllocBuffer(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
ffi::String storage_scope = "global",
ffi::Optional<ffi::Map<ffi::String, ffi::Any>> annotations = std::nullopt);
/*!
* \brief Launch a thread.
* \param var The iteration variable.
* \param extent The extent of environment thread.
* \return The result LaunchThreadFrame.
*/
LaunchThreadFrame LaunchThread(Var var, PrimExpr extent);
/*!
* \brief Launch a new thread.
* \param thread_tag The thread type tag.
* \param extent The extent of environment thread.
* \return The result LaunchThreadFrame.
*/
LaunchThreadFrame LaunchThread(ffi::String thread_tag, PrimExpr extent);
/*!
* \brief Compose TIRx op.
* \param workspace The workspace of the compose op.
* \param config The config of the compose op.
* \param dispatch The optional dispatch variant name.
* \return The result ComposeOpFrame.
*/
ComposeOpFrame ComposeOp(ffi::Map<ffi::String, Buffer> workspace,
ffi::Map<ffi::String, ffi::Any> config,
ffi::Optional<ffi::String> dispatch = std::nullopt);
/*!
* \brief Bind a var to thread env.
* \param thread_tag The thread type tag.
* \param dtype The data type of the variable.
* \return The result variable which gets bound to the thread env.
*/
Var EnvThread(ffi::String thread_tag, PrimType dtype = PrimType::Int(32));
/*!
* \brief Store data in a buffer.
* \param buffer The buffer.
* \param value The value to be stored.
* \param indices The indices location to be stored.
* \param predicate A vector mask of boolean values indicating which lanes of a vector are to be
* stored. The number lanes of the mask must be equal to the number of lanes in value.
*/
void BufferStore(Buffer buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
ffi::Optional<PrimExpr> predicate);
/*!
* \brief Evaluate the input expression.
* \param value The input expression to evaluate.
*/
void Evaluate(Expr value);
/*!
* \brief Create a TIR var that represents a pointer
*
* \param dtype The optional data type of the pointer. If omitted, construct
* an opaque handle.
*
* \param storage_scope The storage scope of the pointer.
*
* \return The pointer.
*/
inline Var Handle(ffi::Optional<PrimType> dtype = std::nullopt,
ffi::String storage_scope = "global") {
Type type_annotation = dtype.has_value() ? Type(PointerType(dtype.value(), storage_scope))
: Type(PointerType::VoidPointerTy(storage_scope));
return tvm::tirx::Var("", type_annotation);
}
inline Var TensorMap() { return tvm::tirx::Var("", PointerType(TensorMapType())); }
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType) \
inline PrimExpr FuncName(ffi::Optional<PrimExpr> expr = std::nullopt) { \
PrimType dtype = DType; \
return expr.has_value() ? tvm::cast(dtype, expr.value()) \
: tvm::tirx::Var("", dtype).as_or_throw<PrimExpr>(); \
}
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, Code) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##8, (PrimType(DLDataType{Code, 8, 1}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##16, (PrimType(DLDataType{Code, 16, 1}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##32, (PrimType(DLDataType{Code, 32, 1}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##64, (PrimType(DLDataType{Code, 64, 1})));
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(BFloat, kDLBfloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Float, kDLFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(UInt, kDLUInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, kDLInt);
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, Code, Size) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x2, (PrimType(DLDataType{Code, Size, 2}))) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, (PrimType(DLDataType{Code, Size, 4}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, (PrimType(DLDataType{Code, Size, 8}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, (PrimType(DLDataType{Code, Size, 16}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, (PrimType(DLDataType{Code, Size, 32}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, (PrimType(DLDataType{Code, Size, 64})));
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, Code) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, Code, 8); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, Code, 16); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, Code, 32); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, Code, 64);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(BFloat, kDLBfloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Float, kDLFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(UInt, kDLUInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, kDLInt);
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, Code, Bits) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType, (PrimType(DLDataType{Code, Bits, 1}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x2, (PrimType(DLDataType{Code, Bits, 2}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, (PrimType(DLDataType{Code, Bits, 4}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, (PrimType(DLDataType{Code, Bits, 8}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, (PrimType(DLDataType{Code, Bits, 16}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, (PrimType(DLDataType{Code, Bits, 32}))); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, (PrimType(DLDataType{Code, Bits, 64})));
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E3M4, kDLFloat8_e3m4, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3, kDLFloat8_e4m3, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3B11FNUZ, kDLFloat8_e4m3b11fnuz, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FN, kDLFloat8_e4m3fn, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FNUZ, kDLFloat8_e4m3fnuz, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2, kDLFloat8_e5m2, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2FNUZ, kDLFloat8_e5m2fnuz, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E8M0FNU, kDLFloat8_e8m0fnu, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E2M3FN, kDLFloat6_e2m3fn, 6);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E3M2FN, kDLFloat6_e3m2fn, 6);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float4E2M1FN, kDLFloat4_e2m1fn, 4);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Boolean, PrimType::Bool());
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Void, PrimType::Void());
#undef TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST
} // namespace tirx
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_TIRX_SCRIPT_BUILDER_IR_H_