def lower_ite(ite_op):
"""Lowered if then else function that calls intrinsic if_then_else.
Unlike a function lowered by create_lower_func, this function
calls the tvm intrinsic if_then_else.
Parameters
----------
ite_op : Op
Takes an if then else op and returns a
call to tir.if_then_else function, passing the op's
arguments. The return type of the call if a uint of the same
width as the custom type is returned.
"""
dtype = ite_op.dtype
t = tvm.DataType(dtype)
assert get_type_registered(t.type_code)
dtype = "uint" + str(t.bits)
if t.lanes > 1:
dtype += "x" + str(t.lanes)
return call_intrin(
dtype,
"tir.if_then_else",
convert(ite_op.args[0]),
convert(ite_op.args[1]),
convert(ite_op.args[2]),
)
def lower_call_pure_extern(op):
"""Lowered call pure extern function that calls intrinsic call_pure_extern.
Unlike a function lowered by create_lower_func, this function
calls the tvm intrinsic call_pure_extern.
Parameters
----------
ite_op : Op
Takes a call_pure_extern op and returns a
call to tir.call_pure_extern function, passing the op's
arguments. The return type of the call if a uint of the same
width as the custom type is returned.
"""
dtype = op.dtype
t = tvm.DataType(dtype)
assert get_type_registered(t.type_code)
dtype = "uint" + str(t.bits)
if t.lanes > 1:
dtype += "x" + str(t.lanes)
return call_intrin(dtype, "tir.call_pure_extern", *op.args)
def test_dtype():
dtype = tvm.DataType("handle")
assert dtype.type_code == tvm.DataTypeCode.HANDLE