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: us...@infra.apache.org