<jupyter_start><jupyter_text>Intrinsics and Math Functions=============================**Author**: `<NAME> `_While TVM supports basic arithmetic operations. In many casesusually we will need more complicated builtin functions.For example :code:`exp` to take the exponential of the function.These functions are target system dependent and may have differentnames of different target platforms. In this tutorial, we will learnhow we can invoke these target specific functions, and how we can unifythe interface via tvm's intrinsic API.<jupyter_code>from __future__ import absolute_import, print_function

import tvm
from tvm import te
import numpy as np<jupyter_output><empty_output><jupyter_text>Direct Declare Extern Math Call-------------------------------The most straight-forward way to call target specific function is viaextern function call construct in tvm.In the following example, we use :any:`tvm.tir.call_pure_extern` to call:code:`__expf` function, which is only available under CUDA.<jupyter_code>n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: tvm.tir.call_pure_extern("float32", "__expf", A[i]), name="B")
s = te.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
f = tvm.build(s, [A, B], "cuda", name="myexp")
print(f.imported_modules[0].get_source())<jupyter_output><empty_output><jupyter_text>Unified Intrinsic Call----------------------The above code verifies that direct external call can be used tocall into device specific functions.However, the above way only works for CUDA target with float type.Ideally, we want to write same code for any device and any data type.TVM intrinsic provides the user a mechanism to achieve this, and thisis the recommended way to solve the problem.The following code use te.exp instead, which create an intrinsic call:py::func:`tvm.te.exp` to do the exponential.<jupyter_code>n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: te.exp(A[i]), name="B")
s = te.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())<jupyter_output><empty_output><jupyter_text>We can find that the code works for both CUDA and opencl.The same te.exp can also be used for float64 data types.<jupyter_code>fopencl = tvm.build(s, [A, B], "opencl", name="myexp")
print(fopencl.imported_modules[0].get_source())<jupyter_output><empty_output><jupyter_text>Intrinsic Lowering Rule-----------------------When :py:func:`tvm.te.exp` is called, TVM creates an intrinsic Call Expr.TVM uses transformation rules to transform the intrinsiccall to device specific extern calls.TVM also allows user to customize the rules during runtime.The following example customizes CUDA lowering rule for :code:`exp`.<jupyter_code>def my_cuda_math_rule(op):
    """Customized CUDA intrinsic lowering rule"""
    assert isinstance(op, tvm.tir.Call)
    name = op.op.name
    assert name.startswith("tir.")
    dispatch_name = name[4:]
    if op.dtype == "float32":
        # call float function
        return tvm.tir.call_pure_extern("float32", "%sf" % dispatch_name, op.args[0])
    elif op.dtype == "float64":
        # call double function
        return tvm.tir.call_pure_extern("float32", dispatch_name, op.args[0])
    else:
        # cannot do translation, return self.
        return op


tvm.target.register_intrin_rule("cuda", "exp", my_cuda_math_rule, override=True)<jupyter_output><empty_output><jupyter_text>Register the rule to TVM with override option to override existing rule.Notice the difference between the printed code from previous one:our new rule uses math function :code:`expf` instead offast math version :code:`__expf`.<jupyter_code>fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())<jupyter_output><empty_output><jupyter_text>Add Your Own Intrinsic----------------------If there is an intrinsic that is not provided by TVM.User can easily add new intrinsic by using the intrinsic rule system.The following example add an intrinsic :code:`mylog` to the system.<jupyter_code>def mylog(x):
    """customized log intrinsic function"""
    return tvm.tir.call_intrin(x.dtype, "tir.mylog", x)


def my_cuda_mylog_rule(op):
    """CUDA lowering rule for log"""
    if op.dtype == "float32":
        return tvm.tir.call_pure_extern("float32", "logf", op.args[0])
    elif op.dtype == "float64":
        return tvm.tir.call_pure_extern("float64", "log", op.args[0])
    else:
        return op


# new op registration is triggered by registering an attribute of the op
tvm.ir.register_op_attr("tir.mylog", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
tvm.target.register_intrin_rule("cuda", "mylog", my_cuda_mylog_rule, override=True)

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: mylog(A[i]), name="B")
s = te.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
fcuda = tvm.build(s, [A, B], "cuda", name="mylog")
print(fcuda.imported_modules[0].get_source())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14725838987713186, 'fraction_numerical': 0.00715202640748212, 'mean_word_length': 4.967177242888402, 'pattern_counts': {'":': 4, '<': 30, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '42397', 'license': 'Apache-2.0', 'n_tokens_mistral': 1849, 'n_tokens_neox': 1786, 'n_words': 515, 'path': '_downloads/c61e59811b10bffc0715e199259f6c52/intrin_math.ipynb', 'repo_name': 'areusch/tvm-docs-refactor', 'size': '5458'}