Skip to content

Commit

Permalink
Handle vectorize for LE statement (apache#3137)
Browse files Browse the repository at this point in the history
* Handle vectorize for LE statement

Fix a new cases introduced by commit 7afbca5

* Add test
  • Loading branch information
wweic authored and Wei Chen committed May 13, 2019
1 parent 87374d1 commit f6c52fa
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
3 changes: 3 additions & 0 deletions src/pass/vectorize_loop.cc
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,9 @@ class Vectorizer : public IRMutator {
Expr Mutate_(const LT* op, const Expr &e) final {
return BinaryVec(op, e);
}
Expr Mutate_(const LE* op, const Expr &e) final {
return BinaryVec(op, e);
}
Expr Mutate_(const GT* op, const Expr &e) final {
return BinaryVec(op, e);
}
Expand Down
24 changes: 24 additions & 0 deletions tests/python/unittest/test_pass_vectorize.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,28 @@ def test_vectorize_with_if():
assert stmt.then_case.value.dtype == "float32x4"
assert isinstance(stmt.else_case, tvm.stmt.For)

def test_vectorize_with_le_cond():
n = tvm.var('n')
ib = tvm.ir_builder.create()
A = ib.pointer("float32", name="A")
with ib.for_range(0, 4, for_type="vectorize") as i:
with ib.if_scope(i <= n):
A[i] = A[i] + 1
stmt = ib.get()
stmt = tvm.ir_pass.VectorizeLoop(stmt)
assert isinstance(stmt, tvm.stmt.For)

def test_vectorize_with_ge_cond():
n = tvm.var('n')
ib = tvm.ir_builder.create()
A = ib.pointer("float32", name="A")
with ib.for_range(0, 4, for_type="vectorize") as i:
with ib.if_scope(i >= n):
A[i] = A[i] + 1
stmt = ib.get()
stmt = tvm.ir_pass.VectorizeLoop(stmt)
assert isinstance(stmt, tvm.stmt.For)

def test_vectorize_if_then_else():
n = tvm.var('n')
x = tvm.var('x')
Expand Down Expand Up @@ -102,3 +124,5 @@ def test_vectorize_if_then_else():
test_vectorize_with_if()
test_vectorize_loop()
test_vectorize_if_then_else()
test_vectorize_with_le_cond()
test_vectorize_with_ge_cond()

0 comments on commit f6c52fa

Please sign in to comment.