/*
 * 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_