Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions include/tvm/tir/builtin.h
Original file line number Diff line number Diff line change
Expand Up @@ -750,6 +750,50 @@ TVM_DLL const Op& start_profile_intrinsic();
*/
TVM_DLL const Op& end_profile_intrinsic();

/*!
* \brief Get a item from any list and return it.
*
* Any anylist_getitem(Handle anylist,
* int index)
* return anylist[index];
* }
*
* \note This intrinsic is only applicable when appearing
* in call_packed and anylist_setitem_call_packed.
*/
TVM_DLL const Op& anylist_getitem();

/*!
* \brief Reset and clear a item in any list.
*
* void anylist_resetitem(Handle anylist,
* int index)
* anylist[index] = nullptr;
* }
*
* \note This intrinsic is only applicable when appearing
* in call_packed and anylist_setitem_call_packed.
*/
TVM_DLL const Op& anylist_resetitem();

/*!
* \brief Set an item into any list by running packed function call.
*
* void anylist_setitem_call_packed(Handle anylist,
* int index,
* name, *args)
*
* anylist[index] = call_packed(name, *args)
* }
* \note This intrinsic can be used in combination with anylist_getitem.
*/
TVM_DLL const Op& anylist_setitem_call_packed();

/*!
* \brief Same as anylist_setitem_call_packed but use C calling convention.
*/
TVM_DLL const Op& anylist_setitem_call_cpacked();

/*! \brief The kind of structure field info used in intrinsic */
enum TVMStructFieldKind : int {
// array head address
Expand Down
8 changes: 8 additions & 0 deletions python/tvm/script/ir_builder/tir/ir.py
Original file line number Diff line number Diff line change
Expand Up @@ -1713,6 +1713,10 @@ def wrapped(*args, **kwargs):
TVMBackendFreeWorkspace = _op_wrapper(_tir_op.TVMBackendFreeWorkspace)
start_profile_intrinsic = _op_wrapper(_tir_op.start_profile_intrinsic)
end_profile_intrinsic = _op_wrapper(_tir_op.end_profile_intrinsic)
anylist_getitem = _op_wrapper(_tir_op.anylist_getitem)
anylist_resetitem = _op_wrapper(_tir_op.anylist_resetitem)
anylist_setitem_call_packed = _op_wrapper(_tir_op.anylist_setitem_call_packed)
anylist_setitem_call_cpacked = _op_wrapper(_tir_op.anylist_setitem_call_cpacked)


def _dtype_forward(func):
Expand Down Expand Up @@ -1988,6 +1992,10 @@ def wrapped(*args, **kwargs):
"start_profile_intrinsic",
"end_profile_intrinsic",
"meta_var",
"anylist_getitem",
"anylist_resetitem",
"anylist_setitem_call_packed",
"anylist_setitem_call_cpacked",
"llvm_lookup_intrinsic_id",
"type_annotation",
"broadcast",
Expand Down
68 changes: 68 additions & 0 deletions python/tvm/tir/op.py
Original file line number Diff line number Diff line change
Expand Up @@ -2931,6 +2931,74 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tir.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 : PrimExpr
The call expression.
"""
return call_intrin("handle", "tir.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 : PrimExpr
The call expression.
"""
return call_intrin("int", "tir.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 : PrimExpr
The call expression.
"""
return call_intrin(
"int", "tir.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 : PrimExpr
The call expression.
"""
return call_intrin(
"int", "tir.anylist_setitem_call_cpacked", list_handle, index, func_name, *args
)


# 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
Expand Down
Loading