basic types and exprs (apache#41)

Author: cyx-6

include/tvm/tir/op.h

@@ -796,6 +796,15 @@ TVM_DLL PrimExpr min(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr>
 TVM_DLL PrimExpr prod(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
                       Span span = Span());
 
+/*!
+ * \brief Calculate fmod(x, y)
+ * \param x Left operand.
+ * \param y Right operand.
+ * \param span The location of this operation in the source.
+ * \return The result expression.
+ */
+TVM_DLL PrimExpr fmod(PrimExpr x, PrimExpr y, Span span = Span());
+
 /*!
  * \brief Calculate floor(x)
  * \param x The input expression.
@@ -896,6 +905,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/script/builder/_ffi_api.py

@@ -17,4 +17,4 @@
 """FFI APIs for tvm.script.builder"""
 import tvm._ffi
 
-tvm._ffi._init_api("script.builder", __name__) # pylint: disable=protected-access
+tvm._ffi._init_api("script.builder", __name__)  # pylint: disable=protected-access

python/tvm/script/builder/tir/__init__.py

@@ -31,3 +31,4 @@
 )
 from .prim_func_frame import arg, prim_func
 from .var import Buffer
+from .op import *

python/tvm/script/builder/tir/_ffi_api.py

@@ -17,4 +17,4 @@
 """FFI APIs for tvm.script.builder.tir"""
 import tvm._ffi
 
-tvm._ffi._init_api("script.builder.tir", __name__) # pylint: disable=protected-access
+tvm._ffi._init_api("script.builder.tir", __name__)  # pylint: disable=protected-access

python/tvm/script/builder/tir/op.py

@@ -0,0 +1,161 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""TVM Script TIR Op"""
+
+from . import _ffi_api
+
+
+from tvm.tir.op import abs, popcount, nextafter, copysign, fmod
+from tvm.tir.op import (
+    floor,
+    floordiv,
+    floormod,
+    ceil,
+    round,
+    trunc,
+    truncdiv,
+    truncmod,
+    nearbyint,
+)
+from tvm.tir.op import (
+    hypot,
+    ldexp,
+    power,
+    exp,
+    exp2,
+    exp10,
+    erf,
+    sqrt,
+    rsqrt,
+    log,
+    log2,
+    log10,
+    log1p,
+    sigmoid,
+)
+from tvm.tir.op import isnan, isfinite, isinf
+from tvm.tir.op import cos, cosh, sin, sinh, tan, tanh
+from tvm.tir.op import acos, acosh, asin, asinh, atan, atanh
+from tvm.tir.op import atan2, clz, comm_reducer, infinity, reinterpret
+from tvm.tir.op import min_value, max_value, if_then_else
+from tvm.tir.op import call_packed, call_extern
+from tvm.tir.expr import Select, Ramp, Broadcast, Shuffle
+from tvm.tir.generic import cast
+
+
+def boolean(expr):
+    return _ffi_api.PrimType("bool", expr)
+
+
+def int8(expr):
+    return _ffi_api.PrimType("int8", expr)
+
+
+def int16(expr):
+    return _ffi_api.PrimType("int16", expr)
+
+
+def int32(expr):
+    return _ffi_api.PrimType("int32", expr)
+
+
+def int64(expr):
+    return _ffi_api.PrimType("int64", expr)
+
+
+def uint8(expr):
+    return _ffi_api.PrimType("uint8", expr)
+
+
+def uint16(expr):
+    return _ffi_api.PrimType("uint16", expr)
+
+
+def uint32(expr):
+    return _ffi_api.PrimType("uint32", expr)
+
+
+def uint64(expr):
+    return _ffi_api.PrimType("uint64", expr)
+
+
+def float8(expr):
+    return _ffi_api.PrimType("float8", expr)
+
+
+def float16(expr):
+    return _ffi_api.PrimType("float16", expr)
+
+
+def float32(expr):
+    return _ffi_api.PrimType("float32", expr)
+
+
+def float64(expr):
+    return _ffi_api.PrimType("float64", expr)
+
+
+def min(a, b, span=None):
+    """Compute the minimum value of two expressions.
+
+    Parameters
+    ----------
+    a : PrimExpr
+        The left hand operand
+
+    b : PrimExpr
+        The right hand operand
+
+    span : Optional[Span]
+        The location of this operator in the source.
+
+    Returns
+    -------
+    res : PrimExpr
+        The result expression.
+
+    Note
+    ----
+    This is the default integer division behavior in C.
+    """
+    return _ffi_api.min(a, b, span)  # type: ignore
+
+
+def max(a, b, span=None):
+    """Compute the maximum value of two expressions.
+
+    Parameters
+    ----------
+    a : PrimExpr
+        The left hand operand
+
+    b : PrimExpr
+        The right hand operand
+
+    span : Optional[Span]
+        The location of this operator in the source.
+
+    Returns
+    -------
+    res : PrimExpr
+        The result expression.
+
+    Note
+    ----
+    This is the default integer division behavior in C.
+    """
+    return _ffi_api.max(a, b, span)  # type: ignore

python/tvm/script/builder/tir/var.py

@@ -20,10 +20,12 @@
 from . import _ffi_api
 
 
-def Buffer( # pylint: disable=invalid-name
+def Buffer(  # pylint: disable=invalid-name
     shape,
     dtype,
     name="buffer",
     storage_scope="",
 ) -> tir.Buffer:
-    return _ffi_api.Buffer(shape, dtype, name, storage_scope) # pylint: disable=no-member # type: ignore
+    return _ffi_api.Buffer(
+        shape, dtype, name, storage_scope
+    )  # pylint: disable=no-member # type: ignore

python/tvm/tir/__init__.py

@@ -44,17 +44,36 @@
 
 from .function import PrimFunc, TensorIntrin, IndexMap
 
-from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_extern
-from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any, min_value, max_value, trace
+from .op import call_packed, call_intrin, call_pure_extern, call_extern
+from .op import (
+    call_llvm_intrin,
+    call_llvm_pure_intrin,
+    ret,
+    all,
+    any,
+    min_value,
+    max_value,
+    trace,
+)
 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 (
+    trunc,
+    abs,
+    round,
+    nextafter,
+    nearbyint,
+    power,
+    popcount,
+    fmod,
+    if_then_else,
+)
 from .op import isnan, isfinite, isinf, copysign
 from .op import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod
-from .op import comm_reducer, min, max, sum
+from .op import comm_reducer, min, max, sum, infinity, reinterpret
 from .op import q_multiply_shift
 
 from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError

python/tvm/tir/op.py

@@ -151,7 +151,10 @@ def call_pure_extern(dtype, func_name, *args, span=None):
         The call expression.
     """
     return Call(
-        dtype, Op.get("tir.call_pure_extern"), convert((StringImm(func_name),) + args), span
+        dtype,
+        Op.get("tir.call_pure_extern"),
+        convert((StringImm(func_name),) + args),
+        span,
     )
 
 
@@ -178,7 +181,10 @@ def call_extern(dtype, func_name, *args, span=None):
         The call expression.
     """
     return Call(
-        dtype, Op.get("tir.call_extern"), convert((StringImm(func_name),) + args), span=span
+        dtype,
+        Op.get("tir.call_extern"),
+        convert((StringImm(func_name),) + args),
+        span=span,
     )
 
 
@@ -210,7 +216,11 @@ def call_llvm_intrin(dtype, name, *args, span=None):
     llvm_id = codegen.llvm_lookup_intrinsic_id(name)
     assert llvm_id != 0, "%s is not an LLVM intrinsic" % name
     return call_intrin(
-        dtype, Op.get("tir.call_llvm_intrin"), tvm.tir.const(llvm_id, "uint32"), *args, span=span
+        dtype,
+        Op.get("tir.call_llvm_intrin"),
+        tvm.tir.const(llvm_id, "uint32"),
+        *args,
+        span=span,
     )
 
 
@@ -394,6 +404,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, span=None) -> 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 _ffi_api.reinterpret(dtype, value, span)  # type: ignore
+
+
 def exp(x):
     """Take exponential of input x.
 

src/script/builder/tir/op.cc

@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include "./op.h"
+
+namespace tvm {
+namespace script {
+namespace builder {
+namespace tir {
+
+PrimExpr prim_type(String type_name, PrimExpr expr) {
+  return cast(DataType(runtime::String2DLDataType(type_name)), expr);
+}
+
+TVM_REGISTER_GLOBAL("script.builder.tir.PrimType").set_body_typed(prim_type);
+TVM_REGISTER_GLOBAL("script.builder.tir.min").set_body_typed([](PrimExpr a, PrimExpr b, Span span) {
+  return tvm::min(a, b, span);
+});
+TVM_REGISTER_GLOBAL("script.builder.tir.max").set_body_typed([](PrimExpr a, PrimExpr b, Span span) {
+  return tvm::max(a, b, span);
+});
+
+}  // namespace tir
+}  // namespace builder
+}  // namespace script
+}  // namespace tvm