[InstCombine] use select-of-constants with set/clear bit mask patterns

[spatel]

Summary

Cond ? (X & ~C) : (X | C) --> (X & ~C) | (Cond ? 0 : C)
Cond ? (X | C) : (X & ~C) --> (X & ~C) | (Cond ? C : 0)

The select-of-constants form results in better codegen. There's an existing test diff that shows a transform that results in an extra IR instruction, but that's an existing problem.

This is motivated by code seen in LLVM itself - see PR37581:
https://bugs.llvm.org/show_bug.cgi?id=37581

define i8 @src(i8 %x, i8 %C, i1 %b)  {
  %notC = xor i8 %C, -1
  %and = and i8 %x, %notC
  %or = or i8 %x, %C
  %cond = select i1 %b, i8 %or, i8 %and
  ret i8 %cond
}

define i8 @tgt(i8 %x, i8 %C, i1 %b)  {
  %notC = xor i8 %C, -1
  %and = and i8 %x, %notC
  %mul = select i1 %b, i8 %C, i8 0
  %or = or i8 %mul, %and
  ret i8 %or
}

http://volta.cs.utah.edu:8080/z/Vt2WVm

[nikic]

LGTM

[sanwou01]

Hi, we're seeing a small (1.0%) regression in omnetpp_r in SPEC INT 2017 on AArch64 with LTO enabled that bisects to this patch. I should be able to reduce omnetpp_r to a small IR example that shows the changed AArch64 codegen, if that's useful. A revert is probably not necessary if all we need is an additional pattern or two in the AArch64 backend.

[spatel]

In D78880#2022752, @sanwou01 wrote:

Hi, we're seeing a small (1.0%) regression in omnetpp_r in SPEC INT 2017 on AArch64 with LTO enabled that bisects to this patch. I should be able to reduce omnetpp_r to a small IR example that shows the changed AArch64 codegen, if that's useful. A revert is probably not necessary if all we need is an additional pattern or two in the AArch64 backend.

Thanks for letting me know. Yes, an IR example would be very helpful in deciding how to proceed.
Note: D68911 has been sitting untouched for ~6 months...

[sanwou01]

Do you think D68911 has a good chance of helping here? I can do a quick test run (quicker than finding a good reproducer) to see if improves.

[spatel]

In D78880#2022851, @sanwou01 wrote:

Do you think D68911 has a good chance of helping here? I can do a quick test run (quicker than finding a good reproducer) to see if improves.

I can't say for sure, but it was the 1st thing that came to mind, and yes, it should be easy to run perf tests.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (C2Log > C1Log) {
     V = Builder.CreateZExtOrTrunc(V, Y->getType());
 
   if (NeedXor)
     V = Builder.CreateXor(V, *C2);
 
   return Builder.CreateOr(V, Y);
 }
 
