[TIR][Transform] HoistIfThenElse added

[ANSHUMAN87]

This is a follow up PR. Please refer #5559.

cc @kevinthesun , @roastduck , @zhiics , @junrushao1994 , @tqchen .

I have tried to cover all the possible scenarios. Please let me know in case i miss anything. TIA!

[roastduck]

Since nearly all programs will be touched by this pass, potential bugs in this pass would be critical. Could you (maybe temporarily) add this pass into the default building procedure and run all the tests? It would greatly reduce potential bugs.

[ANSHUMAN87]

Since nearly all programs will be touched by this pass, potential bugs in this pass would be critical. Could you (maybe temporarily) add this pass into the default building procedure and run all the tests? It would greatly reduce potential bugs.

@roastduck : Thanks a lot for your input!
I will definitely ensure your point during internal test.
Currently I am working on to resolve all the errors reported in CI test suites. Will update once all the issues are resolved.

[ANSHUMAN87]

All the CI issues are resolved now. Also my internal tests shows good result. I think we can start review now. TIA!

[roastduck]

LGTM

[tqchen]

cc @ZihengJiang @merrymercy @Hzfengsy @kevinthesun @junrushao1994 @spectrometerHBH @wpan11nv @kparzysz-quic please help to take a look

[kparzysz-quic]

This is a limited case of "loop unswitching". Please consider a more general solution, where

for (i = 0..N) {
  // statement A
  if (invariant-condition) {
    // statement B
  } else {
    // statement C
  }
  // statement D
}

is transformed into

if (invariant-condition) {
  for (i = 0..N) {
    // statement A
    // statement B
    // statement D
  }
} else {
  for (i = 0..N) {
    // statement A
    // statement C
    // statement D
  }
}

Using the same logic you could unswitch attribute statements, if needed.

[ANSHUMAN87]

This is a limited case of "loop unswitching". Please consider a more general solution, where

for (i = 0..N) {
  // statement A
  if (invariant-condition) {
    // statement B
  } else {
    // statement C
  }
  // statement D
}

is transformed into

if (invariant-condition) {
  for (i = 0..N) {
    // statement A
    // statement B
    // statement D
  }
} else {
  for (i = 0..N) {
    // statement A
    // statement C
    // statement D
  }
}

Using the same logic you could unswitch attribute statements, if needed.

Thanks for bringing up this point!
In fact this is on my top TODO list once current PR is merged.
Here it is not the question of general solution, it is about covering more scenarios.
Current changes covers most of the real-time scenarios (excluding hypothetical scenarios).
This scenario is currently left uncovered intentionally because of following challenges.

1. Invariant identification: What is the best way to identify the if condition variables are not updated inside the block ?
   For example in below case, how we can find out var n is invariant in the loop block:

    var n = 3
    for (i: int32, 0, n: int32) {
      if @tir.likely((n > 23), dtype=bool) {
        n = n + 2      
        data[i] = i
      } else {
        n = n + 3
        data[i] = i + 1
      }
    }

2. Scope identification: Need to scan each and every statement to do this, where logic becomes significantly more complex.

May be we can discuss more about first Challenge. Please suggest if anyone has any optimum solution for it.
Let me know in case the scenario is not clear. TIA!

[ANSHUMAN87]

Gentle ping @tqchen , @ZihengJiang @merrymercy @Hzfengsy @kevinthesun @junrushao1994 @spectrometerHBH @wpan11nv @kparzysz-quic !!!

Let us discuss and bring a conclusion to the open points / challenges mentioned in my previous comment. TIA!

[ANSHUMAN87]

Gentle ping @tqchen , @ZihengJiang @merrymercy @Hzfengsy @kevinthesun @junrushao1994 @spectrometerHBH @wpan11nv @kparzysz-quic !!!

Let us discuss and bring a conclusion to the open points / challenges mentioned in my previous comment. TIA!

also cc @MarisaKirisame !!!

[Hzfengsy]

LGTM

[MarisaKirisame]

0: I dont think a pass have to handle all cases for it to be merged. We can improve upon it incrementally.
1: About loop unswitching, detecting invariant is very hard. But we can simply approximate by detecting expr that didnt have any vars change in the loop. There is rice theorem, so all solution will not be perfect, and we should do the least amount of work that get us good enough performance.
2: you should just build a datastructure/helper function such as boundvar/freevar to help. I dont found dealing with scope too complex in the pass I had written. If you had more specific question/objection please bring it up.

[ANSHUMAN87]

0: I dont think a pass have to handle all cases for it to be merged. We can improve upon it incrementally.
1: About loop unswitching, detecting invariant is very hard. But we can simply approximate by detecting expr that didnt have any vars change in the loop. There is rice theorem, so all solution will not be perfect, and we should do the least amount of work that get us good enough performance.
2: you should just build a datastructure/helper function such as boundvar/freevar to help. I dont found dealing with scope too complex in the pass I had written. If you had more specific question/objection please bring it up.

