Skip to content

Commit

Permalink
[Bugfix] Fix other div zero errors also in rewrite_simplify (apache#8983
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

* fix div zero error of mod, floordiv, floormod in rewrite_simplify

* rewrite the test case of divison by zero to fix ci error

* remove useless tab

* retrigger ci

* remove useless blank to retrigger ci
  • Loading branch information
Sen Yang authored and ylc committed Jan 13, 2022
1 parent f9c6e8c commit 68b4c2f
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 1 deletion.
3 changes: 3 additions & 0 deletions src/arith/rewrite_simplify.cc
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const ModNode* op) {
if (truncmod(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 broadcast(truncmod(b1, c2), lanes).Eval();
}
Expand Down Expand Up @@ -724,6 +725,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) {
if (floordiv(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(floordiv(b1, c2), floordiv(c1, c2), lanes).Eval();
}
Expand Down Expand Up @@ -852,6 +854,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) {
if (floormod(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 broadcast(floormod(b1, c2), lanes).Eval();
}
Expand Down
16 changes: 15 additions & 1 deletion tests/python/unittest/test_arith_rewrite_simplify.py
Original file line number Diff line number Diff line change
Expand Up @@ -934,9 +934,23 @@ def test_shift_left_simplify():

def test_div_zero_simplify():
ck = RewriteChecker()
ramp = tvm.tir.Ramp(1, 1, 2)
broadcast = tvm.tir.Broadcast(0, 2)

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)))
ck.analyzer.rewrite_simplify(tvm.tir.Div(ramp, broadcast))
assert "division by zero" in str(cm.execption)

with pytest.raises(tvm.error.TVMError) as cm:
ck.analyzer.rewrite_simplify(tvm.tir.Mod(ramp, broadcast))
assert "division by zero" in str(cm.execption)

with pytest.raises(tvm.error.TVMError) as cm:
ck.analyzer.rewrite_simplify(tvm.tir.FloorDiv(ramp, broadcast))
assert "division by zero" in str(cm.execption)

with pytest.raises(tvm.error.TVMError) as cm:
ck.analyzer.rewrite_simplify(tvm.tir.FloorMod(ramp, broadcast))
assert "division by zero" in str(cm.execption)


Expand Down

0 comments on commit 68b4c2f

Please sign in to comment.