+/// Canonicalize a set or clear of a masked set of constant bits to
+/// select-of-constants form.
+static Instruction *foldSetClearBits(SelectInst &Sel,
+                                     InstCombiner::BuilderTy &Builder) {
+  Value *Cond = Sel.getCondition();
+  Value *T = Sel.getTrueValue();
+  Value *F = Sel.getFalseValue();
+  Type *Ty = Sel.getType();
+  Value *X;
+  const APInt *NotC, *C;
+
+  // Cond ? (X & ~C) : (X | C) --> (X & ~C) | (Cond ? 0 : C)
+  if (match(T, m_And(m_Value(X), m_APInt(NotC))) &&
+      match(F, m_OneUse(m_Or(m_Specific(X), m_APInt(C)))) && *NotC == ~(*C)) {
+    Constant *Zero = ConstantInt::getNullValue(Ty);
+    Constant *OrC = ConstantInt::get(Ty, *C);
+    Value *NewSel = Builder.CreateSelect(Cond, Zero, OrC, "masksel", &Sel);
+    return BinaryOperator::CreateOr(T, NewSel);
+  }
+
+  // Cond ? (X | C) : (X & ~C) --> (X & ~C) | (Cond ? C : 0)
+  if (match(F, m_And(m_Value(X), m_APInt(NotC))) &&
+      match(T, m_OneUse(m_Or(m_Specific(X), m_APInt(C)))) && *NotC == ~(*C)) {
+    Constant *Zero = ConstantInt::getNullValue(Ty);
+    Constant *OrC = ConstantInt::get(Ty, *C);
+    Value *NewSel = Builder.CreateSelect(Cond, OrC, Zero, "masksel", &Sel);
+    return BinaryOperator::CreateOr(F, NewSel);
+  }
+
+  return nullptr;
+}
+
 /// Transform patterns such as (a > b) ? a - b : 0 into usub.sat(a, b).
 /// There are 8 commuted/swapped variants of this pattern.
 /// TODO: Also support a - UMIN(a,b) patterns.
 static Value *canonicalizeSaturatedSubtract(const ICmpInst *ICI,
                                             const Value *TrueVal,
                                             const Value *FalseVal,
                                             InstCombiner::BuilderTy &Builder) {
   ICmpInst::Predicate Pred = ICI->getPredicate();
Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   if (ICmpInst *ICI = dyn_cast<ICmpInst>(CondVal))
     if (Instruction *Result = foldSelectInstWithICmp(SI, ICI))
       return Result;
 
   if (Instruction *Add = foldAddSubSelect(SI, Builder))
     return Add;
   if (Instruction *Add = foldOverflowingAddSubSelect(SI, Builder))
     return Add;
+  if (Instruction *Or = foldSetClearBits(SI, Builder))
+    return Or;
 
   // Turn (select C, (op X, Y), (op X, Z)) -> (op X, (select C, Y, Z))
   auto *TI = dyn_cast<Instruction>(TrueVal);
   auto *FI = dyn_cast<Instruction>(FalseVal);
   if (TI && FI && TI->getOpcode() == FI->getOpcode())
     if (Instruction *IV = foldSelectOpOp(SI, TI, FI))
       return IV;
 

llvm/test/Transforms/InstCombine/cast.ll

;
   %E = sext i32 %D to i64
   ret i64 %E
 ; lshr+shl will be handled by DAGCombine.
 }
 
 define i64 @test51(i64 %A, i1 %cond) {
 ; ALL-LABEL: @test51(
 ; ALL-NEXT:    [[C:%.*]] = and i64 [[A:%.*]], 4294967294
-; ALL-NEXT:    [[D:%.*]] = or i64 [[A]], 1
-; ALL-NEXT:    [[E:%.*]] = select i1 [[COND:%.*]], i64 [[C]], i64 [[D]]
-; ALL-NEXT:    [[SEXT:%.*]] = shl i64 [[E]], 32
+; ALL-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND:%.*]], true
+; ALL-NEXT:    [[MASKSEL:%.*]] = zext i1 [[NOT_COND]] to i64
+; ALL-NEXT:    [[E:%.*]] = or i64 [[C]], [[MASKSEL]]
+; ALL-NEXT:    [[SEXT:%.*]] = shl nuw i64 [[E]], 32
 ; ALL-NEXT:    [[F:%.*]] = ashr exact i64 [[SEXT]], 32
 ; ALL-NEXT:    ret i64 [[F]]
 ;
   %B = trunc i64 %A to i32
   %C = and i32 %B, -2
   %D = or i32 %B, 1
   %E = select i1 %cond, i32 %C, i32 %D
   %F = sext i32 %E to i64

llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll

;
   %res = mul i32 %select, %select2
   %res2 = mul i32 %res, %and2 ; to bump up the use count of the and
   ret i32 %res2
 }
 
 define i8 @set_bits(i8 %x, i1 %b)  {
 ; CHECK-LABEL: @set_bits(
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X:%.*]], -6
-; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], 5
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[B:%.*]], i8 [[OR]], i8 [[AND]]
+; CHECK-NEXT:    [[MASKSEL:%.*]] = select i1 [[B:%.*]], i8 5, i8 0
+; CHECK-NEXT:    [[COND:%.*]] = or i8 [[AND]], [[MASKSEL]]
 ; CHECK-NEXT:    ret i8 [[COND]]
 ;
   %and = and i8 %x, 250
   %or = or i8 %x, 5
   %cond = select i1 %b, i8 %or, i8 %and
   ret i8 %cond
 }
 
+; Negative test
+
 define i8 @set_bits_not_inverse_constant(i8 %x, i1 %b)  {
 ; CHECK-LABEL: @set_bits_not_inverse_constant(
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X:%.*]], -6
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], 7
 ; CHECK-NEXT:    [[COND:%.*]] = select i1 [[B:%.*]], i8 [[OR]], i8 [[AND]]
 ; CHECK-NEXT:    ret i8 [[COND]]
 ;
   %and = and i8 %x, 250
   %or = or i8 %x, 7
   %cond = select i1 %b, i8 %or, i8 %and
   ret i8 %cond
 }
 
 define i8 @set_bits_extra_use1(i8 %x, i1 %b)  {
 ; CHECK-LABEL: @set_bits_extra_use1(
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X:%.*]], -6
 ; CHECK-NEXT:    call void @use(i8 [[AND]])
-; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], 5
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[B:%.*]], i8 [[OR]], i8 [[AND]]
+; CHECK-NEXT:    [[MASKSEL:%.*]] = select i1 [[B:%.*]], i8 5, i8 0
+; CHECK-NEXT:    [[COND:%.*]] = or i8 [[AND]], [[MASKSEL]]
 ; CHECK-NEXT:    ret i8 [[COND]]
 ;
   %and = and i8 %x, 250
   call void @use(i8 %and)
   %or = or i8 %x, 5
   %cond = select i1 %b, i8 %or, i8 %and
   ret i8 %cond
 }
 
