wehub-resource-sync
3292 lines
72 KiB
Python
3292 lines
72 KiB
Python
# 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.
|
|
# pylint: disable=redefined-builtin, invalid-name, too-many-arguments
|
|
"""Operators used in TIR expression."""
|
|
|
|
from typing import Any
|
|
|
|
import tvm_ffi
|
|
from tvm_ffi import Array
|
|
|
|
import tvm
|
|
from tvm import tirx
|
|
from tvm.ir import Call, Expr, Op, PointerType
|
|
from tvm.ir.base import Span
|
|
from tvm.ir.type import TensorMapType
|
|
from tvm.runtime import const
|
|
|
|
from . import _ffi_api
|
|
from .buffer import Buffer
|
|
from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var
|
|
|
|
tir = tirx # alias for backward compat with upstream tir.convert() calls
|
|
|
|
_DEVICE_INTRIN_PREFIX_TO_NAMESPACE = {
|
|
"cuda_": "cuda",
|
|
"ptx_": "ptx",
|
|
"nvshmem_": "nvshmem",
|
|
"nki_": "nki",
|
|
}
|
|
|
|
|
|
def _canonical_device_intrin_name(func_name: str) -> str:
|
|
"""Return the canonical registry name for statically registered device intrinsics."""
|
|
|
|
if not isinstance(func_name, str) or not func_name.startswith("tirx."):
|
|
return func_name
|
|
basename = func_name[len("tirx.") :]
|
|
if "." in basename:
|
|
return func_name
|
|
for prefix, namespace in _DEVICE_INTRIN_PREFIX_TO_NAMESPACE.items():
|
|
if basename.startswith(prefix):
|
|
return f"tirx.{namespace}.{basename[len(prefix) :]}"
|
|
return func_name
|
|
|
|
|
|
def _primexpr_ty(expr):
|
|
"""Return the runtime primitive type of an expression."""
|
|
ty = getattr(expr, "ty", None)
|
|
if isinstance(ty, tvm.ir.PrimType):
|
|
return ty
|
|
if isinstance(expr, ExprOp):
|
|
return expr.expr_ty()
|
|
raise TypeError(f"Cannot determine Expr type for {type(expr).__name__}")
|
|
|
|
|
|
def _primexpr_dtype(expr):
|
|
"""Return the runtime dtype of a primitive expression without using Expr.dtype."""
|
|
ty = _primexpr_ty(expr)
|
|
if not isinstance(ty, tvm.ir.PrimType):
|
|
raise TypeError(f"Expected PrimType for {type(expr).__name__}, but got {ty}")
|
|
return ty.dtype
|
|
|
|
|
|
def _pack_buffer(buf, span=None):
|
|
"""Build intrinsics that packs the buffer."""
|
|
shape = Call(
|
|
"tirx.tvm_stack_make_shape",
|
|
buf.shape,
|
|
span=span,
|
|
ret_ty=PointerType(tvm.ir.PrimType("int64")),
|
|
)
|
|
strides = (
|
|
Call(
|
|
"tirx.tvm_stack_make_shape",
|
|
buf.strides,
|
|
span=span,
|
|
ret_ty=PointerType(tvm.ir.PrimType("int64")),
|
|
)
|
|
if buf.strides
|
|
else 0
|
|
)
|
|
pack_args = [
|
|
buf.data,
|
|
shape,
|
|
strides,
|
|
len(buf.shape),
|
|
const(0, dtype=buf.dtype),
|
|
buf.elem_offset,
|
|
]
|
|
return Call(Op.get("tirx.tvm_stack_make_array"), pack_args, span=span, ret_ty="handle")
|
|
|
|
|
|
def call_packed_lowered(*args, span=None):
|
|
"""Lowered version of call packed.
|
|
The argument to packed function can be Expr or Buffer.
|
|
The argument is the corresponding POD type when Expr is presented.
|
|
When the argument is Buffer, the corresponding PackedFunc
|
|
will receive an TVMArrayHandle whose content is valid during the callback period.
|
|
If the PackedFunc is a python callback, then the corresponding argument is Tensor.
|
|
|
|
Parameters
|
|
----------
|
|
args : list of Expr or Buffer.
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
|
|
See Also
|
|
--------
|
|
te.extern : Create tensor with extern function call.
|
|
"""
|
|
call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
|
|
return Call(Op.get("tirx.tvm_call_packed_lowered"), call_args, span=span, ret_ty="int32")
|
|
|
|
|
|
def call_cpacked_lowered(*args, span=None):
|
|
"""Lowered version of call c-packed.
|
|
Same as call_packed, except that the first argument is the function name
|
|
(as in call_extern), and the last argument is the resource handle.
|
|
|
|
Parameters
|
|
----------
|
|
args : list of Expr or Buffer.
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
|
|
See Also
|
|
--------
|
|
te.extern : Create tensor with extern function call.
|
|
"""
|
|
call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
|
|
return Call(Op.get("tirx.tvm_call_cpacked_lowered"), call_args, span=span, ret_ty="int32")
|
|
|
|
|
|
def call_packed(*args, span=None):
|
|
"""Build expression by call an external packed function.
|
|
|
|
The argument to packed function can be Expr or Buffer.
|
|
The argument is the corresponding POD type when Expr is presented.
|
|
|
|
When the argument is Buffer, the corresponding PackedFunc
|
|
will receive an TVMArrayHandle whose content is valid during the callback period.
|
|
If the PackedFunc is a python callback, then the corresponding argument is Tensor.
|
|
|
|
Parameters
|
|
----------
|
|
args : list of Expr or Buffer.
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
|
|
See Also
|
|
--------
|
|
te.extern : Create tensor with extern function call.
|
|
"""
|
|
call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
|
|
return Call(Op.get("tirx.tvm_call_packed"), call_args, span=span, ret_ty="int32")
|
|
|
|
|
|
def call_cpacked(*args, span=None):
|
|
"""Build expression by call an external packed function.
|
|
|
|
Same as call_packed, except that the first argument is the function name
|
|
(as in call_extern), and the last argument is the resource handle.
|
|
|
|
Parameters
|
|
----------
|
|
args : list of Expr or Buffer.
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
|
|
See Also
|
|
--------
|
|
te.extern : Create tensor with extern function call.
|
|
"""
|
|
call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
|
|
return Call(Op.get("tirx.tvm_call_cpacked"), call_args, span=span, ret_ty="int32")
|
|
|
|
|
|
def call_intrin(dtype: str | tvm.ir.Type, func_name, *args, attrs=None, span=None):
|
|
"""Build expression by calling an intrinsic function.
|
|
|
|
Intrinsics can be overloaded with multiple data types via
|
|
the intrinsic translation rule.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str or tvm.ir.Type
|
|
The data type of the result.
|
|
|
|
func_name: str
|
|
The intrinsic function name.
|
|
|
|
args : list
|
|
Positional arguments.
|
|
|
|
attrs : Optional[tvm.ir.Attrs or Dict[str, Object]]
|
|
Additional attributes for the call.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
if isinstance(func_name, str):
|
|
func_name = _canonical_device_intrin_name(func_name)
|
|
return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype)
|
|
|
|
|
|
def call_pure_extern(dtype, func_name, *args, span=None):
|
|
"""Build expression by calling a pure extern function.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
func_name: str
|
|
The extern function name.
|
|
|
|
args : list
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return Call(
|
|
Op.get("tirx.call_pure_extern"),
|
|
[func_name, *args],
|
|
span=span,
|
|
ret_ty=dtype,
|
|
)
|
|
|
|
|
|
def call_extern(dtype, func_name, *args, span=None):
|
|
"""Build expression by calling a extern function.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
func_name: str
|
|
The extern function name.
|
|
|
|
args : list
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return Call(
|
|
Op.get("tirx.call_extern"),
|
|
[func_name, *args],
|
|
span=span,
|
|
ret_ty=dtype,
|
|
)
|
|
|
|
|
|
def _require_float_arg(op_name, x):
|
|
x = tirx.convert(x)
|
|
dtype = _primexpr_dtype(x)
|
|
if "float" not in dtype and "bfloat" not in dtype:
|
|
raise TypeError(f"tirx.{op_name} only supports floating-point inputs, but got {dtype}")
|
|
return x
|
|
|
|
|
|
def call_llvm_intrin(dtype, name, *args, span=None):
|
|
"""Build expression by calling a llvm intrinsic function
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
name : str
|
|
The name of the llvm intrinsic function.
|
|
|
|
args : list
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
# pylint: disable=import-outside-toplevel
|
|
from tvm.target import codegen
|
|
|
|
if isinstance(name, str):
|
|
llvm_id = codegen.llvm_lookup_intrinsic_id(name)
|
|
elif isinstance(name, IntImm):
|
|
llvm_id = name.value
|
|
else:
|
|
llvm_id = name
|
|
if llvm_id == 0:
|
|
raise ValueError(f"Unknown llvm intrinsic function {name}")
|
|
return call_intrin(
|
|
dtype,
|
|
Op.get("tirx.call_llvm_intrin"),
|
|
tvm.tirx.const(llvm_id, "uint32"),
|
|
*args,
|
|
span=span,
|
|
)
|
|
|
|
|
|
def call_llvm_pure_intrin(dtype, name, *args, span=None):
|
|
"""Build expression by calling a pure llvm intrinsic function
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
name : str
|
|
The name of the llvm intrinsic function.
|
|
|
|
args : list
|
|
Positional arguments.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
# pylint: disable=import-outside-toplevel
|
|
from tvm.target import codegen
|
|
|
|
if isinstance(name, str):
|
|
llvm_id = codegen.llvm_lookup_intrinsic_id(name)
|
|
elif isinstance(name, IntImm):
|
|
llvm_id = name.value
|
|
else:
|
|
llvm_id = name
|
|
if llvm_id == 0:
|
|
raise ValueError(f"Unknown llvm intrinsic function {name}")
|
|
return call_intrin(
|
|
dtype,
|
|
Op.get("tirx.call_llvm_pure_intrin"),
|
|
tvm.tirx.const(llvm_id, "uint32"),
|
|
*args,
|
|
span=span,
|
|
)
|
|
|
|
|
|
def tvm_stack_alloca(dtype_str, num):
|
|
"""Return new on stack dtype[num]
|
|
|
|
Parameters
|
|
----------
|
|
dtype_str : str
|
|
The data type of array.
|
|
|
|
num : int
|
|
The size of array.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
if dtype_str == "shape":
|
|
ret_ty = PointerType(tvm.ir.PrimType("int64"))
|
|
elif dtype_str == "arg_tcode":
|
|
ret_ty = PointerType(tvm.ir.PrimType("int32"))
|
|
elif dtype_str == "tensormap":
|
|
ret_ty = PointerType(TensorMapType())
|
|
else:
|
|
ret_ty = PointerType(tvm.ir.PrimType("void"))
|
|
return call_intrin(ret_ty, "tirx.tvm_stack_alloca", dtype_str, num)
|
|
|
|
|
|
def tvm_stack_make_shape(*args):
|
|
"""Allocate a shape tuple on stack, return the handle
|
|
|
|
Parameters
|
|
----------
|
|
args : int
|
|
The tuple shape.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(PointerType(tvm.ir.PrimType("int64")), "tirx.tvm_stack_make_shape", *args)
|
|
|
|
|
|
def tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset):
|
|
"""Allocate a Tensor(DLTensor) on stack, return the handle
|
|
|
|
Parameters
|
|
----------
|
|
data : Expr
|
|
The data of array.
|
|
|
|
shape : Expr
|
|
The shape of array.
|
|
|
|
strides : Expr
|
|
The strides of array.
|
|
|
|
ndim : Expr
|
|
The dimensions of array.
|
|
|
|
arr_dtype : Expr
|
|
The data type of array.
|
|
|
|
elem_offse : Expr
|
|
The element offset of array.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
if isinstance(arr_dtype, str | tvm.DataType | tvm.ir.PrimType):
|
|
arr_dtype = const(0, dtype=arr_dtype)
|
|
|
|
return call_intrin(
|
|
"handle",
|
|
"tirx.tvm_stack_make_array",
|
|
data,
|
|
shape,
|
|
strides,
|
|
ndim,
|
|
arr_dtype,
|
|
elem_offset,
|
|
)
|
|
|
|
|
|
def assume(cond=None):
|
|
"""Provide a true statement that can be used for simplifications
|
|
|
|
Parameters
|
|
----------
|
|
cond : Expr
|
|
The constraint condition.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("bool", "tirx.assume", cond)
|
|
|
|
|
|
def undef():
|
|
"""Returns an initialized but arbitrary value
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("int32", "tirx.undef")
|
|
|
|
|
|
def call_tir(global_var: tvm.ir.GlobalVar, *args):
|
|
"""Performs a call into another PrimFunc in the same IRModule
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
assert isinstance(global_var, tvm.ir.GlobalVar)
|
|
|
|
dtype = "void"
|
|
if global_var.ty is not None:
|
|
ret_ty = global_var.ty.ret
|
|
if isinstance(ret_ty, tvm.ir.PrimType):
|
|
dtype = ret_ty
|
|
|
|
return Call(op=global_var, args=args, ret_ty=dtype)
|
|
|
|
|
|
def start_profile_intrinsic(id):
|
|
"""Start profile intrinsic.
|
|
Parameters
|
|
----------
|
|
id : int
|
|
The intrinsic id.
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.start_profile_intrinsic", id)
|
|
|
|
|
|
def end_profile_intrinsic(id):
|
|
"""End profile intrinsic.
|
|
Parameters
|
|
----------
|
|
id : int
|
|
The intrinsic id.
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.end_profile_intrinsic", id)
|
|
|
|
|
|
def tvm_tuple(*value):
|
|
"""Create a tuple structure in value field of AttrStmt
|
|
|
|
Parameters
|
|
----------
|
|
value : Expr
|
|
The value in tuple.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.tvm_tuple", *value)
|
|
|
|
|
|
def handle_add_byte_offset(handle, offset):
|
|
"""Add offset to handle
|
|
|
|
Parameters
|
|
----------
|
|
handle : Expr
|
|
The handle.
|
|
|
|
offset : int
|
|
The offset.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
handle_type = getattr(handle, "ty", None)
|
|
storage_scope = handle_type.storage_scope if isinstance(handle_type, PointerType) else "global"
|
|
return call_intrin(
|
|
PointerType(tvm.ir.PrimType("void"), storage_scope),
|
|
"tirx.handle_add_byte_offset",
|
|
handle,
|
|
offset,
|
|
)
|
|
|
|
|
|
def tvm_struct_get(arr, index, field, dtype):
|
|
"""Get struct field value in array
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The date type of the result.
|
|
|
|
arr : StructType*
|
|
The array of struct.
|
|
|
|
index : int
|
|
The index of struct.
|
|
|
|
field : int
|
|
The field of struct.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(dtype, "tirx.tvm_struct_get", arr, index, field)
|
|
|
|
|
|
def tvm_struct_set(arr, index, field, value):
|
|
"""Set value in struct field in array
|
|
|
|
Parameters
|
|
----------
|
|
arr : StructType*
|
|
The array of struct.
|
|
|
|
index : int
|
|
The index of struct.
|
|
|
|
field : int
|
|
The field of struct.
|
|
|
|
value : Expr
|
|
The value to be set in field.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("int32", "tirx.tvm_struct_set", arr, index, field, value)
|
|
|
|
|
|
def _is_tensormap_var(obj: Var) -> bool:
|
|
return isinstance(obj.ty, PointerType) and isinstance(obj.ty.element_type, TensorMapType)
|
|
|
|
|
|
def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr:
|
|
"""Returns the address of a buffer element or addressable variable.
|
|
|
|
Parameters
|
|
----------
|
|
obj: Union[Buffer, BufferLoad, Var]
|
|
The buffer, buffer load, or addressable variable.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
if isinstance(obj, Buffer):
|
|
n_dim = len(obj.shape)
|
|
buffer_load = BufferLoad(obj, [0] * n_dim)
|
|
return Call("tirx.address_of", [buffer_load], span=span, ret_ty=obj.data.ty)
|
|
elif isinstance(obj, Var):
|
|
if _is_tensormap_var(obj):
|
|
return call_intrin("uint64", "tirx.address_of", obj, span=span)
|
|
if not isinstance(obj.ty, tvm.ir.PrimType):
|
|
raise TypeError(f"address_of expects a scalar or TensorMap Var, but got {obj.ty}")
|
|
return Call("tirx.address_of", [obj], span=span, ret_ty=PointerType(obj.ty))
|
|
elif isinstance(obj, BufferLoad):
|
|
return Call("tirx.address_of", [obj], span=span, ret_ty=obj.buffer.data.ty)
|
|
else:
|
|
raise ValueError(f"Invalid object type: {type(obj)}")
|
|
|
|
|
|
def lookup_param(param_name, span=None):
|
|
"""Returns the param by name
|
|
|
|
Parameters
|
|
----------
|
|
param_name : str
|
|
The name of param.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("handle", "tirx.lookup_param", param_name, span=span)
|
|
|
|
|
|
def tvm_thread_allreduce(*freduce_args):
|
|
"""Perform allreduce inside threadblock.
|
|
|
|
Parameters
|
|
----------
|
|
freduce_args : Expr
|
|
The args.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.tvm_thread_allreduce", *freduce_args)
|
|
|
|
|
|
def tvm_thread_invariant(cond):
|
|
"""Mark condition as thread invariant.
|
|
|
|
Parameters
|
|
----------
|
|
cond : Expr
|
|
The condition.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
assert tvm.ir.is_prim_expr(cond)
|
|
return call_intrin(_primexpr_ty(cond), "tirx.tvm_thread_invariant", cond)
|
|
|
|
|
|
def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1):
|
|
"""Perform synchronization in specified scope.
|
|
|
|
Parameters
|
|
----------
|
|
storage_scope : str
|
|
The storage scope to perform synchronization.
|
|
|
|
is_load : bool
|
|
Whether to perform load synchronization. (for global sync only)
|
|
|
|
num_blocks : int
|
|
The number of blocks to synchronize. (for global sync only)
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.tvm_storage_sync", storage_scope, is_load, num_blocks)
|
|
|
|
|
|
def tvm_kernel_replace_point():
|
|
"""Mark where a transform should replace generated kernel initialization."""
|
|
return call_intrin("void", "tirx.tvm_kernel_replace_point")
|
|
|
|
|
|
def tvm_global_barrier_kinit():
|
|
"""Initialize the global barrier.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.tvm_global_barrier_kinit")
|
|
|
|
|
|
def tvm_warp_shuffle(mask, value, warp_id, width, warp_size):
|
|
"""Exchange value between threads inside a warp.
|
|
|
|
Parameters
|
|
----------
|
|
mask : Expr
|
|
The warp mask indicates active threads inside warp.
|
|
value : Expr
|
|
The value to exchange.
|
|
warp_id : Expr
|
|
The source lane index to fetch value.
|
|
width : Expr
|
|
The width of sub-sections to perform warp shuffle.
|
|
warp_size : Expr
|
|
The warp size.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
_primexpr_ty(value), "tirx.tvm_warp_shuffle", mask, value, warp_id, width, warp_size
|
|
)
|
|
|
|
|
|
def tvm_warp_shuffle_up(mask, value, offset, width, warp_size):
|
|
"""Copy value from a lane with lower (by offset) index relative to caller.
|
|
|
|
Parameters
|
|
----------
|
|
mask : Expr
|
|
The warp mask indicates active threads inside warp.
|
|
value : Expr
|
|
The value to exchange.
|
|
offset : Expr
|
|
The difference between source lane index and destination lane index:
|
|
`offset = dst_lane_idx - src_lane_idx`
|
|
width : Expr
|
|
The width of sub-sections to perform warp shuffle.
|
|
warp_size : Expr
|
|
The warp size.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
_primexpr_ty(value), "tirx.tvm_warp_shuffle_up", mask, value, offset, width, warp_size
|
|
)
|
|
|
|
|
|
def tvm_warp_shuffle_down(mask, value, offset, width, warp_size):
|
|
"""Copy value from a lane with higher (by offset) index relative to caller.
|
|
|
|
Parameters
|
|
----------
|
|
mask : Expr
|
|
The warp mask indicates active threads inside warp.
|
|
value : Expr
|
|
The value to exchange.
|
|
offset : Expr
|
|
The difference between source lane index and destination lane index:
|
|
`offset = src_lane_idx - dst_lane_idx`
|
|
width : Expr
|
|
The width of sub-sections to perform warp shuffle.
|
|
warp_size : Expr
|
|
The warp size.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
_primexpr_ty(value), "tirx.tvm_warp_shuffle_down", mask, value, offset, width, warp_size
|
|
)
|
|
|
|
|
|
def tvm_warp_shuffle_xor(mask, value, lane_mask, width, warp_size):
|
|
"""Copy value from a lane with index computed by `src_lane_idx ^ lane_mask`.
|
|
|
|
Parameters
|
|
----------
|
|
mask : Expr
|
|
The warp mask indicates active threads inside warp.
|
|
value : Expr
|
|
The value to exchange.
|
|
lane_mask : Expr
|
|
The mask to compute source lane index:
|
|
width : Expr
|
|
The width of sub-sections to perform warp shuffle.
|
|
warp_size : Expr
|
|
The warp size.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
_primexpr_ty(value), "tirx.tvm_warp_shuffle_xor", mask, value, lane_mask, width, warp_size
|
|
)
|
|
|
|
|
|
def tvm_warp_activemask():
|
|
"""Return a 32-bit mask indicates currently active threads in a calling warp.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("uint32", "tirx.tvm_warp_activemask")
|
|
|
|
|
|
def type_annotation(dtype):
|
|
"""Create a type annotation expression
|
|
|
|
Parameters
|
|
----------
|
|
dtype : Expr
|
|
The data type.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(dtype, "tirx.type_annotation")
|
|
|
|
|
|
def tvm_access_ptr(ptype, data, offset, extent, rw_mask):
|
|
"""Get head access address with memory access pattern info
|
|
|
|
Parameters
|
|
----------
|
|
ptype : Expr or str
|
|
The data type of pointer. If a ``str``, it is wrapped via
|
|
:func:`type_annotation` so that the lowering rule (which reads
|
|
``args[0].dtype()`` for the cast type) sees the intended dtype
|
|
instead of ``void`` from a raw StringImm.
|
|
|
|
data : DType*
|
|
The data of pointer.
|
|
|
|
offset : int
|
|
The offset of pointer.
|
|
|
|
extent : int
|
|
The extent of pointer.
|
|
|
|
rw_mask : int
|
|
The read write mask.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
if isinstance(ptype, str):
|
|
ptype = type_annotation(ptype)
|
|
data_type = getattr(data, "ty", None)
|
|
storage_scope = data_type.storage_scope if isinstance(data_type, PointerType) else "global"
|
|
return call_intrin(
|
|
PointerType(_primexpr_ty(ptype), storage_scope),
|
|
"tirx.tvm_access_ptr",
|
|
ptype,
|
|
data,
|
|
offset,
|
|
extent,
|
|
rw_mask,
|
|
)
|
|
|
|
|
|
def ptr_byte_offset(data, byte_offset, dtype):
|
|
"""Cast ``data + byte_offset`` to ``dtype*``.
|
|
|
|
``byte_offset`` is always in bytes. Use this when the source CUDA shape
|
|
needs an explicitly typed local pointer derived from a byte-addressed base.
|
|
"""
|
|
if isinstance(dtype, str):
|
|
dtype = type_annotation(dtype)
|
|
data_type = getattr(data, "ty", None)
|
|
storage_scope = data_type.storage_scope if isinstance(data_type, PointerType) else "global"
|
|
return call_intrin(
|
|
PointerType(_primexpr_ty(dtype), storage_scope),
|
|
"tirx.ptr_byte_offset",
|
|
data,
|
|
byte_offset,
|
|
dtype,
|
|
)
|
|
|
|
|
|
def tvm_throw_last_error():
|
|
"""Throw TVMGetLastError()
|
|
|
|
Returns
|
|
-------
|
|
ret : Expr
|
|
The return expression
|
|
"""
|
|
return call_intrin("void", "tirx.tvm_throw_last_error")
|
|
|
|
|
|
def print_buffer(buffer_var, dtype, is_string, is_scalar, dim_num, *shape):
|
|
"""Print out buffer memory during runtime."""
|
|
if len(shape) == 1 and isinstance(shape[0], tuple | list | tvm.ir.Array):
|
|
final_shape_args = list(shape[0])
|
|
else:
|
|
final_shape_args = list(shape)
|
|
if isinstance(dtype, tvm.ir.PrimType):
|
|
dtype = dtype.dtype
|
|
return _ffi_api.print_buffer(
|
|
buffer_var, dtype, is_string, is_scalar, dim_num, *final_shape_args
|
|
)
|
|
|
|
|
|
def cooperative_tensor_fill(
|
|
d: Var,
|
|
index: Expr,
|
|
value: Expr,
|
|
rows: int,
|
|
cols: int,
|
|
):
|
|
return call_intrin("void", "tirx.cooperative_tensor_fill", d, index, value, rows, cols)
|
|
|
|
|
|
def cooperative_tensor_load(
|
|
d: Var,
|
|
index: Expr,
|
|
ptr: Expr,
|
|
stride: Expr,
|
|
rows: int,
|
|
cols: int,
|
|
transpose_matrix: bool = False,
|
|
mma_M: int = 0,
|
|
mma_N: int = 0,
|
|
mma_K: int = 0,
|
|
operand_role: int = 0,
|
|
):
|
|
return call_intrin(
|
|
"void",
|
|
"tirx.cooperative_tensor_load",
|
|
d,
|
|
index,
|
|
ptr,
|
|
stride,
|
|
rows,
|
|
cols,
|
|
transpose_matrix,
|
|
mma_M,
|
|
mma_N,
|
|
mma_K,
|
|
operand_role,
|
|
)
|
|
|
|
|
|
def cooperative_tensor_store(
|
|
d: Expr,
|
|
index: Expr,
|
|
ptr: Expr,
|
|
stride: Expr,
|
|
rows: int,
|
|
cols: int,
|
|
transpose_matrix: bool = False,
|
|
mma_M: int = 0,
|
|
mma_N: int = 0,
|
|
mma_K: int = 0,
|
|
operand_role: int = 0,
|
|
):
|
|
return call_intrin(
|
|
"void",
|
|
"tirx.cooperative_tensor_store",
|
|
d,
|
|
index,
|
|
ptr,
|
|
stride,
|
|
rows,
|
|
cols,
|
|
transpose_matrix,
|
|
mma_M,
|
|
mma_N,
|
|
mma_K,
|
|
operand_role,
|
|
)
|
|
|
|
|
|
def cooperative_tensor_multiply_accumulate(
|
|
d: Var,
|
|
index_d: Expr,
|
|
a: Var,
|
|
index_a: Expr,
|
|
b: Var,
|
|
index_b: Expr,
|
|
c: Var,
|
|
index_c: Expr,
|
|
M: int,
|
|
N: int,
|
|
K: int,
|
|
transpose_a: bool = False,
|
|
transpose_b: bool = False,
|
|
):
|
|
return call_intrin(
|
|
"void",
|
|
"tirx.cooperative_tensor_multiply_accumulate",
|
|
d,
|
|
index_d,
|
|
a,
|
|
index_a,
|
|
b,
|
|
index_b,
|
|
c,
|
|
index_c,
|
|
M,
|
|
N,
|
|
K,
|
|
transpose_a,
|
|
transpose_b,
|
|
)
|
|
|
|
|
|
def vectorlow(dtype, vec):
|
|
"""Get the low level half of the vector
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
vec : list
|
|
The input vector.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(dtype, "tirx.vectorlow", vec)
|
|
|
|
|
|
def vectorhigh(dtype, vec):
|
|
"""Get the high level half of the vector
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
vec : list
|
|
The input vector.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(dtype, "tirx.vectorhigh", vec)
|
|
|
|
|
|
def vectorcombine(dtype, vec1, vec2):
|
|
"""Concat two vectors
|
|
|
|
Parameters
|
|
----------
|
|
vec1 : list
|
|
The input vector.
|
|
|
|
vec2 : list
|
|
The input vector.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(dtype, "tirx.vectorcombine", vec1, vec2)
|
|
|
|
|
|
def dp4a(vec1, vec2, acc=0):
|
|
"""Dot product of two int8x4 vectors and add an optional accumulator
|
|
|
|
Parameters
|
|
----------
|
|
vec1 : int8x4
|
|
The input vector.
|
|
|
|
vec2 : int8x4
|
|
The input vector.
|
|
|
|
acc : int32
|
|
The accumulator.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("int32", "tirx.dp4a", vec1, vec2, acc)
|
|
|
|
|
|
def ret(val, span=None):
|
|
"""Create a tir return expression
|
|
|
|
Parameters
|
|
----------
|
|
val : Expr
|
|
The returned tir expression, whose data type is int, float or void pointer.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
ret : Expr
|
|
The return expression
|
|
"""
|
|
if not isinstance(val, Expr):
|
|
val = tirx.convert(val)
|
|
return Call(Op.get("tirx.ret"), [val], span=span, ret_ty=val.ty)
|
|
|
|
|
|
def any(*args, span=None):
|
|
"""Create a new experssion of the union of all conditions in the arguments
|
|
|
|
Parameters
|
|
----------
|
|
args : list
|
|
List of symbolic boolean expressions
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
expr: Expr
|
|
Expression
|
|
"""
|
|
if not args:
|
|
raise ValueError("Any must take at least 1 argument")
|
|
if len(args) == 1:
|
|
return args[0]
|
|
val = _ffi_api._OpOr(args[0], args[1], span) # type: ignore
|
|
for i in range(2, len(args)):
|
|
val = _ffi_api._OpOr(val, args[i], span) # type: ignore
|
|
return val
|
|
|
|
|
|
def all(*args, span=None):
|
|
"""Create a new expression of the intersection of all conditions in the
|
|
arguments
|
|
|
|
Parameters
|
|
----------
|
|
args : list
|
|
List of symbolic boolean expressions
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
expr: Expr
|
|
Expression
|
|
"""
|
|
if not args:
|
|
raise ValueError("Any must take at least 1 argument")
|
|
if len(args) == 1:
|
|
return args[0]
|
|
val = _ffi_api._OpAnd(args[0], args[1], span) # type: ignore
|
|
for i in range(2, len(args)):
|
|
val = _ffi_api._OpAnd(val, args[i], span) # type: ignore
|
|
return val
|
|
|
|
|
|
@tvm_ffi.register_global_func("tvm.default_trace_action")
|
|
def _tvm_default_trace_action(*args):
|
|
print(list(args))
|
|
|
|
|
|
def trace(args, trace_action="tvm.default_trace_action"):
|
|
"""Trace tensor data at the runtime.
|
|
|
|
The trace function allows to trace specific tensor at the
|
|
runtime. The tracing value should come as last argument.
|
|
The trace action should be specified, by default
|
|
tvm.default_trace_action is used.
|
|
|
|
Parameters
|
|
----------
|
|
args : list of Expr or Buffers.
|
|
Positional arguments.
|
|
|
|
trace_action : str.
|
|
The name of the trace action.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
|
|
See Also
|
|
--------
|
|
tvm.tirx.call_packed : Creates packed function.
|
|
"""
|
|
if not isinstance(args, list):
|
|
raise Exception("tvm.tirx.trace consumes the args as list type")
|
|
call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
|
|
call_args.insert(0, tvm.tirx.StringImm(trace_action))
|
|
tracing_value = args[-1]
|
|
ret_ty = tracing_value.ty if isinstance(tracing_value, Expr) else tracing_value.dtype
|
|
return tvm.ir.Call(Op.get("tirx.tvm_call_trace_packed"), call_args, ret_ty=ret_ty)
|
|
|
|
|
|
def min_value(dtype, span=None):
|
|
"""minimum value of dtype
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
value : tvm.Expr
|
|
The minimum value of dtype.
|
|
"""
|
|
return _ffi_api.min_value(dtype, span) # type: ignore
|
|
|
|
|
|
def max_value(dtype: str, span: Span | None = None) -> Any:
|
|
"""maximum value of dtype
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
value : tvm.Expr
|
|
The maximum value of dtype.
|
|
"""
|
|
return _ffi_api.max_value(dtype, span) # type: ignore
|
|
|
|
|
|
def infinity(dtype: str, span: Span | None = None) -> Any:
|
|
"""infinity value of dtype
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
value : tvm.Expr
|
|
The infinity value of dtype.
|
|
"""
|
|
return _ffi_api.infinity(dtype, span) # type: ignore
|
|
|
|
|
|
def reinterpret(dtype, value, span: Span | None = None) -> Expr:
|
|
"""Reinterpret a value as an exact primitive or pointer type.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str or tvm.ir.Type
|
|
The data type.
|
|
|
|
value : Expr
|
|
The input value.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
value : tvm.Expr
|
|
The reinterpret cast value of dtype.
|
|
"""
|
|
if isinstance(dtype, str):
|
|
dtype = (
|
|
PointerType(tvm.ir.PrimType("void")) if dtype == "handle" else tvm.ir.PrimType(dtype)
|
|
)
|
|
return _ffi_api.reinterpret(dtype, value, span) # type: ignore
|
|
|
|
|
|
def exp(x):
|
|
"""Take exponential of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.exp", x)
|
|
|
|
|
|
def exp2(x):
|
|
"""Calculate 2**x
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.exp2", x)
|
|
|
|
|
|
def exp10(x):
|
|
"""Calculate 10**x
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.exp10", x)
|
|
|
|
|
|
def fma(x, y, z):
|
|
"""Take fused multiply-add of input x, y, z.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
First input argument.
|
|
|
|
y : Expr
|
|
Second input argument.
|
|
|
|
z : Expr
|
|
Third input argument.
|
|
|
|
Returns
|
|
-------
|
|
out : Expr
|
|
The result of x * y + z.
|
|
"""
|
|
x = tir.convert(x)
|
|
y = tir.convert(y)
|
|
z = tir.convert(z)
|
|
return call_intrin(_primexpr_ty(x), "tirx.fma", x, y, z)
|
|
|
|
|
|
def erf(x):
|
|
"""Take gauss error function of the input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.erf", x)
|
|
|
|
|
|
def tanh(x):
|
|
"""Take hyperbolic tanh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.tanh", x)
|
|
|
|
|
|
def sigmoid(x):
|
|
"""Quick function to get sigmoid
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.sigmoid", x)
|
|
|
|
|
|
def log(x):
|
|
"""Take log of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.log", x)
|
|
|
|
|
|
def log2(x):
|
|
"""Take log2 of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.log2", x)
|
|
|
|
|
|
def log10(x):
|
|
"""Take log10 of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.log10", x)
|
|
|
|
|
|
def log1p(x):
|
|
"""Take log(x + 1) with respect to input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.log1p", x)
|
|
|
|
|
|
def tan(x):
|
|
"""Take tan of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = _require_float_arg("tan", x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.tan", x)
|
|
|
|
|
|
def cos(x):
|
|
"""Take cos of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = _require_float_arg("cos", x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.cos", x)
|
|
|
|
|
|
def cosh(x):
|
|
"""Take cosh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.cosh", x)
|
|
|
|
|
|
def acos(x):
|
|
"""Take acos of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.acos", x)
|
|
|
|
|
|
def acosh(x):
|
|
"""Take acos of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.acosh", x)
|
|
|
|
|
|
def sin(x):
|
|
"""Take sin of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = _require_float_arg("sin", x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.sin", x)
|
|
|
|
|
|
def sinh(x):
|
|
"""Take sinh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.sinh", x)
|
|
|
|
|
|
def asin(x):
|
|
"""Take asin of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.asin", x)
|
|
|
|
|
|
def asinh(x):
|
|
"""Take asinh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.asinh", x)
|
|
|
|
|
|
def atan(x):
|
|
"""Take atan of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.atan", x)
|
|
|
|
|
|
def atanh(x):
|
|
"""Take atanh of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.atanh", x)
|
|
|
|
|
|
def atan2(x1, x2):
|
|
"""Take arctan2(x1, x2).
|
|
|
|
Parameters
|
|
----------
|
|
x1 : Expr
|
|
Input argument.
|
|
|
|
x2 : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x1 = tir.convert(x1)
|
|
x2 = tir.convert(x2)
|
|
return call_intrin(_primexpr_ty(x1), "tirx.atan2", x1, x2)
|
|
|
|
|
|
def sqrt(x):
|
|
"""Take square root of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.sqrt", x)
|
|
|
|
|
|
def rsqrt(x):
|
|
"""Take reciprocal of square root of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.rsqrt", x)
|
|
|
|
|
|
def clz(x):
|
|
"""Count leading zero bits of an integer x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input 32 or 64 bit integer.
|
|
The result is undefined if the input is 0.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return call_intrin("int32", "tirx.clz", x)
|
|
|
|
|
|
def floor(x: ExprWithOp, span=None):
|
|
"""Take floor of float input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.floor(x, span) # type: ignore
|
|
|
|
|
|
def ceil(x, span=None):
|
|
"""Take ceil of float input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.ceil(x, span) # type: ignore
|
|
|
|
|
|
def trunc(x, span=None):
|
|
"""Get truncated value of the input.
|
|
|
|
The truncated value of the scalar x is the
|
|
nearest integer i which is closer to zero than x is.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.trunc(x, span) # type: ignore
|
|
|
|
|
|
def abs(x, span=None):
|
|
"""Get absolute value of the input element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.abs(x, span) # type: ignore
|
|
|
|
|
|
def bitwise_and(x, y, span=None):
|
|
"""Take bitwise and of two values
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Left operand
|
|
|
|
y : Expr
|
|
Right operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.bitwise_and(x, y, span)
|
|
|
|
|
|
def bitwise_not(x, span=None):
|
|
"""Take bitwise not of input value
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.bitwise_not(x, span)
|
|
|
|
|
|
def bitwise_or(x, y, span=None):
|
|
"""Take bitwise or of two values
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Left operand
|
|
|
|
y : Expr
|
|
Right operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.bitwise_or(x, y, span)
|
|
|
|
|
|
def bitwise_xor(x, y, span=None):
|
|
"""Take bitwise xor of two values
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Left operand
|
|
|
|
y : Expr
|
|
Right operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.bitwise_xor(x, y, span)
|
|
|
|
|
|
def round(x, span=None):
|
|
"""Round elements of the array to the nearest integer.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.round(x, span) # type: ignore
|
|
|
|
|
|
def nearbyint(x, span=None):
|
|
"""Round elements of the array to the nearest integer.
|
|
This intrinsic uses llvm.nearbyint instead of llvm.round
|
|
which is faster but will results different from te.round.
|
|
Notably nearbyint rounds according to the rounding mode,
|
|
whereas te.round (llvm.round) ignores that.
|
|
For differences between the two see:
|
|
https://en.cppreference.com/w/cpp/numeric/math/round
|
|
https://en.cppreference.com/w/cpp/numeric/math/nearbyint
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.nearbyint(x, span) # type: ignore
|
|
|
|
|
|
def nextafter(x1, x2):
|
|
"""Return the next floating-point value after x1 towards x2.
|
|
|
|
Parameters
|
|
----------
|
|
x1 : Expr
|
|
Input argument.
|
|
|
|
x2 : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x1 = tir.convert(x1)
|
|
x2 = tir.convert(x2)
|
|
return call_intrin(_primexpr_ty(x1), "tirx.nextafter", x1, x2) # type: ignore
|
|
|
|
|
|
def hypot(x1, x2):
|
|
"""Equivalent to sqrt(x1**2 + x2**2), element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x1 : Expr
|
|
Input argument.
|
|
|
|
x2 : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x1 = tir.convert(x1)
|
|
x2 = tir.convert(x2)
|
|
return call_intrin(_primexpr_ty(x1), "tirx.hypot", x1, x2) # type: ignore
|
|
|
|
|
|
def copysign(x1, x2):
|
|
"""Change the sign of x1 to that of x2, element-wise.
|
|
|
|
Parameters
|
|
----------
|
|
x1 : Expr
|
|
Input argument.
|
|
|
|
x2 : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x1 = tir.convert(x1)
|
|
x2 = tir.convert(x2)
|
|
return call_intrin(_primexpr_ty(x1), "tirx.copysign", x1, x2) # type: ignore
|
|
|
|
|
|
def ldexp(x1, x2):
|
|
"""Returns x1 * (2 ** x2).
|
|
|
|
Parameters
|
|
----------
|
|
x1 : Expr
|
|
Input argument.
|
|
|
|
x2 : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x1 = tir.convert(x1)
|
|
x2 = tir.convert(x2)
|
|
return call_intrin(_primexpr_ty(x1), "tirx.ldexp", x1, x2) # type: ignore
|
|
|
|
|
|
def likely(cond, span=None):
|
|
"""Mark condition as likely.
|
|
|
|
Parameters
|
|
----------
|
|
|
|
cond : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The marked expression.
|
|
"""
|
|
return _ffi_api.likely(cond, span) # type: ignore
|
|
|
|
|
|
def filter(var, pred, *, span=None): # pylint: disable=redefined-builtin
|
|
"""Thread-set filter escape hatch.
|
|
|
|
Use this wrapper only when the predicate is *not* in the canonical
|
|
thread-filter grammar (see ``src/tirx/analysis/filter_canonical.h``).
|
|
Canonical predicates -- pure conjunctions of ``scopeid_var <op> const``
|
|
comparisons plus bare ``T.ptx.elect_sync()`` calls -- are recognized by
|
|
the lowering pass directly from ``if cond:``, so the wrapper is redundant
|
|
for them.
|
|
|
|
When wrapped: ``var`` (a ``ScopeIdDef``-declared scope identifier) tells
|
|
the compiler which active-set axis to collapse to a singleton when the
|
|
opaque predicate evaluates true; ``pred`` is preserved verbatim and
|
|
evaluated at runtime.
|
|
|
|
The legacy three-argument range form ``filter(var, lo, hi)`` has been
|
|
removed -- write ``lo <= var and var < hi`` (or ``var == lo`` when
|
|
``hi == lo + 1``) at the call site instead.
|
|
"""
|
|
return call_intrin("bool", "tirx.filter", var, pred, span=span)
|
|
|
|
|
|
def selector(var, pred, span=None):
|
|
"""Analysis-only active-thread selector.
|
|
|
|
``selector(var, pred)`` denotes the unique value of ``var`` in the current
|
|
active domain for which ``pred`` is true. It is intended for compiler
|
|
metadata and should not survive to executable codegen.
|
|
"""
|
|
return call_intrin(_primexpr_ty(var), "tirx.selector", var, pred, span=span)
|
|
|
|
|
|
def isnan(x, span=None):
|
|
"""Check if input value is Nan.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.isnan(x, span) # type: ignore
|
|
|
|
|
|
def isnullptr(x, span=None):
|
|
"""Check if input value is nullptr.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return call_intrin("bool", "tirx.isnullptr", x, span=span) # type: ignore
|
|
|
|
|
|
def isfinite(x, span=None):
|
|
"""Check if input value is finite.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.isfinite(x, span) # type: ignore
|
|
|
|
|
|
def isinf(x, span=None):
|
|
"""Check if input value is infinite.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.isinf(x, span) # type: ignore
|
|
|
|
|
|
def power(x, y, span=None):
|
|
"""x power y
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
y : Expr
|
|
The exponent
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
z : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api._OpPow(x, y, span) # type: ignore
|
|
|
|
|
|
def pow(x, y, span=None):
|
|
"""x power y
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
y : Expr
|
|
The exponent
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
z : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api._OpPow(x, y, span) # type: ignore
|
|
|
|
|
|
def popcount(x):
|
|
"""Count the number of set bits in input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
return call_intrin(_primexpr_ty(x), "tirx.popcount", x)
|
|
|
|
|
|
def q_multiply_shift(x, y, q, s):
|
|
"""Execute a multiplication between two Q-numbers x and y
|
|
followed by a right shift s. The mathematical expression is:
|
|
|
|
out = round(x*y*2^-s)
|
|
|
|
More about Q-numbers here: https://en.wikipedia.org/wiki/Q_(number_format)
|
|
The rounding rule is to the nearest value, rounding half up
|
|
(i.e., round(x.1) = x and round (x.5) = x+1)
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
First Q-number
|
|
y : Expr
|
|
Second Q-number
|
|
q : Expr
|
|
Number of fractional bits in x and y. Needs to be > 0
|
|
s : Expr
|
|
Integer shift
|
|
|
|
Returns
|
|
-------
|
|
y : Expr
|
|
The result.
|
|
"""
|
|
return call_intrin("int32", "tirx.q_multiply_shift", x, y, q, s)
|
|
|
|
|
|
def q_multiply_shift_per_axis(
|
|
x: Expr,
|
|
y: Expr,
|
|
ls: Expr,
|
|
rs: Expr,
|
|
q: IntImm,
|
|
is_lshift_required: IntImm,
|
|
is_rshift_required: IntImm,
|
|
):
|
|
"""Execute a multiplication between two Q-numbers x and y
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
First Q-number.
|
|
y : Expr
|
|
Second Q-number.
|
|
ls : Expr
|
|
Integer left shift.
|
|
rs : Expr
|
|
Integer right shift.
|
|
q : IntImm
|
|
Number of fractional bits in x and y. Needs to be > 0.
|
|
is_lshift_required : IntImm
|
|
Whether we need to do left shift or not.
|
|
is_rshift_required : IntImm
|
|
Whether we need to do right shift or not.
|
|
|
|
Returns
|
|
-------
|
|
z : Expr
|
|
The result.
|
|
"""
|
|
return call_intrin(
|
|
"int32",
|
|
"tirx.q_multiply_shift_per_axis",
|
|
x,
|
|
y,
|
|
ls,
|
|
rs,
|
|
q,
|
|
is_lshift_required,
|
|
is_rshift_required,
|
|
)
|
|
|
|
|
|
def shift_left(x, y, span=None):
|
|
"""Return the result of x left shifted by y bits.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
y : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
z : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.left_shift(x, y, span)
|
|
|
|
|
|
def shift_right(x, y, span=None):
|
|
"""Return the result of x right shifted by y bits.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
|
|
y : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
z : Expr
|
|
The result.
|
|
"""
|
|
return _ffi_api.right_shift(x, y, span)
|
|
|
|
|
|
def fmod(x, y):
|
|
"""Return the remainder of x divided by y with the same sign as x.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
Input argument.
|
|
y : Expr
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
z : Expr
|
|
The result.
|
|
"""
|
|
x = tir.convert(x)
|
|
y = tir.convert(y)
|
|
return call_intrin(_primexpr_ty(x), "tirx.fmod", x, y)
|
|
|
|
|
|
def if_then_else(cond, t, f, span=None):
|
|
"""Conditional selection expression.
|
|
|
|
Parameters
|
|
----------
|
|
cond : Expr
|
|
The condition
|
|
|
|
t : Expr
|
|
The result expression if cond is true.
|
|
|
|
f : Expr
|
|
The result expression if cond is false.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
result : Node
|
|
The result of conditional expression.
|
|
|
|
Note
|
|
----
|
|
Unlike Select, if_then_else will not execute
|
|
the branch that does not satisfy the condition.
|
|
You can use it to guard against out of bound access.
|
|
Unlike Select, if_then_else cannot be vectorized
|
|
if some lanes in the vector have different conditions.
|
|
"""
|
|
return _ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore
|
|
|
|
|
|
def div(a, b, span=None):
|
|
"""Compute a / b as in C/C++ semantics.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand, known to be non-negative.
|
|
|
|
b : Expr
|
|
The right hand operand, known to be non-negative.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
Note
|
|
----
|
|
When operands are integers, returns truncdiv(a, b, span).
|
|
"""
|
|
return _ffi_api._OpDiv(a, b, span) # type: ignore
|
|
|
|
|
|
def indexdiv(a, b, span=None):
|
|
"""Compute floor(a / b) where a and b are non-negative.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand, known to be non-negative.
|
|
|
|
b : Expr
|
|
The right hand operand, known to be non-negative.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
|
|
Note
|
|
----
|
|
Use this function to split non-negative indices.
|
|
This function may take advantage of operands'
|
|
non-negativeness.
|
|
"""
|
|
return _ffi_api._OpIndexDiv(a, b, span) # type: ignore
|
|
|
|
|
|
def indexmod(a, b, span=None):
|
|
"""Compute the remainder of indexdiv. a and b are non-negative.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand, known to be non-negative.
|
|
|
|
b : Expr
|
|
The right hand operand, known to be non-negative.
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
|
|
Note
|
|
----
|
|
Use this function to split non-negative indices.
|
|
This function may take advantage of operands'
|
|
non-negativeness.
|
|
"""
|
|
return _ffi_api._OpIndexMod(a, b, span) # type: ignore
|
|
|
|
|
|
def truncdiv(a, b, span=None):
|
|
"""Compute the truncdiv of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
|
|
Note
|
|
----
|
|
This is the default integer division behavior in C.
|
|
"""
|
|
return _ffi_api._OpTruncDiv(a, b, span) # type: ignore
|
|
|
|
|
|
def truncmod(a, b, span=None):
|
|
"""Compute the truncmod of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
|
|
Note
|
|
----
|
|
This is the default integer division behavior in C.
|
|
"""
|
|
return _ffi_api._OpTruncMod(a, b, span) # type: ignore
|
|
|
|
|
|
def floordiv(a, b, span=None):
|
|
"""Compute the floordiv of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
"""
|
|
return _ffi_api._OpFloorDiv(a, b, span) # type: ignore
|
|
|
|
|
|
def logaddexp(a, b, span=None):
|
|
"""Compute the logaddexp of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
"""
|
|
return _ffi_api._OpLogAddExp(a, b, span) # type: ignore
|
|
|
|
|
|
def floormod(a, b, span=None):
|
|
"""Compute the floormod of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
"""
|
|
return _ffi_api._OpFloorMod(a, b, span) # type: ignore
|
|
|
|
|
|
def ceildiv(lhs, rhs, span=None):
|
|
"""Generic ceildiv operator.
|
|
|
|
Parameters
|
|
----------
|
|
lhs : object
|
|
The left operand.
|
|
rhs : object
|
|
The right operand.
|
|
span : Optional[Span]
|
|
The location of this operator in the source.
|
|
|
|
Returns
|
|
-------
|
|
op : tvm.Expr
|
|
The result Expr of ceildiv operaton.
|
|
"""
|
|
return _ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore
|
|
|
|
|
|
def comm_reducer(fcombine, fidentity, name="reduce"):
|
|
"""Create a commutative reducer for reduction.
|
|
|
|
Parameters
|
|
----------
|
|
fcombine : function(Expr -> Expr -> Expr)
|
|
A binary function which takes two Expr as input to return a Expr.
|
|
|
|
fidentity : function(str -> Expr)
|
|
A function which takes a type string as input to return a const Expr.
|
|
|
|
Returns
|
|
-------
|
|
reducer : function
|
|
A function which creates a reduce expression over axis.
|
|
There are two ways to use it:
|
|
|
|
1. accept (expr, axis, where) to produce an Reduce Expr on
|
|
specified axis;
|
|
2. simply use it with multiple Exprs.
|
|
|
|
Example
|
|
-------
|
|
.. code-block:: python
|
|
|
|
n = te.var("n")
|
|
m = te.var("m")
|
|
mysum = te.comm_reducer(lambda x, y: x+y,
|
|
lambda t: tvm.tirx.const(0, dtype=t), name="mysum")
|
|
A = te.placeholder((n, m), name="A")
|
|
k = te.reduce_axis((0, m), name="k")
|
|
B = te.compute((n,), lambda i: mysum(A[i, k], axis=k), name="B")
|
|
"""
|
|
|
|
def _reduce_directly(*args):
|
|
num = len(args)
|
|
# process `where` is None
|
|
if num == 3 and args[2] is None:
|
|
num = 2
|
|
res = args[0]
|
|
for i in range(num - 1):
|
|
res = fcombine(res, args[i + 1])
|
|
return res
|
|
|
|
def _make_reduce(expr, axis, where=None, init=None):
|
|
code = fcombine.__code__
|
|
assert fcombine.__code__.co_argcount == 2
|
|
expr = tir.convert(expr)
|
|
if init is not None:
|
|
init = tir.convert(init)
|
|
if isinstance(expr, Array):
|
|
size = len(expr)
|
|
lhs = []
|
|
rhs = []
|
|
dtypes = []
|
|
for i in range(size):
|
|
dtype = _primexpr_dtype(expr[i])
|
|
dtypes.append(dtype)
|
|
lname = code.co_varnames[0] + "_" + str(i)
|
|
lhs.append(Var(lname, dtype))
|
|
rname = code.co_varnames[1] + "_" + str(i)
|
|
rhs.append(Var(rname, dtype))
|
|
if init is None:
|
|
init = []
|
|
result = fcombine(lhs, rhs)
|
|
id_elem = fidentity(*dtypes)
|
|
else:
|
|
assert tvm.ir.is_prim_expr(expr)
|
|
size = 1
|
|
dtype = _primexpr_dtype(expr)
|
|
lvar = Var(code.co_varnames[0], dtype)
|
|
rvar = Var(code.co_varnames[1], dtype)
|
|
result = [fcombine(lvar, rvar)]
|
|
id_elem = [fidentity(dtype)]
|
|
lhs = [lvar]
|
|
rhs = [rvar]
|
|
expr = [expr]
|
|
if init is not None:
|
|
init = [init]
|
|
combiner = CommReducer(lhs, rhs, result, id_elem)
|
|
if not isinstance(axis, list | tuple | tvm.ir.Array):
|
|
axis = [axis]
|
|
if where is None:
|
|
where = tir.convert(True)
|
|
if init is None:
|
|
outputs = tuple(
|
|
tvm.tirx.Reduce(combiner, expr, axis, where, i, []) for i in range(size)
|
|
)
|
|
else:
|
|
outputs = tuple(
|
|
tvm.tirx.Reduce(combiner, expr, axis, where, i, init) for i in range(size)
|
|
)
|
|
return outputs[0] if size == 1 else outputs
|
|
|
|
# pylint: disable=keyword-arg-before-vararg
|
|
def reducer(expr, axis, where=None, init=None, *args):
|
|
if isinstance(axis, tvm.tirx.IterVar | list | tuple):
|
|
assert not args
|
|
return _make_reduce(expr, axis, where, init)
|
|
|
|
if where is None:
|
|
assert not args
|
|
assert init is None
|
|
return _reduce_directly(expr, axis)
|
|
elif init is None:
|
|
assert not args
|
|
return _reduce_directly(expr, axis, where)
|
|
else:
|
|
return _reduce_directly(expr, axis, where, init, *args)
|
|
|
|
doc_str = """Create a {0} expression over axis.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The source expression.
|
|
axis : IterVar
|
|
The reduction IterVar axis
|
|
where : optional, Expr
|
|
Filtering predicate of the reduction.
|
|
Returns
|
|
-------
|
|
value : Expr
|
|
The result value.
|
|
|
|
Example
|
|
-------
|
|
.. code-block:: python
|
|
|
|
m = te.var("m")
|
|
n = te.var("n")
|
|
A = te.placeholder((m, n), name="A")
|
|
k = te.reduce_axis((0, n), name="k")
|
|
|
|
# there are two way to use this {0} reducer:
|
|
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
|
|
# tvm.{0} represents tvm.te.{0} or tvm.tirx.{0}.
|
|
B = te.compute((m,), lambda i: tvm.{0}(A[i, k], axis=k), name="B")
|
|
|
|
# mode 2, simply use it with multiple Exprs:
|
|
{0}_res = tvm.{0}(m, n)
|
|
"""
|
|
reducer.__doc__ = doc_str.format(name)
|
|
return reducer
|
|
|
|
|
|
def TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint):
|
|
"""Backend function to allocate temporal workspace
|
|
|
|
Parameters
|
|
----------
|
|
device_type : int
|
|
The device type which the space will be allocated.
|
|
|
|
device_id : int
|
|
The device id which the space will be allocated.
|
|
|
|
nbytes : int
|
|
The size of the space requested.
|
|
|
|
dtype_code_hint : int
|
|
The type code of the array elements. Only used in certain backends such as OpenGL.
|
|
|
|
dtype_bits_hint : int
|
|
The type bits of the array elements. Only used in certain backends such as OpenGL.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"handle",
|
|
"tirx.TVMBackendAllocWorkspace",
|
|
device_type,
|
|
device_id,
|
|
nbytes,
|
|
dtype_code_hint,
|
|
dtype_bits_hint,
|
|
)
|
|
|
|
|
|
def TVMBackendFreeWorkspace(device_type, device_id, ptr):
|
|
"""Backend function to free temporal workspace.
|
|
|
|
Parameters
|
|
----------
|
|
device_type : int
|
|
The device type which the space will be allocated.
|
|
|
|
device_id : int
|
|
The device id which the space will be allocated.
|
|
|
|
ptr : Var
|
|
The result allocated space pointer.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("int32", "tirx.TVMBackendFreeWorkspace", device_type, device_id, ptr)
|
|
|
|
|
|
def anylist_getitem(list_handle, index):
|
|
"""Returns an item from any list.
|
|
list_handle: Var
|
|
The handle to anylist
|
|
index : int
|
|
The index
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("handle", "tirx.anylist_getitem", list_handle, index)
|
|
|
|
|
|
def anylist_resetitem(list_handle, index):
|
|
"""Reset an item from any list.
|
|
list_handle: Var
|
|
The handle to anylist
|
|
index : int
|
|
The index
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("int", "tirx.anylist_resetitem", list_handle, index)
|
|
|
|
|
|
def anylist_setitem_call_packed(list_handle, index, func_name, *args):
|
|
"""Set anylist item by result of packed call.
|
|
list_handle: Var
|
|
The handle to anylist
|
|
index : int
|
|
The index
|
|
func_name: str
|
|
The name of the function to be called.
|
|
args:
|
|
Extra arguments
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"int", "tirx.anylist_setitem_call_packed", list_handle, index, func_name, *args
|
|
)
|
|
|
|
|
|
def anylist_setitem_call_cpacked(list_handle, index, func_name, *args):
|
|
"""Set anylist item by result of packed call.
|
|
list_handle: Var
|
|
The handle to anylist
|
|
index : int
|
|
The index
|
|
func_name: str
|
|
The name of the function to be called.
|
|
args:
|
|
Extra arguments
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"int", "tirx.anylist_setitem_call_cpacked", list_handle, index, func_name, *args
|
|
)
|
|
|
|
|
|
def vscale():
|
|
"""Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
|
|
(https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
Call to the vscale intrinsic
|
|
"""
|
|
return call_intrin("int32", "tirx.vscale")
|
|
|
|
|
|
def get_active_lane_mask(dtype, base, limit):
|
|
"""
|
|
Calculate a predicate mask given an upper bound (limit) and a current value (base).
|
|
|
|
It will be lowered to the llvm.get.active.lane.mask intrinsic.
|
|
(https://llvm.org/docs/LangRef.html#llvm-get-active-lane-mask-intrinsics)
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the result.
|
|
|
|
base : Expr
|
|
An expression reprsenting the base.
|
|
|
|
limit : Expr
|
|
An expression representing the limit.
|
|
"""
|
|
return call_intrin(dtype, "tirx.get_active_lane_mask", base, limit)
|
|
|
|
|
|
def get_vscale_expr(dtype: str | tvm_ffi.dtype, min_size: int = 128) -> Expr:
|
|
"""
|
|
Create a datatype dependent scalable expression.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : Union[str, tvm_ffi.DataType]
|
|
Element data type.
|
|
min_size : int
|
|
The minimum size of the scalable vector in bits.
|
|
"""
|
|
if isinstance(dtype, str):
|
|
dtype = tvm_ffi.dtype(dtype)
|
|
return min_size // dtype.bits * vscale()
|
|
|
|
|
|
def ignore_loop_partition(predicate) -> Expr:
|
|
"""
|
|
Annotate a predicate not be considered as target condition of loop partition.
|
|
|
|
Parameters
|
|
----------
|
|
predicate : Expr
|
|
The annotated predicate expression.
|
|
"""
|
|
return call_intrin("bool", "tirx.ignore_loop_partition", predicate)
|
|
|
|
|
|
# pylint: disable=unnecessary-lambda
|
|
sum = comm_reducer(lambda x, y: x + y, lambda t: const(0, dtype=t), name="sum")
|
|
min = comm_reducer(lambda x, y: _ffi_api._OpMin(x, y, None), max_value, name="min") # type: ignore
|
|
max = comm_reducer(lambda x, y: _ffi_api._OpMax(x, y, None), min_value, name="max") # type: ignore
|
|
|
|
|
|
def tvm_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout):
|
|
"""TVM intrinsic for tensor core load operators
|
|
|
|
Parameters
|
|
----------
|
|
fragment : Var
|
|
The wmma fragment.
|
|
|
|
m : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
n : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
k : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
index : Expr
|
|
The fragment index.
|
|
|
|
buffer_ptr : Expr
|
|
The fragment buffer pointer.
|
|
|
|
stride : Expr
|
|
The fragment stride.
|
|
|
|
layout : Literal["row_major", "column_major"]
|
|
The fragment layout.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"void", "tirx.tvm_load_matrix_sync", fragment, m, n, k, index, buffer_ptr, stride, layout
|
|
)
|
|
|
|
|
|
def tvm_mma_sync(
|
|
fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c
|
|
):
|
|
"""TVM intrinsic for tensor core mma_sync operators
|
|
|
|
Parameters
|
|
----------
|
|
fragment_d : Var
|
|
The wmma fragment_d.
|
|
|
|
index_d : Expr
|
|
The fragment_d index.
|
|
|
|
fragment_a : Var
|
|
The wmma fragment_a.
|
|
|
|
index_a : Expr
|
|
The fragment_a index.
|
|
|
|
fragment_b : Var
|
|
The wmma fragment_b.
|
|
|
|
index_b : Expr
|
|
The fragment_b index.
|
|
|
|
fragment_c : Var
|
|
The wmma fragment_c.
|
|
|
|
index_c : Expr
|
|
The fragment_c index.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"void",
|
|
"tirx.tvm_mma_sync",
|
|
fragment_d,
|
|
index_d,
|
|
fragment_a,
|
|
index_a,
|
|
fragment_b,
|
|
index_b,
|
|
fragment_c,
|
|
index_c,
|
|
)
|
|
|
|
|
|
def tvm_bmma_sync(
|
|
fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c
|
|
):
|
|
"""TVM intrinsic for tensor core bmma_sync operators
|
|
|
|
Parameters
|
|
----------
|
|
fragment_d : Var
|
|
The bwmma fragment_d.
|
|
|
|
index_d : Expr
|
|
The fragment_d index.
|
|
|
|
fragment_a : Var
|
|
The bwmma fragment_a.
|
|
|
|
index_a : Expr
|
|
The fragment_a index.
|
|
|
|
fragment_b : Var
|
|
The bwmma fragment_b.
|
|
|
|
index_b : Expr
|
|
The fragment_b index.
|
|
|
|
fragment_c : Var
|
|
The bwmma fragment_c.
|
|
|
|
index_c : Expr
|
|
The fragment_c index.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"void",
|
|
"tirx.tvm_bmma_sync",
|
|
fragment_d,
|
|
index_d,
|
|
fragment_a,
|
|
index_a,
|
|
fragment_b,
|
|
index_b,
|
|
fragment_c,
|
|
index_c,
|
|
)
|
|
|
|
|
|
def tvm_fill_fragment(fragment, m, n, k, index, value):
|
|
"""TVM intrinsic for tensor core fill_fragment operators
|
|
|
|
Parameters
|
|
----------
|
|
fragment : Var
|
|
The wmma fragment
|
|
|
|
m : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
n : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
k : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
index : Expr
|
|
The fragment index.
|
|
|
|
value : Expr
|
|
The value to be filled in fragment.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("void", "tirx.tvm_fill_fragment", fragment, m, n, k, index, value)
|
|
|
|
|
|
def tvm_store_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout):
|
|
"""TVM intrinsic for tensor core store operators
|
|
|
|
Parameters
|
|
----------
|
|
fragment : Var
|
|
The wmma fragment.
|
|
|
|
m : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
n : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
k : UIntImm
|
|
The shape of wmma fragment.
|
|
|
|
index : Expr
|
|
The fragment index.
|
|
|
|
buffer_ptr : Expr
|
|
The fragment buffer pointer.
|
|
|
|
stride : Expr
|
|
The fragment stride.
|
|
|
|
layout : Literal["row_major", "column_major"]
|
|
The fragment layout.
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin(
|
|
"void", "tirx.tvm_store_matrix_sync", fragment, m, n, k, index, buffer_ptr, stride, layout
|
|
)
|
|
|
|
|
|
def thread_return():
|
|
"""TVM intrinsic to call thread_return()
|
|
|
|
Returns
|
|
-------
|
|
call : Expr
|
|
The call expression.
|
|
"""
|
|
return call_intrin("", "tirx.thread_return")
|
|
|
|
|
|
def continue_loop(span=None):
|
|
"""Create a tir intrinsic call to represent continue expression
|
|
|
|
Parameters
|
|
----------
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
ret : Expr
|
|
The continue expression
|
|
"""
|
|
|
|
return _ffi_api.continue_loop(span)
|
|
|
|
|
|
def break_loop(span=None):
|
|
"""Create a tir intrinsic call to represent break expression
|
|
|
|
Parameters
|
|
----------
|
|
span : Optional[Span]
|
|
The location of this operator in the source code.
|
|
|
|
Returns
|
|
-------
|
|
ret : Expr
|
|
The break expression
|
|
"""
|
|
|
|
return _ffi_api.break_loop(span)
|