[SimpleLoopUnswitch] Port partially invariant unswitch from LoopUnswitch to SimpleLoopUnswitch

[jaykang10]

Summary

The partially invariant unswitch has been implemented on LoopUnswitch. https://reviews.llvm.org/D93764 We need to port the feature to SimpleLoopUnswitch for new pass manager. It is related to https://bugs.llvm.org/show_bug.cgi?id=49128

[fhahn]

Thanks for putting up this patch! Would it be possible to move the detection code (hasPartialIVCondition) somewhere, so both version of unswitching can use the same code?

[jaykang10]

In D99354#2652524, @fhahn wrote:

Thanks for putting up this patch! Would it be possible to move the detection code (hasPartialIVCondition) somewhere, so both version of unswitching can use the same code?

Yep, I also wanted to share it so I have tried to keep the existing implementation as much as possible... I am not sure which place is good for this... maybe, somewhere in lib/Transform/Utils?

[jaykang10]

Moved hasPartialIVCondition to LoopUtils.h following comment of @fhahn

[lebedev.ri]

Thank you for looking into this!
Unhelpful comment: have you considered splitting this into two patches, moving code from LoopUnswitch to a common place, and enhancing SimpleLoopUnswitch ?

[jaykang10]

In D99354#2655155, @lebedev.ri wrote:

Thank you for looking into this!
Unhelpful comment: have you considered splitting this into two patches, moving code from LoopUnswitch to a common place, and enhancing SimpleLoopUnswitch ?

You are right! Let me split this patch into two patches.

[jaykang10]

Following comment of @lebedev.ri, split previous patch into two patches. This one works on top of https://reviews.llvm.org/D99490

[lebedev.ri]

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll

1 ↗