+; Negative test
+
 define i8 @set_bits_extra_use2(i8 %x, i1 %b)  {
 ; CHECK-LABEL: @set_bits_extra_use2(
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X:%.*]], -6
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], 5
 ; CHECK-NEXT:    call void @use(i8 [[OR]])
 ; CHECK-NEXT:    [[COND:%.*]] = select i1 [[B:%.*]], i8 [[OR]], i8 [[AND]]
 ; CHECK-NEXT:    ret i8 [[COND]]
 ;
   %and = and i8 %x, 250
   %or = or i8 %x, 5
   call void @use(i8 %or)
   %cond = select i1 %b, i8 %or, i8 %and
   ret i8 %cond
 }
 
 define <2 x i8> @clear_bits(<2 x i8> %x, <2 x i1> %b)  {
 ; CHECK-LABEL: @clear_bits(
 ; CHECK-NEXT:    [[AND:%.*]] = and <2 x i8> [[X:%.*]], <i8 37, i8 37>
-; CHECK-NEXT:    [[OR:%.*]] = or <2 x i8> [[X]], <i8 -38, i8 -38>
-; CHECK-NEXT:    [[COND:%.*]] = select <2 x i1> [[B:%.*]], <2 x i8> [[AND]], <2 x i8> [[OR]]
+; CHECK-NEXT:    [[MASKSEL:%.*]] = select <2 x i1> [[B:%.*]], <2 x i8> zeroinitializer, <2 x i8> <i8 -38, i8 -38>
+; CHECK-NEXT:    [[COND:%.*]] = or <2 x i8> [[AND]], [[MASKSEL]]
 ; CHECK-NEXT:    ret <2 x i8> [[COND]]
 ;
   %and = and <2 x i8> %x, <i8 37, i8 37>
   %or = or <2 x i8> %x, <i8 218, i8 218>
   %cond = select <2 x i1> %b, <2 x i8> %and, <2 x i8> %or
   ret <2 x i8> %cond
 }
 
+; Negative test
+
 define <2 x i8> @clear_bits_not_inverse_constant(<2 x i8> %x, <2 x i1> %b)  {
 ; CHECK-LABEL: @clear_bits_not_inverse_constant(
 ; CHECK-NEXT:    [[AND:%.*]] = and <2 x i8> [[X:%.*]], <i8 undef, i8 37>
 ; CHECK-NEXT:    [[OR:%.*]] = or <2 x i8> [[X]], <i8 -38, i8 -38>
 ; CHECK-NEXT:    [[COND:%.*]] = select <2 x i1> [[B:%.*]], <2 x i8> [[AND]], <2 x i8> [[OR]]
 ; CHECK-NEXT:    ret <2 x i8> [[COND]]
 ;
   %and = and <2 x i8> %x, <i8 undef, i8 37>
   %or = or <2 x i8> %x, <i8 218, i8 218>
   %cond = select <2 x i1> %b, <2 x i8> %and, <2 x i8> %or
   ret <2 x i8> %cond
 }
 
 define <2 x i8> @clear_bits_extra_use1(<2 x i8> %x, i1 %b)  {
 ; CHECK-LABEL: @clear_bits_extra_use1(
 ; CHECK-NEXT:    [[AND:%.*]] = and <2 x i8> [[X:%.*]], <i8 37, i8 37>
 ; CHECK-NEXT:    call void @use_vec(<2 x i8> [[AND]])
-; CHECK-NEXT:    [[OR:%.*]] = or <2 x i8> [[X]], <i8 -38, i8 -38>
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[B:%.*]], <2 x i8> [[AND]], <2 x i8> [[OR]]
+; CHECK-NEXT:    [[MASKSEL:%.*]] = select i1 [[B:%.*]], <2 x i8> zeroinitializer, <2 x i8> <i8 -38, i8 -38>
+; CHECK-NEXT:    [[COND:%.*]] = or <2 x i8> [[AND]], [[MASKSEL]]
 ; CHECK-NEXT:    ret <2 x i8> [[COND]]
 ;
   %and = and <2 x i8> %x, <i8 37, i8 37>
   call void @use_vec(<2 x i8> %and)
   %or = or <2 x i8> %x, <i8 218, i8 218>
   %cond = select i1 %b, <2 x i8> %and, <2 x i8> %or
   ret <2 x i8> %cond
 }
 
+; Negative test
+
 define i8 @clear_bits_extra_use2(i8 %x, i1 %b)  {
 ; CHECK-LABEL: @clear_bits_extra_use2(
 ; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X:%.*]], -6
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[X]], 5
 ; CHECK-NEXT:    call void @use(i8 [[OR]])
 ; CHECK-NEXT:    [[COND:%.*]] = select i1 [[B:%.*]], i8 [[AND]], i8 [[OR]]
 ; CHECK-NEXT:    ret i8 [[COND]]
 ;
   %and = and i8 %x, 250
   %or = or i8 %x, 5
   call void @use(i8 %or)
   %cond = select i1 %b, i8 %and, i8 %or
   ret i8 %cond
 }