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 Poistional arguments. span : Optional[Span] The location of this operator in the source code. Returns ------- call : PrimExpr The call expression. """ # pylint: disable=import-outside-toplevel from tvm.target import codegen llvm_id = codegen.llvm_lookup_intrinsic_id(name) assert llvm_id != 0, "%s is not an LLVM intrinsic" % name return call_intrin( dtype, Op.get("tir.call_llvm_pure_intrin"), tvm.tir.const(llvm_id, "uint32"), *args, span=span, )
def call_llvm_intrin(dtype, name, *args): """Build expression by calling an llvm intrinsic function Parameters ---------- dtype : str The data type of the result. name : str The name of the llvm intrinsic function. args : list Poistional arguments. Returns ------- call : PrimExpr The call expression. """ # pylint: disable=import-outside-toplevel from tvm.target import codegen llvm_id = codegen.llvm_lookup_intrinsic_id(name) assert llvm_id != 0, "%s is not an LLVM intrinsic" % name return call_pure_intrin(dtype, Op.get("tir.call_llvm_intrin"), tvm.tir.const(llvm_id, 'uint32'), *args)
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 : PrimExpr The call expression. """ return Call(dtype, Op.get("tir.call_extern"), convert((StringImm(func_name), ) + args), span=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_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 _pack_buffer(buf): """Build intrinsics that packs the buffer. """ shape = Call("handle", "tir.tvm_stack_make_shape", buf.shape) strides = Call("handle", "tir.tvm_stack_make_shape", buf.strides) if buf.strides else 0 pack_args = [ buf.data, shape, strides, len(buf.shape), const(0, dtype=buf.dtype), buf.elem_offset ] return Call("handle", Op.get("tir.tvm_stack_make_array"), pack_args)
def call_extern(dtype, func_name, *args): """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. Returns ------- call : PrimExpr The call expression. """ return Call(dtype, Op.get("tir.call_extern"), convert((StringImm(func_name), ) + args), Call.Extern)