def _rule_float_suffix(op):
"""Intrinsic rule: Add float suffix if it is float32.
This is an example intrinsic generation rule.
Parameters
----------
op : PrimExpr
The call expression of original intrinsic.
Returns
-------
ret : PrimExpr
The translated intrinsic rule.
Return same op if no translation is possible.
See Also
--------
register_intrin_rule : The registeration function for intrin rule.
"""
name = op.name
assert name.startswith("tir.")
prefix = name[4:]
if op.dtype == "float32":
return call_pure_extern(op.dtype, "%sf" % prefix, *op.args)
if op.dtype == "float64":
return call_pure_extern(op.dtype, prefix, *op.args)
return op
def _rule_float_direct(op):
"""Intrinsic rule: Directly call pure extern function for floats.
This is an example intrinsic generation rule.
Parameters
----------
op : PrimExpr
The call expression of original intrinsic.
Returns
-------
ret : PrimExpr
The translated intrinsic rule.
Return same op if no translation is possible.
See Also
--------
register_intrin_rule : The registeration function for intrin rule.
"""
if str(op.dtype).startswith("float"):
return call_pure_extern(op.dtype, op.op.name[4:], *op.args)
return None
def reduce_op(x, y): assert x.dtype == y.dtype , "Reduing elements that don't have same data type: %s v.s. %s" % (x.dtype, y.dtype) return tir.call_pure_extern(x.dtype, name, x, y, *args[1:])
