wehub-resource-sync
577 lines
23 KiB
C++
577 lines
23 KiB
C++
/*
|
|
* 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_
|