Skip to content

Commit

Permalink
[Bugfix] Fix div zero error in rewrite_simplify (#8961)
Browse files Browse the repository at this point in the history
* fix div zero error in rewrite_simplify

* update the style to fix ci error

* remove useless code and comment
  • Loading branch information
Sen Yang authored Sep 9, 2021
1 parent b8fcad8 commit e47fc6a
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 16 deletions.
1 change: 1 addition & 0 deletions src/arith/rewrite_simplify.cc
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const DivNode* op) {
if ((div(ramp(b1, c1, lanes), broadcast(c2, lanes))).Match(ret)) {
int64_t c1val = c1.Eval()->value;
int64_t c2val = c2.Eval()->value;
ICHECK(c2val != 0) << "division by zero";
if (c1val % c2val == 0) {
return ramp(div(b1, c2), div(c1, c2), lanes).Eval();
}
Expand Down
26 changes: 10 additions & 16 deletions tests/python/unittest/test_arith_rewrite_simplify.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import pytest
import tvm
from tvm import te

Expand Down Expand Up @@ -931,20 +932,13 @@ def test_shift_left_simplify():
ck.verify(z, tvm.tir.const(1 << 10, "int32"))


def test_div_zero_simplify():
ck = RewriteChecker()

with pytest.raises(tvm.error.TVMError) as cm:
ck.analyzer.rewrite_simplify(tvm.tir.Div(tvm.tir.Ramp(1, 1, 2), tvm.tir.Broadcast(0, 2)))
assert "division by zero" in str(cm.execption)


if __name__ == "__main__":
test_floordiv_index_simplify()
test_floormod_index_simplify()
test_cmp_simplify()
test_vector_simplify()
test_add_index_simplify()
test_sub_index_simplify()
test_mul_index_simplify()
test_div_index_simplify()
test_max_index_simplify()
test_min_index_simplify()
test_mod_index_simplify()
test_select_simplify()
test_logical_simplify()
test_let_simplify()
test_cast_simplify()
test_shift_left_simplify()
pytest.main([__file__])

0 comments on commit e47fc6a

Please sign in to comment.