(On Diff #333792)

Precommit this?

[jaykang10]

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll

1 ↗

(On Diff #333792)

Sorry... This test file is a copy of test/Transforms/LoopUnswitch/partial-unswitch.ll using SimpleLoopUnswitch. The output of SimpleLoopUnswitch is slightly different with LoopUnswitch one so I have run the script to generate assertion. I have checked the each tests' output. If it is not good to use the script, I will add the assertions manually. Please let me know.

[lebedev.ri]

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll

1 ↗

(On Diff #333792)

No i mean, just commit this test file as-is, obviously after regenerating the check lines to pass on main.

[jaykang10]

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll

1 ↗

(On Diff #333792)

Ah, sorry. let me commit this test separately after checking it again.

[jaykang10]

Following comment of @lebedev.ri, moved the test to separate patch.

[jaykang10]

Rebased

[jaykang10]

Rebase

[jaykang10]

Any comments please?

[jaykang10]

Any objection to push this change please? or can someone let me know what I have to do something more for this change please?

[jaykang10]

Rebased

[jaykang10]

@fhahn Can you review this change when you have time please?

[fhahn]

Does this patch give the expected speedup on omnetpp?

In D99354#2676137, @jaykang10 wrote:

@fhahn Can you review this change when you have time please?

please keep the 'common courtesy ping rate' of 1 week in mind (https://llvm.org/docs/CodeReview.html#code-reviews-speed-and-reciprocity), also considering that doing a proper review can take a substantial amount of time for a reviewer and there can be a number of reasons for a delayed response (like holidays, other urgent work).

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

212

Invariants here is not really accurate I think. Those are the instructions we need to duplicate outside the loop, right?

219

nit: perhaps a bit simpler NewInst->insertBefore(BB->getTerminator())

2016

struct not needed.

2029–2031

I think we should probably update IVConditionInfo to provide a better way to check if there are partially invariant conditions, e.g. an isPartiallyInvariant accessor. WDYT?

2149–2160

nit: In the partially invariant case, if UnswitchedSuccBB is an exit block, do not ....

2151

no need for llvm::, also can we just capture [&SuccBB]?

2375

We don't need to execute this loop at all for the PartiallyInvariant case, right?

2692

struct not needed.

2738

nit: no llvm:: should be needed.

2742

nit: no llvm:: should be needed

2871

Do we need to check here that we only do this for blocks where we partially unswitch on their conditions?

2950–2951

this change seems unrelated?

3154

can you add a comment here to explain the check?

[jaykang10]

In D99354#2678873, @fhahn wrote:

Does this patch give the expected speedup on omnetpp?

Yep, It makes expected speedup on omnetpp.

In D99354#2676137, @jaykang10 wrote:

@fhahn Can you review this change when you have time please?

please keep the 'common courtesy ping rate' of 1 week in mind (https://llvm.org/docs/CodeReview.html#code-reviews-speed-and-reciprocity), also considering that doing a proper review can take a substantial amount of time for a reviewer and there can be a number of reasons for a delayed response (like holidays, other urgent work).

I am sorry for inconvenient. I will follow the review rule.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

212

Yep, I will change it to ToToDuplicate.

219

The BB is empty block so we can not use BB.getTerminator() here.

2016

Yep, I will remove it.

2029–2031

I agree with you. It is better to provide isPartiallyInvariant in IVConditionInfo. If possible, it could be good to create a separate patch for it after pushing this patch.

2149–2160

Yep, I will update it.

2151

Yep, I will update it.

2375

Yep, you are right! I will update it.

2692

Yep, I will update it.

2738

Yep, I will update it.

2742

Yep, I will update it.

2871

We need to check the cost of the successors which are duplicated. The partially unswitched blocks on their conditions are only duplicated.

2950–2951

Ah, I ran clang-format. It is result of clang-format.

3154

Yep, I will add a comment.

[jaykang10]

Following comments of @fhahn, updated patch.

[fhahn]

Thanks for the latest update! I gave the patch a try on SPEC2006 on X86 with -O3 -flto and it looks like there's a crash when building 445.gobmk.

[jaykang10]

In D99354#2691707, @fhahn wrote:

Thanks for the latest update! I gave the patch a try on SPEC2006 on X86 with -O3 -flto and it looks like there's a crash when building 445.gobmk.

Ah, sorry, I could make a mistake while I rebase and update the change. Let me check.

[jaykang10]

Fixed a bug

[fhahn]

In D99354#2694514, @jaykang10 wrote:

Fixed a bug

Could you add a test case for the problem you fixed?

[jaykang10]

@fhahn There was a bug. I have fixed it. I have re-run spec2006 on x86 and it was fine.

[jaykang10]

In D99354#2694518, @fhahn wrote:

In D99354#2694514, @jaykang10 wrote:

Fixed a bug

Could you add a test case for the problem you fixed?

Yep, let me try to add it.

[jaykang10]

Added a test case for previous bug

[jaykang10]

@fhahn Is it ok to push this change? If you need something more, please let me know.

[jaykang10]

Rebased

[jaykang10]

@fhahn Can we push this change please? If you need something more, please let me know.

[jaykang10]

@fhahn I think it is ready to push this change.

[jaykang10]

@fhahn Sorry for ping.

[fhahn]

One final round of comments, after that I think this LG.

Could you also add a test for the threshold (like llvm/test/Transforms/LoopUnswitch/partial-unswitch-mssa-threshold.ll) and the interesting MemorySSA update cases from llvm/test/Transforms/LoopUnswitch/partial-unswitch-update-memoryssa.ll?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

210

nit:

Copy a set of loop invariant values \p ToDuplicate and insert them at the end of \p BB....?

... and conditional branch on the copied condition.` We only branch on a single value.

2950–2951

ok, can you undo it? As it is not related to the change.

[fhahn]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2029–2031

Sounds good, looking forward to a follow-up!

2050–2051

message in assert needs updating?

2737–2752

enough to capture [this]?

2746

Why TinyPtrVector here? In most case we will have at least 2 instructions here anyways? Probably better to use a SmallVector?

2749

emplace_back?

[jaykang10]

In D99354#2719131, @fhahn wrote:

One final round of comments, after that I think this LG.

Could you also add a test for the threshold (like llvm/test/Transforms/LoopUnswitch/partial-unswitch-mssa-threshold.ll) and the interesting MemorySSA update cases from llvm/test/Transforms/LoopUnswitch/partial-unswitch-update-memoryssa.ll?

Yep, let me try to add the tests.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

210

Yep, I will update it.

2029–2031

Yep

2050–2051

Yep, I will update it.

2737–2752

This code is inside static function rather than class's member function so we can not use [this]. I will update it with [&L].

2746

The SimpleUnswitchPass uses TinyPtrVector for UnswitchCandidates. In order to follow the interface, TinyPtrVector is being used here.

2749

TinyPtrVector does not have emplace_back.

2950–2951

Yep, I will undo it.

[jaykang10]

Following comments of @fhahn, updated code and added tests.

[fhahn]

LGTM, thanks! A few small comments that can be addressed before committing without further review I think.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

3154

nit: perhaps unswitched using a partially invariant condition, ...?

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll

1104 ↗

(On Diff #341113)

Could you add a brief comment and a more descriptive name for the test?

1143 ↗

(On Diff #341113)

It would be great if you could use descriptive labels in the test both for basic blocks and values, to make it easier to read in case people need to take a look in the future.

1150 ↗

(On Diff #341113)

please avoid using null as pointer, as this UB

[jaykang10]

In D99354#2726745, @fhahn wrote:

LGTM, thanks! A few small comments that can be addressed before committing without further review I think.

Thanks @fhahn! After updating code following your comments, I will push this change.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

3154

Yep, I will update it.

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll

1104 ↗

(On Diff #341113)

Yep, I will add them.

1143 ↗

(On Diff #341113)

Yep, I will update it.

1150 ↗

(On Diff #341113)

Yep, I will update it.

[jaykang10]

Following comments of @fhahn, updated code and test

[fhahn]

It looks like this patch causes SimpleLoopUnswitch to not terminate on some inputs, see https://bugs.llvm.org/show_bug.cgi?id=50302 for an example

[fhahn]

I think this causes https://bugs.llvm.org/show_bug.cgi?id=50279 as well

[jaykang10]

In D99354#2750756, @fhahn wrote:

I think this causes https://bugs.llvm.org/show_bug.cgi?id=50279 as well

@fhahn Thanks for letting me know. Let me have a look.

[fhahn]

In D99354#2750822, @jaykang10 wrote:

In D99354#2750756, @fhahn wrote:

I think this causes https://bugs.llvm.org/show_bug.cgi?id=50279 as well

@fhahn Thanks for letting me know. Let me have a look.

Thanks. If it's not straight-forward to resolve, it would be best to revert the patch for now.

[jaykang10]

In D99354#2754363, @fhahn wrote:

In D99354#2750822, @jaykang10 wrote:

In D99354#2750756, @fhahn wrote:

I think this causes https://bugs.llvm.org/show_bug.cgi?id=50279 as well

@fhahn Thanks for letting me know. Let me have a look.

Thanks. If it's not straight-forward to resolve, it would be best to revert the patch for now.

@fhahn Thanks for kind suggestion. I have figured out what causes the endless compiling. It looks we need to move partially invariant instructions rather than duplicating them. Once I resolve this issue, I need to check the score of benchmarks again. Therefore, I would post a patch for review tomorrow or the day after tomorrow.

[fhahn]

In D99354#2754454, @jaykang10 wrote:

In D99354#2754363, @fhahn wrote:

In D99354#2750822, @jaykang10 wrote:

In D99354#2750756, @fhahn wrote:

I think this causes https://bugs.llvm.org/show_bug.cgi?id=50279 as well

@fhahn Thanks for letting me know. Let me have a look.

Thanks. If it's not straight-forward to resolve, it would be best to revert the patch for now.

@fhahn Thanks for kind suggestion. I have figured out what causes the endless compiling. It looks we need to move partially invariant instructions rather than duplicating them. Once I resolve this issue, I need to check the score of benchmarks again. Therefore, I would post a patch for review tomorrow or the day after tomorrow.

Great, thanks! It sounds like it would make sense to revert until this is resolved than, so there's no need to rush :)

[jaykang10]

In D99354#2754816, @fhahn wrote:

In D99354#2754454, @jaykang10 wrote:

In D99354#2754363, @fhahn wrote:

In D99354#2750822, @jaykang10 wrote:

In D99354#2750756, @fhahn wrote:

I think this causes https://bugs.llvm.org/show_bug.cgi?id=50279 as well

@fhahn Thanks for letting me know. Let me have a look.

Thanks. If it's not straight-forward to resolve, it would be best to revert the patch for now.

@fhahn Thanks for kind suggestion. I have figured out what causes the endless compiling. It looks we need to move partially invariant instructions rather than duplicating them. Once I resolve this issue, I need to check the score of benchmarks again. Therefore, I would post a patch for review tomorrow or the day after tomorrow.

Great, thanks! It sounds like it would make sense to revert until this is resolved than, so there's no need to rush :)

Yep, I will revert it now. Additionally, I have seen SimpleLoopUnswitch with this patch does not affect performance number of omnetpp in SPEC2017. I need to figure out it too. Once I fixed these issues, I will let you know.

[fhahn]

In D99354#2756257, @jaykang10 wrote:

Yep, I will revert it now. Additionally, I have seen SimpleLoopUnswitch with this patch does not affect performance number of omnetpp in SPEC2017. I need to figure out it too. Once I fixed these issues, I will let you know.

Thanks!

[jaykang10]

The below bugs are fixed.
https://bugs.llvm.org/show_bug.cgi?id=50279
https://bugs.llvm.org/show_bug.cgi?id=50302

Even if this patch does not add the processed loop to loop pass manager again, the later loop transformations create loop with the header, which has partially invariant instructions, and add it to loop pass manager again. The loop pass manager runs SimpleLoopUnswitchPass with the loop again and it sometimes causes endless unswitch with some CFGs.

In order to avoid the endless unswitch, the new change removes the partially invariant instruction from header as below.

//    +--------------------+
//    |     preheader      |
//    |  %a = load %ptr    |
//    +--------------------+
//             |     /----------------\
//    +--------v----v------+          |
//    |      header        |---\      |
//    | %c = phi %a, %b    |   |      |
//    | %cond = cmp %c, .. |   |      |
//    | br %cond           |   |      |
//    +--------------------+   |      |
//             |               |      |
//    +--------v-----------+   |      |
//    |  store %ptr        |   |      |
//    +--------------------+   |      |
//             |               |      |
//    +--------v-----------<---/      |
//    |       latch        >----------/
//    |  %b = load %ptr    |
//    +--------------------+

[jaykang10]

@fhahn I have updated the diff which fixes the bugs. I have added phi node to header in order to remove the partially invariant instructions from header. If you feel something wrong from it, please let me know.

For the performance impact on the omnetpp of spec2017, it looks the different inlining is being happened on the benchmark between legacy and new pass manager and it causes fewer chances for unswitching loop. Maybe, I could try to figure out what it causes the different inlining between legacy and new pass manager later.

[jaykang10]

@fhahn Can we push this patch again please? I have checked the benchmarks and the scores are fine.

[fhahn]

In D99354#2763024, @jaykang10 wrote:

The below bugs are fixed.
https://bugs.llvm.org/show_bug.cgi?id=50279
https://bugs.llvm.org/show_bug.cgi?id=50302

Even if this patch does not add the processed loop to loop pass manager again, the later loop transformations create loop with the header, which has partially invariant instructions, and add it to loop pass manager again. The loop pass manager runs SimpleLoopUnswitchPass with the loop again and it sometimes causes endless unswitch with some CFGs.

In order to avoid the endless unswitch, the new change removes the partially invariant instruction from header as below.

//    +--------------------+
//    |     preheader      |
//    |  %a = load %ptr    |
//    +--------------------+
//             |     /----------------\
//    +--------v----v------+          |
//    |      header        |---\      |
//    | %c = phi %a, %b    |   |      |
//    | %cond = cmp %c, .. |   |      |
//    | br %cond           |   |      |
//    +--------------------+   |      |
//             |               |      |
//    +--------v-----------+   |      |
//    |  store %ptr        |   |      |
//    +--------------------+   |      |
//             |               |      |
//    +--------v-----------<---/      |
//    |       latch        >----------/
//    |  %b = load %ptr    |
//    +--------------------+

Hm, this seems quite fragile, as another pass may decide to move the load to the header again. In the legacy pass manager, the loop is annotated with metadata to avoid partially unswithcing the same loop multiple times (https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Scalar/LoopUnswitch.cpp#L849) Can we do the same here?

llvm/test/Transforms/SimpleLoopUnswitch/endless-unswitch.ll

1 ↗

(On Diff #345815)

We should avoid adding tests using -O3 here. Can we reproduce the issue by running simple-loop-unswtich twice? Also it would be good to clean up the test a bit.

[jaykang10]

In D99354#2773344, @fhahn wrote:

In D99354#2763024, @jaykang10 wrote:

The below bugs are fixed.
https://bugs.llvm.org/show_bug.cgi?id=50279
https://bugs.llvm.org/show_bug.cgi?id=50302

Even if this patch does not add the processed loop to loop pass manager again, the later loop transformations create loop with the header, which has partially invariant instructions, and add it to loop pass manager again. The loop pass manager runs SimpleLoopUnswitchPass with the loop again and it sometimes causes endless unswitch with some CFGs.

In order to avoid the endless unswitch, the new change removes the partially invariant instruction from header as below.

//    +--------------------+
//    |     preheader      |
//    |  %a = load %ptr    |
//    +--------------------+
//             |     /----------------\
//    +--------v----v------+          |
//    |      header        |---\      |
//    | %c = phi %a, %b    |   |      |
//    | %cond = cmp %c, .. |   |      |
//    | br %cond           |   |      |
//    +--------------------+   |      |
//             |               |      |
//    +--------v-----------+   |      |
//    |  store %ptr        |   |      |
//    +--------------------+   |      |
//             |               |      |
//    +--------v-----------<---/      |
//    |       latch        >----------/
//    |  %b = load %ptr    |
//    +--------------------+

Hm, this seems quite fragile, as another pass may decide to move the load to the header again.

um... I thought the load would not be moved to header later because it is loop invariant and we usually try to move it to outside loop.

In the legacy pass manager, the loop is annotated with metadata to avoid partially unswithcing the same loop multiple times (https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Scalar/LoopUnswitch.cpp#L849) Can we do the same here?

Yep, I have seen the metadata. I was just not sure whether people will accept it or not. Let me try to add it.

llvm/test/Transforms/SimpleLoopUnswitch/endless-unswitch.ll

1 ↗

(On Diff #345815)

um... in order to reproduce the endless unswitch, it needs other passes. Let me try to reduce the test.

[jaykang10]

Following comments of @fhahn, updated code.

• Added metadata to avoid endless unswitch
• Reduced test case

[jaykang10]

@fhahn I have followed the LoopUnSwitch pass's metadata and it looks ok with benchmark scores. If you feel something wrong from the change, please let me know.

[jaykang10]

@fhahn Can we push this change again please?

[jaykang10]

@fhahn Sorry for ping.

[fhahn]

Thanks for the update! looks good, as this in line with the legacy loop-unswitch implementation. I think the last piece of work remaining is to update the tests to check that the metadata is emitted? Also, could you add a test with an partially unswitchable loop, that already has llvm.loop.unswitch.partial.disable and ensure it is not unswitched?

[jaykang10]

Thanks for comments @fhahn! Let me add tests for the metadata.

[jaykang10]

Following comment of @fhahn, added tests for metadata with llvm.loop.unswitch.partial.disable.

[fhahn]

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch-check-metadata.ll

9 ↗

(On Diff #348957)

Can this test just be part of llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll?

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch-generate-metadata.ll

8 ↗

(On Diff #348957)

can you also include checks for the generated IR, ensuring that the whole metadata chain is correct (!llvm.loop attached to the right unswitched loop and contains the disable metadata? Also, can this just be part of llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll? I think we can just extend the checks there, rather than adding a new file + function.

[fhahn]

LGTM thanks, with the additional suggestions for the tests.

[jaykang10]

In D99354#2793045, @fhahn wrote:

LGTM thanks, with the additional suggestions for the tests.

Thanks @fhahn! After updating tests, I will push it.

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch-check-metadata.ll

9 ↗

(On Diff #348957)

Yep, I will add the metadata check in llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll.

llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch-generate-metadata.ll

8 ↗

(On Diff #348957)

Yep, I will add the metadata check in llvm/test/Transforms/SimpleLoopUnswitch/partial-unswitch.ll.

[jaykang10]

Following comments from @fhahn, updated tests.

[jonpa]

Great to see progress on this :-)

I don't see the speedup on omnetpp that we saw before on SystemZ. Do I need to pass some option or adjust some target threshold value somewhere?

[jaykang10]

In D99354#2794815, @jonpa wrote:

Great to see progress on this :-)

I don't see the speedup on omnetpp that we saw before on SystemZ. Do I need to pass some option or adjust some target threshold value somewhere?

@jonpa Thanks for checking the performance number on SystemZ.

I was able to see the speed up on omnetpp of SPEC2006 but there was no speedup on omentpp of SPEC2017. I guess you are looking SPEC2017 one.

It seemed there is different inlining between new pass manager and legacy pass manager and it causes to reduce partial loop unswitch opportunities on omnetpp of SPEC2017.

I am trying to get the speedup with the partial loop unswitch on omnetpp of SPEC2017 again. Once I fix it, I will let you know.

[jonpa]

In D99354#2795599, @jaykang10 wrote:

In D99354#2794815, @jonpa wrote:

Great to see progress on this :-)

I don't see the speedup on omnetpp that we saw before on SystemZ. Do I need to pass some option or adjust some target threshold value somewhere?

@jonpa Thanks for checking the performance number on SystemZ.

I was able to see the speed up on omnetpp of SPEC2006 but there was no speedup on omentpp of SPEC2017. I guess you are looking SPEC2017 one.

yes

It seemed there is different inlining between new pass manager and legacy pass manager and it causes to reduce partial loop unswitch opportunities on omnetpp of SPEC2017.

I am trying to get the speedup with the partial loop unswitch on omnetpp of SPEC2017 again. Once I fix it, I will let you know.

awesome!

[fhahn]

In D99354#2763024, @jaykang10 wrote:

The below bugs are fixed.
https://bugs.llvm.org/show_bug.cgi?id=50279
https://bugs.llvm.org/show_bug.cgi?id=50302

@jaykang10 can you verify the patch fixes the linked issue & close them?

[jaykang10]

In D99354#2802472, @fhahn wrote:

In D99354#2763024, @jaykang10 wrote:

The below bugs are fixed.
https://bugs.llvm.org/show_bug.cgi?id=50279
https://bugs.llvm.org/show_bug.cgi?id=50302

@jaykang10 can you verify the patch fixes the linked issue & close them?

Thanks for reminding me! @fhahn

Yep, I have checked those bug are fixed. I will close them.

[fhahn]

@jaykang10 it looks like this patch regressed codegen in some cases: https://bugs.llvm.org/show_bug.cgi?id=51139. It would be great if you could take a look.

[jaykang10]

In D99354#2903817, @fhahn wrote:

@jaykang10 it looks like this patch regressed codegen in some cases: https://bugs.llvm.org/show_bug.cgi?id=51139. It would be great if you could take a look.

Ah, let me have a look.

[xbolva00]

In D99354#2912791, @jaykang10 wrote:

In D99354#2903817, @fhahn wrote:

@jaykang10 it looks like this patch regressed codegen in some cases: https://bugs.llvm.org/show_bug.cgi?id=51139. It would be great if you could take a look.

Ah, let me have a look.

Maybe revert while investigating?

[jaykang10]

In D99354#2949059, @xbolva00 wrote:

In D99354#2912791, @jaykang10 wrote:

In D99354#2903817, @fhahn wrote:

@jaykang10 it looks like this patch regressed codegen in some cases: https://bugs.llvm.org/show_bug.cgi?id=51139. It would be great if you could take a look.

Ah, let me have a look.

Maybe revert while investigating?

As I mentioned on https://bugs.llvm.org/show_bug.cgi?id=51141, I am not sure the unswitch pass has to detect the cases which the loop load pre, sccp or other passes can optimize and the pass do not unswitch the loop...
I am creating a patch to fix https://bugs.llvm.org/show_bug.cgi?id=51141 first. The patch does not fix https://bugs.llvm.org/show_bug.cgi?id=51139. Once I create it, let me add you as reviewer.

[jaykang10]

In D99354#2949165, @jaykang10 wrote:

In D99354#2949059, @xbolva00 wrote:

In D99354#2912791, @jaykang10 wrote:

In D99354#2903817, @fhahn wrote:

@jaykang10 it looks like this patch regressed codegen in some cases: https://bugs.llvm.org/show_bug.cgi?id=51139. It would be great if you could take a look.

Ah, let me have a look.

Maybe revert while investigating?

As I mentioned on https://bugs.llvm.org/show_bug.cgi?id=51141, I am not sure the unswitch pass has to detect the cases which the loop load pre, sccp or other passes can optimize and the pass do not unswitch the loop...
I am creating a patch to fix https://bugs.llvm.org/show_bug.cgi?id=51141 first. The patch does not fix https://bugs.llvm.org/show_bug.cgi?id=51139. Once I create it, let me add you as reviewer.

For https://bugs.llvm.org/show_bug.cgi?id=51139, I have added below comment on it.

In this case, the inliner pass fails to inline the function g after unswitch because of the cost as below debug message.

NOT Inlining (cost=250, threshold=250), Call:   call void @g(i32 %2) #3

Originally, after inlining function g, the JumpThreading pass made the block with call @foo dead and SimplifyCFG pass deleted it.

If you add `always_inline` attribute to the function g's prototype as below, you can see the callfoo is gone.

void g(int h) __attribute__((always_inline));

In this case, as I mentioned previously, I am not sure the unswitch pass has to check the inline cost or something like that... If someone has idea about it, please let me know.

[jaykang10]

In D99354#2951701, @jaykang10 wrote:

In D99354#2949165, @jaykang10 wrote:

In D99354#2949059, @xbolva00 wrote:

In D99354#2912791, @jaykang10 wrote:

In D99354#2903817, @fhahn wrote:

@jaykang10 it looks like this patch regressed codegen in some cases: https://bugs.llvm.org/show_bug.cgi?id=51139. It would be great if you could take a look.

Ah, let me have a look.

Maybe revert while investigating?

As I mentioned on https://bugs.llvm.org/show_bug.cgi?id=51141, I am not sure the unswitch pass has to detect the cases which the loop load pre, sccp or other passes can optimize and the pass do not unswitch the loop...
I am creating a patch to fix https://bugs.llvm.org/show_bug.cgi?id=51141 first. The patch does not fix https://bugs.llvm.org/show_bug.cgi?id=51139. Once I create it, let me add you as reviewer.

For https://bugs.llvm.org/show_bug.cgi?id=51139, I have added below comment on it.

In this case, the inliner pass fails to inline the function g after unswitch because of the cost as below debug message.

NOT Inlining (cost=250, threshold=250), Call:   call void @g(i32 %2) #3

Originally, after inlining function g, the JumpThreading pass made the block with call @foo dead and SimplifyCFG pass deleted it.

If you add `always_inline` attribute to the function g's prototype as below, you can see the callfoo is gone.

void g(int h) __attribute__((always_inline));

In this case, as I mentioned previously, I am not sure the unswitch pass has to check the inline cost or something like that... If someone has idea about it, please let me know.

The SimpleUnswitchPass checks the cost of duplicating the loop so we could adjust it. Let me try to adjust the cost calculation for partially invariant unswitch.

[rui.zhang]

In D99354#2795599, @jaykang10 wrote:

I was able to see the speed up on omnetpp of SPEC2006 but there was no speedup on omentpp of SPEC2017. I guess you are looking SPEC2017 one.

@jaykang10, can you share the key loops that contribute to the speedup you observed from partial unswitch in spec2k6 omnetpp by any chance? Thanks.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

static cl::opt<bool> UnswitchGuards(
     "simple-loop-unswitch-guards", cl::init(true), cl::Hidden,
     cl::desc("If enabled, simple loop unswitching will also consider "
              "llvm.experimental.guard intrinsics as unswitch candidates."));
 static cl::opt<bool> DropNonTrivialImplicitNullChecks(
     "simple-loop-unswitch-drop-non-trivial-implicit-null-checks",
     cl::init(false), cl::Hidden,
     cl::desc("If enabled, drop make.implicit metadata in unswitched implicit "
              "null checks to save time analyzing if we can keep it."));
+static cl::opt<unsigned>
+    MSSAThreshold("simple-loop-unswitch-memoryssa-threshold",
+                  cl::desc("Max number of memory uses to explore during "
+                           "partial unswitching analysis"),
+                  cl::init(100), cl::Hidden);
 
 /// Collect all of the loop invariant input values transitively used by the
 /// homogeneous instruction graph from a given root.
 ///
 /// This essentially walks from a root recursively through loop variant operands
 /// which have the exact same opcode and finds all inputs which are loop
 /// invariant. For some operations these can be re-associated and unswitched out
 /// of the loop entirely.
static void buildPartialUnswitchConditionalBranch(BasicBlock &BB,
   IRBuilder<> IRB(&BB);
 
   Value *Cond = Direction ? IRB.CreateOr(Invariants) :
     IRB.CreateAnd(Invariants);
   IRB.CreateCondBr(Cond, Direction ? &UnswitchedSucc : &NormalSucc,
                    Direction ? &NormalSucc : &UnswitchedSucc);
 }
 
+/// Copy a set of loop invariant values, and conditionally branch on them.
+static void buildPartialInvariantUnswitchConditionalBranch(
+    BasicBlock &BB, ArrayRef<Value *> Invariants, bool Direction,
+    BasicBlock &UnswitchedSucc, BasicBlock &NormalSucc, Loop &L,
+    MemorySSAUpdater *MSSAU) {
+  ValueToValueMapTy VMap;
+  for (auto *Val : reverse(Invariants)) {
+    Instruction *Inst = cast<Instruction>(Val);
+    Instruction *NewInst = Inst->clone();
+    BB.getInstList().insert(BB.end(), NewInst);
+    RemapInstruction(NewInst, VMap,
+                     RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
+    VMap[Val] = NewInst;
+
+    if (!MSSAU)
+      continue;
+
+    MemorySSA *MSSA = MSSAU->getMemorySSA();
+    if (auto *MemUse =
+            dyn_cast_or_null<MemoryUse>(MSSA->getMemoryAccess(Inst))) {
+      auto *DefiningAccess = MemUse->getDefiningAccess();
+      // Get the first defining access before the loop.
+      while (L.contains(DefiningAccess->getBlock())) {
+        // If the defining access is a MemoryPhi, get the incoming
+        // value for the pre-header as defining access.
+        if (auto *MemPhi = dyn_cast<MemoryPhi>(DefiningAccess))
+          DefiningAccess =
+              MemPhi->getIncomingValueForBlock(L.getLoopPreheader());
+        else
+          DefiningAccess = cast<MemoryDef>(DefiningAccess)->getDefiningAccess();
+      }
+      MSSAU->createMemoryAccessInBB(NewInst, DefiningAccess,
+                                    NewInst->getParent(),
+                                    MemorySSA::BeforeTerminator);
+    }
+  }
+
+  IRBuilder<> IRB(&BB);
+  Value *Cond = VMap[Invariants[0]];
+  IRB.CreateCondBr(Cond, Direction ? &UnswitchedSucc : &NormalSucc,
+                   Direction ? &NormalSucc : &UnswitchedSucc);
+}
+
 /// Rewrite the PHI nodes in an unswitched loop exit basic block.
 ///
 /// Requires that the loop exit and unswitched basic block are the same, and
 /// that the exiting block was a unique predecessor of that block. Rewrites the
 /// PHI nodes in that block such that what were LCSSA PHI nodes become trivial
 /// PHI nodes from the old preheader that now contains the unswitched
 /// terminator.
 static void rewritePHINodesForUnswitchedExitBlock(BasicBlock &UnswitchedBB,
for (DomTreeNode *ChildN : *N) {
              "Cannot visit a node twice when walking a tree!");
       DomWorklist.push_back(ChildN);
     }
   } while (!DomWorklist.empty());
 }
 
 static void unswitchNontrivialInvariants(
     Loop &L, Instruction &TI, ArrayRef<Value *> Invariants,
-    SmallVectorImpl<BasicBlock *> &ExitBlocks, DominatorTree &DT, LoopInfo &LI,
-    AssumptionCache &AC, function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
+    SmallVectorImpl<BasicBlock *> &ExitBlocks,
+    struct IVConditionInfo &PartialIVInfo, DominatorTree &DT, LoopInfo &LI,
+    AssumptionCache &AC,
+    function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
     ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
   auto *ParentBB = TI.getParent();
   BranchInst *BI = dyn_cast<BranchInst>(&TI);
   SwitchInst *SI = BI ? nullptr : cast<SwitchInst>(&TI);
 
   // We can only unswitch switches, conditional branches with an invariant
-  // condition, or combining invariant conditions with an instruction.
+  // condition, or combining invariant conditions with an instruction or
+  // partially invariant instructions.
   assert((SI || (BI && BI->isConditional())) &&
          "Can only unswitch switches and conditional branch!");
-  bool FullUnswitch = SI || BI->getCondition() == Invariants[0];
+  bool PartiallyInvariant = !PartialIVInfo.InstToDuplicate.empty();
+  bool FullUnswitch =
+      SI || (BI->getCondition() == Invariants[0] && !PartiallyInvariant);
   if (FullUnswitch)
     assert(Invariants.size() == 1 &&
            "Cannot have other invariants with full unswitching!");
   else
     assert(isa<Instruction>(BI->getCondition()) &&
            "Partial unswitching requires an instruction as the condition!");
 
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
 
   // Constant and BBs tracking the cloned and continuing successor. When we are
   // unswitching the entire condition, this can just be trivially chosen to
   // unswitch towards `true`. However, when we are unswitching a set of
-  // invariants combined with `and` or `or`, the combining operation determines
-  // the best direction to unswitch: we want to unswitch the direction that will
-  // collapse the branch.
+  // invariants combined with `and` or `or` or partially invariant instructions,
+  // the combining operation determines the best direction to unswitch: we want
+  // to unswitch the direction that will collapse the branch.
   bool Direction = true;
   int ClonedSucc = 0;
   if (!FullUnswitch) {
     if (!match(BI->getCondition(), m_LogicalOr())) {
-      assert(match(BI->getCondition(), m_LogicalAnd()) &&
-             "Only `or`, `and`, an `select` instructions can combine "
-             "invariants being unswitched.");
+      assert(
+          (match(BI->getCondition(), m_LogicalAnd()) || PartiallyInvariant) &&
+          "Only `or`, `and`, an `select` instructions can combine invariants "
+          "being unswitched. Partially invariant instructions can also be "
+          "unswitched.");
+      if (match(BI->getCondition(), m_LogicalAnd()) ||
+          (PartiallyInvariant && !PartialIVInfo.KnownValue->isOneValue())) {
         Direction = false;
         ClonedSucc = 1;
       }
     }
+  }
 
   BasicBlock *RetainedSuccBB =
       BI ? BI->getSuccessor(1 - ClonedSucc) : SI->getDefaultDest();
   SmallSetVector<BasicBlock *, 4> UnswitchedSuccBBs;
   if (BI)
     UnswitchedSuccBBs.insert(BI->getSuccessor(ClonedSucc));
   else
     for (auto Case : SI->cases())
static void unswitchNontrivialInvariants(
   // Keep track of the dominator tree updates needed.
   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
 
   // Clone the loop for each unswitched successor.
   SmallVector<std::unique_ptr<ValueToValueMapTy>, 4> VMaps;
   VMaps.reserve(UnswitchedSuccBBs.size());
   SmallDenseMap<BasicBlock *, BasicBlock *, 4> ClonedPHs;
   for (auto *SuccBB : UnswitchedSuccBBs) {
+    // In Partially invariant case, if UnswithcedSuccBB is exit block, do not
+    // clone loop and assigned the UnswitchedSuccBB to ClonedPHs.
+    if (PartiallyInvariant && llvm::any_of(ExitBlocks, [&](BasicBlock *ExitBB) {
+          return ExitBB == SuccBB;
+        }))
+      ClonedPHs[SuccBB] = SuccBB;
+    else {
       VMaps.emplace_back(new ValueToValueMapTy());
       ClonedPHs[SuccBB] = buildClonedLoopBlocks(
           L, LoopPH, SplitBB, ExitBlocks, ParentBB, SuccBB, RetainedSuccBB,
           DominatingSucc, *VMaps.back(), DTUpdates, AC, DT, LI, MSSAU);
     }
+  }
 
   // Drop metadata if we may break its semantics by moving this instr into the
   // split block.
   if (TI.getMetadata(LLVMContext::MD_make_implicit)) {
     if (DropNonTrivialImplicitNullChecks)
       // Do not spend time trying to understand if we can keep it, just drop it
       // to save compile time.
       TI.setMetadata(LLVMContext::MD_make_implicit, nullptr);
if (FullUnswitch) {
     BranchInst::Create(RetainedSuccBB, ParentBB);
   } else {
     assert(BI && "Only branches have partial unswitching.");
     assert(UnswitchedSuccBBs.size() == 1 &&
            "Only one possible unswitched block for a branch!");
     BasicBlock *ClonedPH = ClonedPHs.begin()->second;
     // When doing a partial unswitch, we have to do a bit more work to build up
     // the branch in the split block.
+    if (PartiallyInvariant)
+      buildPartialInvariantUnswitchConditionalBranch(
+          *SplitBB, Invariants, Direction, *ClonedPH, *LoopPH, L, MSSAU);
+    else
       buildPartialUnswitchConditionalBranch(*SplitBB, Invariants, Direction,
                                             *ClonedPH, *LoopPH);
     DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
 
     if (MSSAU) {
       DT.applyUpdates(DTUpdates);
       DTUpdates.clear();
 
       // Perform MSSA cloning updates.
       for (auto &VMap : VMaps)
if (BI) {
     // we cant just go replace them with constants in both branches.
     //
     // For 'AND' we infer that true branch ("continue") means true
     // for each invariant operand.
     // For 'OR' we can infer that false branch ("continue") means false
     // for each invariant operand.
     // So it happens that for multiple-partial case we dont replace
     // in the unswitched branch.
-    bool ReplaceUnswitched = FullUnswitch || (Invariants.size() == 1);
+    bool ReplaceUnswitched =
+        FullUnswitch || (Invariants.size() == 1) || PartiallyInvariant;
 
     ConstantInt *UnswitchedReplacement =
         Direction ? ConstantInt::getTrue(BI->getContext())
                   : ConstantInt::getFalse(BI->getContext());
     ConstantInt *ContinueReplacement =
         Direction ? ConstantInt::getFalse(BI->getContext())
                   : ConstantInt::getTrue(BI->getContext());
     for (Value *Invariant : Invariants)
       // Use make_early_inc_range here as set invalidates the iterator.
       for (Use &U : llvm::make_early_inc_range(Invariant->uses())) {
         Instruction *UserI = dyn_cast<Instruction>(U.getUser());
-        if (!UserI)
+        if (!UserI || PartiallyInvariant)
           continue;
 
         // Replace it with the 'continue' side if in the main loop body, and the
         // unswitched if in the cloned blocks.
         if (DT.dominates(LoopPH, UserI->getParent()))
           U.set(ContinueReplacement);
         else if (ReplaceUnswitched &&
                  DT.dominates(ClonedPH, UserI->getParent()))
#endif
 
   // Now that we've unswitched something, make callbacks to report the changes.
   // For that we need to merge together the updated loops and the cloned loops
   // and check whether the original loop survived.
   SmallVector<Loop *, 4> SibLoops;
   for (Loop *UpdatedL : llvm::concat<Loop *>(NonChildClonedLoops, HoistedLoops))
     if (UpdatedL->getParentLoop() == ParentL)
       SibLoops.push_back(UpdatedL);
-  UnswitchCB(IsStillLoop, SibLoops);
+  UnswitchCB(IsStillLoop, PartiallyInvariant, SibLoops);
 
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
 
   if (BI)
     ++NumBranches;
   else
     ++NumSwitches;
static int CalculateUnswitchCostMultiplier(
 
   LLVM_DEBUG(dbgs() << "  Computed multiplier  " << CostMultiplier
                     << " (siblings " << SiblingsMultiplier << " * clones "
                     << (1 << ClonesPower) << ")"
                     << " for unswitch candidate: " << TI << "\n");
   return CostMultiplier;
 }
 
-static bool
-unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
-                      AssumptionCache &AC, TargetTransformInfo &TTI,
-                      function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
-                      ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
+static bool unswitchBestCondition(
+    Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
+    AAResults &AA, TargetTransformInfo &TTI,
+    function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
   // Collect all invariant conditions within this loop (as opposed to an inner
   // loop which would be handled when visiting that inner loop).
   SmallVector<std::pair<Instruction *, TinyPtrVector<Value *>>, 4>
       UnswitchCandidates;
 
   // Whether or not we should also collect guards in the loop.
   bool CollectGuards = false;
   if (UnswitchGuards) {
     auto *GuardDecl = L.getHeader()->getParent()->getParent()->getFunction(
         Intrinsic::getName(Intrinsic::experimental_guard));
     if (GuardDecl && !GuardDecl->use_empty())
       CollectGuards = true;
   }
 
+  struct IVConditionInfo PartialIVInfo;
   for (auto *BB : L.blocks()) {
     if (LI.getLoopFor(BB) != &L)
       continue;
 
     if (CollectGuards)
       for (auto &I : *BB)
         if (isGuard(&I)) {
           auto *Cond = cast<IntrinsicInst>(&I)->getArgOperand(0);
if (!BI || !BI->isConditional() || isa<Constant>(BI->getCondition()) ||
       continue;
 
     if (L.isLoopInvariant(BI->getCondition())) {
       UnswitchCandidates.push_back({BI, {BI->getCondition()}});
       continue;
     }
 
     Instruction &CondI = *cast<Instruction>(BI->getCondition());
-    if (!match(&CondI, m_CombineOr(m_LogicalAnd(), m_LogicalOr())))
-      continue;
-
+    if (match(&CondI, m_CombineOr(m_LogicalAnd(), m_LogicalOr()))) {
       TinyPtrVector<Value *> Invariants =
           collectHomogenousInstGraphLoopInvariants(L, CondI, LI);
       if (Invariants.empty())
         continue;
 
       UnswitchCandidates.push_back({BI, std::move(Invariants)});
+      continue;
+    }
+  }
+
+  if (MSSAU &&
+      !llvm::any_of(UnswitchCandidates, [&](auto &TerminatorAndInvariants) {
+        return TerminatorAndInvariants.first == L.getHeader()->getTerminator();
+      })) {
+    MemorySSA *MSSA = MSSAU->getMemorySSA();
+    if (auto Info = llvm::hasPartialIVCondition(&L, MSSAThreshold, MSSA, &AA)) {
+      LLVM_DEBUG(
+          dbgs() << "simple-loop-unswitch: Found partially invariant condition "
+                 << *Info->InstToDuplicate[0] << "\n");
+      PartialIVInfo = *Info;
+      TinyPtrVector<Value *> ValsToDuplicate;
+      for (auto *Inst : Info->InstToDuplicate)
+        ValsToDuplicate.push_back(Inst);
+      UnswitchCandidates.push_back(
+          {L.getHeader()->getTerminator(), std::move(ValsToDuplicate)});
+    }
   }
 
   // If we didn't find any candidates, we're done.
   if (UnswitchCandidates.empty())
     return false;
 
   // Check if there are irreducible CFG cycles in this loop. If so, we cannot
   // easily unswitch non-trivial edges out of the loop. Doing so might turn the
auto ComputeUnswitchedCost = [&](Instruction &TI,
 
     InstructionCost Cost = 0;
     for (BasicBlock *SuccBB : successors(&BB)) {
       // Don't count successors more than once.
       if (!Visited.insert(SuccBB).second)
         continue;
 
       // If this is a partial unswitch candidate, then it must be a conditional
-      // branch with a condition of either `or`, `and`, or their corresponding
-      // select forms. In that case, one of the successors is necessarily
-      // duplicated, so don't even try to remove its cost.
+      // branch with a condition of either `or`, `and`, their corresponding
+      // select forms or partially invariant instructions. In that case, one of
+      // the successors is necessarily duplicated, so don't even try to remove
+      // its cost.
       if (!FullUnswitch) {
         auto &BI = cast<BranchInst>(TI);
         if (match(BI.getCondition(), m_LogicalAnd())) {
           if (SuccBB == BI.getSuccessor(1))
             continue;
-        } else {
-          assert(match(BI.getCondition(), m_LogicalOr()) &&
-                 "Only `and` and `or` conditions can result in a partial "
-                 "unswitch!");
+        } else if (match(BI.getCondition(), m_LogicalOr())) {
           if (SuccBB == BI.getSuccessor(0))
             continue;
+        } else if (!PartialIVInfo.InstToDuplicate.empty()) {
+          if (PartialIVInfo.KnownValue->isOneValue() &&
+              SuccBB == BI.getSuccessor(1))
+            continue;
+          else if (!PartialIVInfo.KnownValue->isOneValue() &&
+                   SuccBB == BI.getSuccessor(0))
+            continue;
         }
       }
 
       // This successor's domtree will not need to be duplicated after
       // unswitching if the edge to the successor dominates it (and thus the
       // entire tree). This essentially means there is no other path into this
       // subtree and so it will end up live in only one clone of the loop.
       if (SuccBB->getUniquePredecessor() ||
if (BestUnswitchCost >= UnswitchThreshold) {
     return false;
   }
 
   // If the best candidate is a guard, turn it into a branch.
   if (isGuard(BestUnswitchTI))
     BestUnswitchTI = turnGuardIntoBranch(cast<IntrinsicInst>(BestUnswitchTI), L,
                                          ExitBlocks, DT, LI, MSSAU);
 
-  LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = "
-                    << BestUnswitchCost << ") terminator: " << *BestUnswitchTI
-                    << "\n");
+  LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = " << BestUnswitchCost
+                    << ") terminator: " << *BestUnswitchTI << "\n");
   unswitchNontrivialInvariants(L, *BestUnswitchTI, BestUnswitchInvariants,
-                               ExitBlocks, DT, LI, AC, UnswitchCB, SE, MSSAU);
+                               ExitBlocks, PartialIVInfo, DT, LI, AC,
+                               UnswitchCB, SE, MSSAU);
   return true;
 }
 
 /// Unswitch control flow predicated on loop invariant conditions.
 ///
 /// This first hoists all branches or switches which are trivial (IE, do not
 /// require duplicating any part of the loop) out of the loop body. It then
 /// looks at other loop invariant control flows and tries to unswitch those as
 /// well by cloning the loop if the result is small enough.
 ///
-/// The `DT`, `LI`, `AC`, `TTI` parameters are required analyses that are also
-/// updated based on the unswitch.
-/// The `MSSA` analysis is also updated if valid (i.e. its use is enabled).
+/// The `DT`, `LI`, `AC`, `AA`, `TTI` parameters are required analyses that are
+/// also updated based on the unswitch. The `MSSA` analysis is also updated if
+/// valid (i.e. its use is enabled).
 ///
 /// If either `NonTrivial` is true or the flag `EnableNonTrivialUnswitch` is
 /// true, we will attempt to do non-trivial unswitching as well as trivial
 /// unswitching.
 ///
 /// The `UnswitchCB` callback provided will be run after unswitching is
 /// complete, with the first parameter set to `true` if the provided loop
 /// remains a loop, and a list of new sibling loops created.
 ///
 /// If `SE` is non-null, we will update that analysis based on the unswitching
 /// done.
-static bool unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI,
-                         AssumptionCache &AC, TargetTransformInfo &TTI,
-                         bool NonTrivial,
-                         function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
-                         ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
+static bool
+unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
+             AAResults &AA, TargetTransformInfo &TTI, bool NonTrivial,
+             function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
   assert(L.isRecursivelyLCSSAForm(DT, LI) &&
          "Loops must be in LCSSA form before unswitching.");
 
   // Must be in loop simplified form: we need a preheader and dedicated exits.
   if (!L.isLoopSimplifyForm())
     return false;
 
   // Try trivial unswitch first before loop over other basic blocks in the loop.
   if (unswitchAllTrivialConditions(L, DT, LI, SE, MSSAU)) {
     // If we unswitched successfully we will want to clean up the loop before
     // processing it further so just mark it as unswitched and return.
-    UnswitchCB(/*CurrentLoopValid*/ true, {});
+    UnswitchCB(/*CurrentLoopValid*/ true, false, {});
     return true;
   }
 
   // Check whether we should continue with non-trivial conditions.
   // EnableNonTrivialUnswitch: Global variable that forces non-trivial
   //                           unswitching for testing and debugging.
   // NonTrivial: Parameter that enables non-trivial unswitching for this
   //             invocation of the transform. But this should be allowed only
unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
   // For non-trivial unswitching, because it often creates new loops, we rely on
   // the pass manager to iterate on the loops rather than trying to immediately
   // reach a fixed point. There is no substantial advantage to iterating
   // internally, and if any of the new loops are simplified enough to contain
   // trivial unswitching we want to prefer those.
 
   // Try to unswitch the best invariant condition. We prefer this full unswitch to
   // a partial unswitch when possible below the threshold.
-  if (unswitchBestCondition(L, DT, LI, AC, TTI, UnswitchCB, SE, MSSAU))
+  if (unswitchBestCondition(L, DT, LI, AC, AA, TTI, UnswitchCB, SE, MSSAU))
     return true;
 
   // No other opportunities to unswitch.
   return false;
 }
 
 PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
                                               LoopStandardAnalysisResults &AR,
                                               LPMUpdater &U) {
   Function &F = *L.getHeader()->getParent();
   (void)F;
 
   LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << L
                     << "\n");
 
   // Save the current loop name in a variable so that we can report it even
   // after it has been deleted.
   std::string LoopName = std::string(L.getName());
 
   auto UnswitchCB = [&L, &U, &LoopName](bool CurrentLoopValid,
+                                        bool PartiallyInvariant,
                                         ArrayRef<Loop *> NewLoops) {
     // If we did a non-trivial unswitch, we have added new (cloned) loops.
     if (!NewLoops.empty())
       U.addSiblingLoops(NewLoops);
 
     // If the current loop remains valid, we should revisit it to catch any
     // other unswitch opportunities. Otherwise, we need to mark it as deleted.
-    if (CurrentLoopValid)
+    if (CurrentLoopValid) {
+      if (!PartiallyInvariant)
         U.revisitCurrentLoop();
-    else
+    } else
       U.markLoopAsDeleted(L, LoopName);
   };
 
   Optional<MemorySSAUpdater> MSSAU;
   if (AR.MSSA) {
     MSSAU = MemorySSAUpdater(AR.MSSA);
     if (VerifyMemorySSA)
       AR.MSSA->verifyMemorySSA();
   }
-  if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.TTI, NonTrivial, UnswitchCB,
-                    &AR.SE, MSSAU.hasValue() ? MSSAU.getPointer() : nullptr))
+  if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.AA, AR.TTI, NonTrivial,
+                    UnswitchCB, &AR.SE,
+                    MSSAU.hasValue() ? MSSAU.getPointer() : nullptr))
     return PreservedAnalyses::all();
 
   if (AR.MSSA && VerifyMemorySSA)
     AR.MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
bool SimpleLoopUnswitchLegacyPass::runOnLoop(Loop *L, LPPassManager &LPM) {
   Function &F = *L->getHeader()->getParent();
 
   LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << *L
                     << "\n");
 
   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+  auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
   auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
   MemorySSA *MSSA = nullptr;
   Optional<MemorySSAUpdater> MSSAU;
   if (EnableMSSALoopDependency) {
     MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
     MSSAU = MemorySSAUpdater(MSSA);
   }
 
   auto *SEWP = getAnalysisIfAvailable<ScalarEvolutionWrapperPass>();
   auto *SE = SEWP ? &SEWP->getSE() : nullptr;
 
-  auto UnswitchCB = [&L, &LPM](bool CurrentLoopValid,
+  auto UnswitchCB = [&L, &LPM](bool CurrentLoopValid, bool PartiallyInvariant,
                                ArrayRef<Loop *> NewLoops) {
     // If we did a non-trivial unswitch, we have added new (cloned) loops.
     for (auto *NewL : NewLoops)
       LPM.addLoop(*NewL);
 
     // If the current loop remains valid, re-add it to the queue. This is
     // a little wasteful as we'll finish processing the current loop as well,
     // but it is the best we can do in the old PM.
-    if (CurrentLoopValid)
+    if (CurrentLoopValid) {
+      if (!PartiallyInvariant)
         LPM.addLoop(*L);
-    else
+    } else
       LPM.markLoopAsDeleted(*L);
   };
 
   if (MSSA && VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
-  bool Changed = unswitchLoop(*L, DT, LI, AC, TTI, NonTrivial, UnswitchCB, SE,
-                              MSSAU.hasValue() ? MSSAU.getPointer() : nullptr);
+  bool Changed =
+      unswitchLoop(*L, DT, LI, AC, AA, TTI, NonTrivial, UnswitchCB, SE,
 
   if (MSSA && VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(DT.verify(DominatorTree::VerificationLevel::Fast));