Thanks @MarisaKirisame for your enlightening response.

I am in total agreement with all your points. :)

[ANSHUMAN87]

I think we can go ahead and merge the PR, as most of the scenarios are handled.
I will keep loop unswitching as open point, once it is handled will raise new PR for it.
And we can keep adding new useful scenarios, as and when someone reports it.

@tqchen : Would you please share your opinion on this. TIA!

[tqchen]

Need explicit approval https://tvm.apache.org/docs/contribute/code_review.html#approve-and-request-changes-explicitly from @kparzysz-quic @MarisaKirisame .

The most important thing is the code clearity(others can understand the logic) and correctness.

cc @ZihengJiang @junrushao1994 @vinx13 @zhiics it would be great if you can also take a look

[kparzysz-quic]

I think this is ok for now.

[kevinthesun]

I'll take a look in the next few days.

[kevinthesun]

One thing I think it is good to have in this PR is to get some benchmark data, since we now enable this pass by default. IMO this pass is especially valuable when tackling dynamic kernels in GPU which introduces a lot of branchings. It's great to see how much performance improvement we can have. For CPU, we need to make sure this pass doesn't introduce regression for some common workloads, such as resnet.

[ANSHUMAN87]

@kevinthesun : Thanks a lot for your input! I believe all the cases which are covered now, does not degrade performance in any case either CPU or GPU :)
Do you have any suggestion how to obtain the benchmark data(like some test case or some existing tools) ?

[MarisaKirisame]

Everything is green, but since @kevinthesun want to review it I will wait to merge for a few day.

[ANSHUMAN87]

One thing I think it is good to have in this PR is to get some benchmark data, since we now enable this pass by default. IMO this pass is especially valuable when tackling dynamic kernels in GPU which introduces a lot of branchings. It's great to see how much performance improvement we can have. For CPU, we need to make sure this pass doesn't introduce regression for some common workloads, such as resnet.

@kevinthesun : I have verified the inference time for Resnet50 on CPU. There is no performance impact. In fact i did not find anything as Hoisting candidate.

Hoisting Disabled :

Hoisting Enabled:

Hope it helps. Please let me know, if i have mistaken anything. TIA!

[kevinthesun]

One thing I think it is good to have in this PR is to get some benchmark data, since we now enable this pass by default. IMO this pass is especially valuable when tackling dynamic kernels in GPU which introduces a lot of branchings. It's great to see how much performance improvement we can have. For CPU, we need to make sure this pass doesn't introduce regression for some common workloads, such as resnet.

@kevinthesun : I have verified the inference time for Resnet50 on CPU. There is no performance impact. In fact i did not find anything as Hoisting candidate.

Hoisting Disabled :

Hoisting Enabled:

Hope it helps. Please let me know, if i have mistaken anything. TIA!

Usually there are two cases which might involve this pass: 1) Loop tiling with non-factor split. 2) Dynamic shape op. If I remember correctly, a conv2d with symbolic batch size will generate an IR with a lot of hoist candidates. Due to the limitation of nvcc, performance for such as kernel is quite terrible and this pass is able to handle this. It would be nice if we can verify this case for GPU and add a unit test.

[ANSHUMAN87]

@kevinthesun : Thanks for your response!
I am familiar with the scenario you are referring for conv op. I agree to the possibilities of hoisting in that case.
The same concern was raised by @roastduck too, i believe.
Unfortunately due to some intermediate pass dependency on the positioning of For and If node, i temporarily disabled this hoisting scenario with Attr nodes(IterVar & GlobalVar).

I think it would be better, i enable this hoisting scenario too and move the position of Hoisting Pass to the end of list during lowering. May be after that we can take the performance data.

Please let me know your opinion on this. TIA!

[kevinthesun]

@ANSHUMAN87 Thanks for clarification. Though we might not need to do so in this PR, it would be great if we can bring this in since AFAIK gpu is the major case for this pass in tvm.

[kevinthesun]

LGTM

[MarisaKirisame]

Thanks @ANSHUMAN87 @roastduck @kevinthesun @kparzysz-quic @hzfan

[ANSHUMAN87]

Thanks a lot @MarisaKirisame , @roastduck @kevinthesun @kparzysz-quic @Hzfengsy, @tqchen !

This PR has 2 open points as per discussion with all the members participated.

Summarizing as below:

Open point 1:- Support hoisting for Attr Nodes with IterVar & GlobalVar(GPU kernel case).
Open point 2:- Support for loop unswitching case.

I have kept it on my TODO list. Will ensure support of these cases in my future PRs.
However i request to all if anyone finds any more scenarios which need to be covered, please raise an issue and tag me, so that i can handle that too. TIA!