roastduck opened a new pull request #5465:
URL: https://github.com/apache/incubator-tvm/pull/5465
I added a new pass to convert `tvm_if_then_else` intrinsics to `IfThenElse`
statments. For example:
```c++
a[i] = tvm.te.if_then_else(x < y, x, y)
```
is now
```c++
if (x < y) {
a[i] = x;
} else {
a[i] = y;
}
```
Rationale:
- `tvm_if_then_else` intrinsics can now enjoy the optimizations of the
`LoopPartition` pass after tagged as `likely`. (My condition-heavy sparse
application gains up to 2.85x speedup)
- Code lines containing `tvm_if_then_else` is often too long to read. This
improves readability.
Changes:
- Add an `IfThenElseIntrinToStmt` pass, which runs just before
`LoopPartition`. (**Question 1: Should I run this pass with `simple_mode` on?**
Currently no. I don't know if the users like the IR to be more similar to what
they write, or the lines to be not too long)
- Added a test for the pass.
- Enable access to the `likely` node in the Python API.
- Fixed the breaking `ThreadSync` pass since it did not allow synchronizing
threads in a conditional branch, which may be legal when there are
synchronizations in both the then-case and the else-case, like the code piece
below. Currently I simply deleted the branch check. (**Question 2: Any better
way?**)
```cuda
if (...) {
__syncthread();
a[i] = x;
} else {
__syncthread();
a[i] = y;
}
```
- Fixed the breaking `LowerWarpMemory` pass by adding an extra branching,
which handles the situation where there are only allocates of but no stores
into a warp memory.
Algorithm for `IfThenElseIntrinToStmt`:
1. Find a statement.
2. Find the out-most `tvm_if_then_else` intrinsic in it.
3. Put the statement into a new `IfThenElse` statement, and remove the
out-most `tvm_if_then_else` intrinsic.
4. Recurse until no more `tvm_if_then_else`.
Requesting reviews and thanks in advance. @tqchen @ZihengJiang @ajtulloch.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
