Expose Missing TIR Builtins to Python (#12466)

This PR exposes the following TIR operation in python:

`address_of`: tested [here](https://github.com/apache/tvm/blob/d2f9f254d275df256dbcbc5a9f8b3a07cee1d81f/tests/python/unittest/test_tvmscript_roundtrip.py#L3247)
`lookup_param`: tested [here](https://github.com/apache/tvm/blob/d2f9f254d275df256dbcbc5a9f8b3a07cee1d81f/tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py#L171)
`infinity`: add new unittest
`reinterpret`: tested [here](https://github.com/apache/tvm/blob/d2f9f254d275df256dbcbc5a9f8b3a07cee1d81f/tests/python/unittest/test_tvmscript_roundtrip.py#L2991)
`isnullptr`: tested [here](https://github.com/apache/tvm/blob/d2f9f254d275df256dbcbc5a9f8b3a07cee1d81f/tests/python/unittest/test_tvmscript_roundtrip.py#L260)

Co-Authored-By: yongwww <[email protected]>

Author: cyx-6

include/tvm/tir/op.h

@@ -701,6 +701,7 @@ TVM_DECLARE_INTRIN_UNARY(rsqrt);
 TVM_DECLARE_INTRIN_UNARY(log);
 TVM_DECLARE_INTRIN_UNARY(log2);
 TVM_DECLARE_INTRIN_UNARY(log10);
+TVM_DECLARE_INTRIN_UNARY(log1p);
 TVM_DECLARE_INTRIN_UNARY(popcount);
 TVM_DECLARE_INTRIN_UNARY(tan);
 TVM_DECLARE_INTRIN_UNARY(cos);

python/tvm/tir/__init__.py

@@ -50,14 +50,15 @@
 from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any, min_value, max_value, trace
 from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
 from .op import tvm_tuple, tvm_struct_get, tvm_struct_set
-from .op import assume, undef
+from .op import address_of, lookup_param, assume, undef
+from .op import infinity, reinterpret
 from .op import exp, exp2, exp10, log, log2, log10, log1p, ldexp, clz
 from .op import sin, sinh, asin, asinh
 from .op import cos, cosh, acos, acosh
 from .op import tan, tanh, atan, atan2, atanh
 from .op import erf, sigmoid, sqrt, rsqrt, floor, ceil, hypot
 from .op import trunc, abs, round, nextafter, nearbyint, power, popcount, fmod, if_then_else
-from .op import likely, isnan, isfinite, isinf, copysign
+from .op import likely, isnan, isnullptr, isfinite, isinf, copysign
 from .op import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod, ceildiv
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift

python/tvm/tir/op.py

@@ -466,6 +466,44 @@ def tvm_struct_set(arr, index, field, value):
     return call_intrin("handle", "tir.tvm_struct_set", arr, index, field, value)
 
 
+def address_of(buffer_load, span=None):
+    """Returns the address of an element in the buffer
+
+    Parameters
+    ----------
+    buffer_load: BufferLoad
+        The buffer load.
+
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    call : PrimExpr
+        The call expression.
+    """
+    return call_intrin("handle", "tir.address_of", buffer_load, span=span)
+
+
+def lookup_param(param_name, span=None):
+    """Returns the param by name
+
+    Parameters
+    ----------
+    param_name : str
+        The name of param.
+
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    call : PrimExpr
+        The call expression.
+    """
+    return call_intrin("handle", "tir.lookup_param", param_name, span=span)
+
+
 def ret(val):
     """Create a tir return expression
 
@@ -610,6 +648,47 @@ def max_value(dtype: str, span: Optional[Span] = None) -> Any:
     return _ffi_api.max_value(dtype, span)  # type: ignore
 
 
+def infinity(dtype: str, span: Optional[Span] = None) -> Any:
+    """infinity value of dtype
+
+    Parameters
+    ----------
+    dtype : str
+        The data type.
+
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    value : tvm.Expr
+        The infinity value of dtype.
+    """
+    return _ffi_api.infinity(dtype, span)  # type: ignore
+
+
+def reinterpret(dtype, value) -> Any:
+    """infinity value of dtype
+
+    Parameters
+    ----------
+    dtype : str
+        The data type.
+
+    value : PrimExpr
+        The input value.
+
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    value : tvm.Expr
+        The reinterpret cast value of dtype.
+    """
+    return call_intrin(dtype, "tir.reinterpret", value)
+
+
 def exp(x):
     """Take exponential of input x.
 
@@ -1253,6 +1332,25 @@ def isnan(x, span=None):
     return _ffi_api.isnan(x, span)  # type: ignore
 
 
+def isnullptr(x, span=None):
+    """Check if input value is nullptr.
+
+    Parameters
+    ----------
+    x : PrimExpr
+        Input argument.
+
+    span : Optional[Span]
+        The location of this operator in the source code.
+
+    Returns
+    -------
+    y : PrimExpr
+        The result.
+    """
+    return call_intrin("bool", "tir.isnullptr", x, span=span)  # type: ignore
+
+
 def isfinite(x, span=None):
     """Check if input value is finite.
 

src/tir/op/op.cc

@@ -929,6 +929,8 @@ TVM_REGISTER_GLOBAL("tir.min_value").set_body_typed(min_value);
 
 TVM_REGISTER_GLOBAL("tir.max_value").set_body_typed(max_value);
 
+TVM_REGISTER_GLOBAL("tir.infinity").set_body_typed(infinity);
+
 TVM_REGISTER_GLOBAL("tir.abs").set_body_typed(tvm::abs);
 
 TVM_REGISTER_GLOBAL("tir.likely").set_body_typed(tvm::likely);

tests/python/unittest/test_tir_nodes.py

@@ -301,6 +301,12 @@ def test_divide_by_zero():
             pass
 
 
+def test_infinity():
+    assert str(tvm.tir.infinity("float16")) == "inff16"
+    assert str(tvm.tir.infinity("float32")) == "inff32"
+    assert str(tvm.tir.infinity("float64")) == "inff64"
+
+
 def test_isnan():
     x = te.var("x", "float32")
     assert str(tvm.tir.isnan(x)) == "@tir.isnan(x: float32, dtype=bool)"

tests/python/unittest/test_tir_op_types.py

@@ -39,6 +39,29 @@ def test_tir_op_tvm_struct_set():
     assert expr.op.name == "tir.tvm_struct_set"
 
 
+def test_tir_op_address_of():
+    buffer = tir.decl_buffer((128), "float32")
+    expr = tir.address_of(buffer[0])
+    assert expr.op.name == "tir.address_of"
+
+
+def test_tir_op_lookup_param():
+    expr = tir.lookup_param("p0")
+    assert expr.op.name == "tir.lookup_param"
+
+
+def test_tir_op_reinterpret():
+    x = tir.Var("x", dtype="int32")
+    expr = tir.reinterpret("float32", x)
+    assert expr.op.name == "tir.reinterpret"
+
+
+def test_tir_op_isnullptr():
+    x = tir.Var("x", dtype="int32")
+    expr = tir.isnullptr(x)
+    assert expr.op.name == "tir.isnullptr"
+
+
 def test_tir_op_call_assume():
     x = tir.Var("x", dtype="int32")
     expr = tir.assume(cond=x)
@@ -60,6 +83,10 @@ def test_tir_op_call_likely():
     test_tir_op_tvm_tuple()
     test_tir_op_tvm_struct_get()
     test_tir_op_tvm_struct_set()
+    test_tir_op_address_of()
+    test_tir_op_lookup_param()
+    test_tir_op_reinterpret()
+    test_tir_op_isnullptr()
     test_tir_op_call_assume()
     test_tir_op_call_undef()
     test_tir_op_call_likely()