Missed Optimization: fold disjoint bitwise ORs (with select/zext i1) into a single constant select

[Leo0506]

define i8 @src(i1 %arg0) {
  %v0 = zext i1 %arg0 to i8
  %v1 = select i1 %arg0, i8 0, i8 8
  %v2 = or disjoint i8 %v1, %v0
  %v3 = or disjoint i8 %v2, 16
  ret i8 %v3
}

define i8 @tgt(i1 %arg0) {
  %v1 = select i1 %arg0, i8 17, i8 24
  ret i8 %v1
}

alive2: https://alive2.llvm.org/ce/z/vSQli0
godbolt: https://godbolt.org/z/543Mvx7MG
Pattern Found in: https://github.com/dtcxzyw/llvm-opt-benchmark/blob/main/bench/llvm/optimized/CGExpr.ll

[spaits]

I did some investigation.
I found out, if we wouldn't require one operand to be constant in InstCombinerImpl::foldBinOpIntoSelectOrPhi this optimization would be done:

llvm-project/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Lines 2263 to 2266 in 0b01b96

	Instruction *InstCombinerImpl::foldBinOpIntoSelectOrPhi(BinaryOperator &I) {
	if (!isa<Constant>(I.getOperand(1)))
	return nullptr;

This constraint was also mentioned in #116073 by @nikic .

I think this issue gives us a reason to think about getting rid of that constraint. I don't think it would affect program correctness, since I don't seem to find any places in that code that would actually require a constant (thanks to #116073). The compile time impact is another question. That should be measured.

Can please I ask your opinions @dtcxzyw @nikic ?

[nikic]

Yes, if compile-time allows we should get rid of that limitation.

[spaits]

Great. Thank you for your response. I will remove the constant check, check the changes in tests and adjust them and create a PR. Then we can measure the compile time impact of the PR.

[spaits]

Sorry I couldn't get back earlier, I didn't have as much time on my hand as I have imagined. Here are my findings so far.

What happens if we delete the constant requirement.

I have deleted the constant requirement. It turns out, in some tests
it will cause an infinite loop.

Here is an example from llvm/test/Transforms/InstCombine/logical-select.ll.
This is the problemaitc function:

define i32 @fold_inverted_icmp_preds(i32 %a, i32 %b, i32 %c, i32 %d) {
  %cmp1 = icmp slt i32 %a, %b
  %sel1 = select i1 %cmp1, i32 %c, i32 0
  %cmp2 = icmp sge i32 %a, %b
  %sel2 = select i1 %cmp2, i32 %d, i32 0
  %or = or i32 %sel1, %sel2
  ret i32 %or
}

When enabling both parameters to be non-const we will end up with the following
IR after a few iterations:

%cmp1 = icmp slt i32 %a, %b
%cmp2.not = icmp slt i32 %a, %b
%1 = select i1 %cmp2.not, i32 0, i32 %d
%2 = select i1 %cmp1, i32 %c, i32 0
%or = or i32 %1, %2
ret i32 %or

When FoldOpIntoSelect runs for the last or, it will be able to create
a new select. simplifyOperationIntoSelectOperand will be successful for the false value.
The resulting LLVM-IR will look like this:

%cmp1 = icmp slt i32 %a, %b
%cmp2.not = icmp slt i32 %a, %b
%1 = select i1 %cmp1, i32 %c, i32 0
%2 = or i32 %d, %1
%or = select i1 %cmp2.not, i32 %1, i32 %2
ret i32 %or

Then instcombine runs again. Now llvm::InstCombinerImpl::foldSelectIntoOp will run. It will fold the visited select, into and or instruction and we will have the following LLVM-IR:

%cmp1 = icmp slt i32 %a, %b
%cmp2.not = icmp slt i32 %a, %b
%1 = select i1 %cmp1, i32 %c, i32 0
%2 = select i1 %cmp2.not, i32 0, i32 %d
%or = or i32 %1, %2
ret i32 %or

On this LLVM-IR, FoldOpIntoSelect will be able to run again, and we have an infinite loop.

What's next?

I would be happy, if we could delete the const constraint for FoldOpIntoSelect. So I think I should fix the looping. Also there are many other tests timing out:

llvm/test/Transforms/InstCombine/conditional-variable-length-signext-after-high-bit-extract.ll
llvm/test/Transforms/InstCombine/icmp-and-lowbit-mask.ll
llvm/test/Transforms/InstCombine/logical-select-inseltpoison.ll
llvm/test/Transforms/InstCombine/logical-select.ll
llvm/test/Transforms/InstCombine/vec_sext.ll
llvm/test/Transforms/InstCombine/conditional-variable-length-signext-after-high-bit-extract.ll
llvm/test/Transforms/InstCombine/icmp-and-lowbit-mask.ll
llvm/test/Transforms/InstCombine/logical-select-inseltpoison.ll
llvm/test/Transforms/InstCombine/logical-select.ll
llvm/test/Transforms/InstCombine/vec_sext.ll

And these are just the tests I have encountered on my laptop, that is not the strongest one (so the testing has ended when all the processes (10) launched by lit were "blocked" by a timing out tc.) There are probably more tests timing out. I didn't have the time to investigate all.

To me, both FoldOpIntoSelect and foldSelectIntoOp seem to be rational with what they do. I don't really have an idea where should I start experimenting.

Also I have created a combine in a commit, that addresses this pattern before finding out that tis optimization can be achived by deleting the mentioned constraint on constants. spaits@b6d04e2#diff-5a47d869f7afa461d83e87efbcc1a1b2edf68d03d3033751dcedd57272dbba10R4093
This combine could also be generalized for binary operators I think.

I would apprichiate your opinions on this. How would you approach this issue? Should we still want to enable FoldOpIntoSelect for both non-const operand instructions? If so, should I try to make adjustments to FoldOpIntoSelect or foldSelectIntoOp? Or I should just create the function that matches the pattern for binary operators?

[nikic]

Ah yeah, that sounds familiar. I think we tried this before and ran into the same issue :)

Maybe setting SimplifyBothArms=true for the non-constant case would work?

[spaits]

Finally I had a little bit of time this weekend to experiment with this ticket.

I have added SimplifyBothArms=true. It fixes the timeouts for me, but when that is added the optimization described in the issue doesn't happen. This change also generates a bit better code in some cases in the check-llvm-transforms-instcombine suite. I have created a (draft) PR #166102 for this.

I have also decide to try to prevent this cycle in FoldOpIntoSelect. I figured the loops could be prevented by checking whether any of the new select's arm is another select and disabling the optimization in that case. It works for some tests, but the following two test cases still time out:

llvm/test/Transforms/InstCombine/icmp-and-lowbit-mask.ll
llvm/test/Transforms/InstCombine/conditional-variable-length-signext-after-high-bit-extract.ll

I will wish to investigate them when I have time again.

[spaits]

I did some even further investigation on the issue. I have checked why FoldOpIntoSelect handles the code in the issue succesfully. It turns out the following code in simplifyOrInst is the one doing the work:

; thread #1, name = 'opt', stop reason = step over
;     frame #0: 0x000055555b7a8ace opt`simplifyOrInst(Op0=0x0000555560389b40, Op1=0x0000555560389fa0, Q=0x00007fffffff9910, MaxRecurse=3) at InstructionSimplify.cpp:2364:12
;    2361   // X | X = X
;    2362   // X | 0 = X
;    2363   if (Op0 == Op1 || match(Op1, m_Zero()))
; -> 2364     return Op0;
;    2365
;    2366   if (Value *R = simplifyOrLogic(Op0, Op1))
;    2367     return R;

So enabling foldBinOpIntoSelectOrPhi for more cases won't get us this optimization in all cases.
Here is an example, where the code I wrote in spaits@b6d04e2#diff-5a47d869f7afa461d83e87efbcc1a1b2edf68d03d3033751dcedd57272dbba10R4093 will do the optimization but foldBinOpIntoSelectOrPhi will not:

define i8 @src(i1 %arg0) {
  %v0 = zext i1 %arg0 to i8
  %v1 = select i1 %arg0, i8 21, i8 42
  %v2 = or i8 %v1, %v0
  %v3 = or i8 %v2, 16
  ret i8 %v3
}

Alive2:
https://alive2.llvm.org/ce/z/WGnVos

Also, I think the code in my patch could be genralized for and and xor too.