This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d805ae3bd9 [TIR] Expose: `call_packed_lowered`, `call_cpacked_lowered`
(#12425)
d805ae3bd9 is described below
commit d805ae3bd99a71255dc7f7d1d5aa0746ab2ed21e
Author: Yaxing Cai <[email protected]>
AuthorDate: Sun Aug 14 21:58:52 2022 -0700
[TIR] Expose: `call_packed_lowered`, `call_cpacked_lowered` (#12425)
Added the following operations in TIR:
- call_packed_lowered
- call_cpacked_lowered
Co-Authored-By: yongwww <[email protected]>
---
python/tvm/tir/__init__.py | 1 +
python/tvm/tir/op.py | 55 ++++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 56 insertions(+)
diff --git a/python/tvm/tir/__init__.py b/python/tvm/tir/__init__.py
index c64b7dfe71..b782cf0425 100644
--- a/python/tvm/tir/__init__.py
+++ b/python/tvm/tir/__init__.py
@@ -45,6 +45,7 @@ from .stmt import BufferRegion, MatchBufferRegion, Block,
BlockRealize
from .function import PrimFunc, TensorIntrin, IndexMap
+from .op import call_packed_lowered, call_cpacked_lowered
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern,
call_extern
from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any,
min_value, max_value, trace
from .op import exp, exp2, exp10, log, log2, log10, log1p, ldexp, clz
diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
index 17005b04a4..4f43437a11 100644
--- a/python/tvm/tir/op.py
+++ b/python/tvm/tir/op.py
@@ -42,6 +42,61 @@ def _pack_buffer(buf, span=None):
return Call("handle", Op.get("tir.tvm_stack_make_array"), pack_args, span)
+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 recieve an TVMArrayHandle whose content is valid during the callback
period.
+ If the PackedFunc is a python callback, then the corresponding argument is
NDArray.
+
+ Parameters
+ ----------
+ args : list of Expr or Buffer.
+ Positional arguments.
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ call : PrimExpr
+ 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("int32", Op.get("tir.tvm_call_packed_lowered"), call_args,
span)
+
+
+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 : PrimExpr
+ 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("int32", Op.get("tir.tvm_call_cpacked_lowered"), call_args,
span)
+
+
def call_packed(*args, span=None):
"""Build expression by call an external packed function.
