You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2022/08/16 07:35:39 UTC
[tvm] branch main updated: [TIR] Expose Misc TIR operations to python (#12435)
This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c477c763c3 [TIR] Expose Misc TIR operations to python (#12435)
c477c763c3 is described below
commit c477c763c37adf29b34528ca52d231d622719b3e
Author: Yaxing Cai <ca...@gmail.com>
AuthorDate: Tue Aug 16 00:35:31 2022 -0700
[TIR] Expose Misc TIR operations to python (#12435)
This PR exposes the following TIR operation in python:
- `assume`: tested [here](https://github.com/apache/tvm/blob/bb513866ad70fa20eb0c37ca339d330d6a76c747/tests/python/unittest/test_tir_transform_remove_assume.py#L34)
- `undef`: tested [here](https://github.com/apache/tvm/blob/bb513866ad70fa20eb0c37ca339d330d6a76c747/tests/python/unittest/test_tir_transform_remove_undef.py#L63)
- `likely`: tested [here](https://github.com/apache/tvm/blob/bb513866ad70fa20eb0c37ca339d330d6a76c747/tests/python/unittest/test_tir_schedule_compute_at.py#L849)
Co-Authored-By: yongwww <yo...@gmail.com>
---
python/tvm/tir/__init__.py | 3 +-
python/tvm/tir/op.py | 47 ++++++++++++++++++++++++++++++
src/tir/op/op.cc | 2 ++
tests/python/unittest/test_tir_op_types.py | 42 ++++++++++++++++++++++++++
4 files changed, 93 insertions(+), 1 deletion(-)
diff --git a/python/tvm/tir/__init__.py b/python/tvm/tir/__init__.py
index 77a27cd63e..f2bde89a21 100644
--- a/python/tvm/tir/__init__.py
+++ b/python/tvm/tir/__init__.py
@@ -49,13 +49,14 @@ from .op import call_packed_lowered, call_cpacked_lowered
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 tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
+from .op import assume, undef
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 isnan, isfinite, isinf, copysign
+from .op import likely, isnan, 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
diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
index 91fdd15448..b5f14cdc61 100644
--- a/python/tvm/tir/op.py
+++ b/python/tvm/tir/op.py
@@ -373,6 +373,33 @@ def tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset):
)
+def assume(cond=None):
+ """Provide a true statement that can be used for simplifications
+
+ Parameters
+ ----------
+ cond : Expr
+ The constraint condition.
+
+ Returns
+ -------
+ call : PrimExpr
+ The call expression.
+ """
+ return call_intrin("int32", "tir.assume", cond)
+
+
+def undef():
+ """Returns an initialized but arbitrary value
+
+ Returns
+ -------
+ call : PrimExpr
+ The call expression.
+ """
+ return call_intrin("int32", "tir.undef")
+
+
def ret(val):
"""Create a tir return expression
@@ -1121,6 +1148,26 @@ def ldexp(x1, x2):
return call_intrin(x1.dtype, "tir.ldexp", x1, x2) # type: ignore
+def likely(cond, span=None):
+ """Mark condition as likely.
+
+ Parameters
+ ----------
+
+ cond : PrimExpr
+ Input argument.
+
+ span : Optional[Span]
+ The location of this operator in the source code.
+
+ Returns
+ -------
+ y : PrimExpr
+ The marked expression.
+ """
+ return _ffi_api.likely(cond, span) # type: ignore
+
+
def isnan(x, span=None):
"""Check if input value is Nan.
diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
index 114571218b..7879c9fee9 100644
--- a/src/tir/op/op.cc
+++ b/src/tir/op/op.cc
@@ -931,6 +931,8 @@ TVM_REGISTER_GLOBAL("tir.max_value").set_body_typed(max_value);
TVM_REGISTER_GLOBAL("tir.abs").set_body_typed(tvm::abs);
+TVM_REGISTER_GLOBAL("tir.likely").set_body_typed(tvm::likely);
+
TVM_REGISTER_GLOBAL("tir.isnan").set_body_typed(tvm::isnan);
TVM_REGISTER_GLOBAL("tir.isfinite").set_body_typed(tvm::isfinite);
diff --git a/tests/python/unittest/test_tir_op_types.py b/tests/python/unittest/test_tir_op_types.py
new file mode 100644
index 0000000000..e805d9dea9
--- /dev/null
+++ b/tests/python/unittest/test_tir_op_types.py
@@ -0,0 +1,42 @@
+# 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.
+# pylint: disable=missing-docstring
+import tvm
+from tvm import tir
+
+
+def test_tir_op_call_assume():
+ x = tir.Var("x", dtype="int32")
+ expr = tir.assume(cond=x)
+ assert expr.op.name == "tir.assume"
+
+
+def test_tir_op_call_undef():
+ expr = tir.undef()
+ assert expr.op.name == "tir.undef"
+
+
+def test_tir_op_call_likely():
+ x = tir.Var("x", dtype="int32")
+ expr = tir.likely(cond=x)
+ assert expr.op.name == "tir.likely"
+
+
+if __name__ == "__main__":
+ test_tir_op_call_assume()
+ test_tir_op_call_undef()
+ test_tir_op_call_likely()