This is an automated email from the ASF dual-hosted git repository.
tlopex 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 bfa07828a6 [BugFix][Relax] Add legalize for isnan, isinf, isfinite
(#19492)
bfa07828a6 is described below
commit bfa07828a6ee0faa427dc1e3199d793016fb8910
Author: as4230 <[email protected]>
AuthorDate: Sat May 2 05:56:33 2026 -0400
[BugFix][Relax] Add legalize for isnan, isinf, isfinite (#19492)
The Relax ops `relax.isnan`, `relax.isinf`, and `relax.isfinite` are
registered in the op registry and emit valid IR, but lack FLegalize
entries so LegalizeOps() leaves them unlowered and relax.build() crashes
at codegen with:
InternalError: CodeGenVM cannot handle this intrinsic now:
Op(relax.isnan)
Same crash on `isinf` and `isfinite`, on both LLVM and CUDA targets.
The TOPI implementations already exist (python/tvm/topi/math.py). The
fix is three register_legalize calls following the
_call_topi_without_attr pattern used by the other unary ops in the same
file.
Fixes #19452.
---
python/tvm/relax/transform/legalize_ops/unary.py | 3 +++
tests/python/relax/test_transform_legalize_ops_unary.py | 3 +++
2 files changed, 6 insertions(+)
diff --git a/python/tvm/relax/transform/legalize_ops/unary.py
b/python/tvm/relax/transform/legalize_ops/unary.py
index fc6b96a438..4d09c6d61c 100644
--- a/python/tvm/relax/transform/legalize_ops/unary.py
+++ b/python/tvm/relax/transform/legalize_ops/unary.py
@@ -38,6 +38,9 @@ register_legalize("relax.cos",
_call_topi_without_attr(topi.cos, "tir_cos"))
register_legalize("relax.cosh", _call_topi_without_attr(topi.cosh, "tir_cosh"))
register_legalize("relax.exp", _call_topi_without_attr(topi.exp, "tir_exp"))
register_legalize("relax.floor", _call_topi_without_attr(topi.floor,
"tir_floor"))
+register_legalize("relax.isfinite", _call_topi_without_attr(topi.isfinite,
"tir_isfinite"))
+register_legalize("relax.isinf", _call_topi_without_attr(topi.isinf,
"tir_isinf"))
+register_legalize("relax.isnan", _call_topi_without_attr(topi.isnan,
"tir_isnan"))
register_legalize("relax.log", _call_topi_without_attr(topi.log, "tir_log"))
register_legalize("relax.logical_not",
_call_topi_without_attr(topi.logical_not, "tir_logical_not"))
register_legalize("relax.negative", _call_topi_without_attr(topi.negative,
"tir_negative"))
diff --git a/tests/python/relax/test_transform_legalize_ops_unary.py
b/tests/python/relax/test_transform_legalize_ops_unary.py
index e32f95a3c4..4f8ee67d99 100644
--- a/tests/python/relax/test_transform_legalize_ops_unary.py
+++ b/tests/python/relax/test_transform_legalize_ops_unary.py
@@ -89,6 +89,9 @@ def _test_symbolic_shape(name: str, relax_op: Callable,
te_func: Callable, dtype
("exp", R.exp, topi.exp, "float32"),
("floor", R.floor, topi.floor, "float32"),
("floor", R.floor, topi.identity, "int32"),
+ ("isfinite", R.isfinite, topi.isfinite, "float32"),
+ ("isinf", R.isinf, topi.isinf, "float32"),
+ ("isnan", R.isnan, topi.isnan, "float32"),
("log", R.log, topi.log, "float32"),
("negative", R.negative, topi.negative, "float32"),
("round", R.round, topi.round, "float32"),
