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:])