[LICM][WIP] Transform and hoist select instructions if possible

[StephenFan]

Summary

Try to hoist select instruction that match the following and
%c and %cond are loop invariants:
Loop:

%x = add i32 %y, %c
%z = select i1 %cond, i32 %y, i32 %x

>

OutOfLoop:

%a = select i1 %cond, i32 0, i32 %c
...

Loop:

%z = add i32 %y, %a

TBH, I want to optimize the program:

define dso_local i32 @test(i32 noundef %x) local_unnamed_addr #0 {
entry:
  br label %for.cond

for.cond:                                         ; preds = %for.inc, %entry
  %y.0 = phi i32 [ 0, %entry ], [ %y.1, %for.inc ]
  %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
  %cmp = icmp ult i32 %i.0, 1000
  br i1 %cmp, label %for.body, label %for.cond.cleanup

for.cond.cleanup:                                 ; preds = %for.cond
  ret i32 %y.0

for.body:                                         ; preds = %for.cond
  %tobool.not = icmp eq i32 %x, 0
  br i1 %tobool.not, label %for.inc, label %if.then

if.then:                                          ; preds = %for.body
  %add = add nsw i32 %y.0, 10
  br label %for.inc

for.inc:                                          ; preds = %for.body, %if.then
  %y.1 = phi i32 [ %add, %if.then ], [ %y.0, %for.body ]
  %inc = add nuw nsw i32 %i.0, 1
  br label %for.cond
}

Idealy, it can be optimized to

define dso_local i32 @test(i32 noundef %x) local_unnamed_addr #0 {
entry:
  %tobool.not = icmp eq i32 %x, 0
  %0 = select i1 %tobool.not, i32 0, i32 10000
  ret i32 %0
}

But SimplifyCFG transforms for.body and if.then to a select instruction. Which is:

define dso_local i32 @test(i32 noundef %x) local_unnamed_addr {
entry:
  br label %for.cond

for.cond:                                         ; preds = %for.body, %entry
  %y.0 = phi i32 [ 0, %entry ], [ %spec.select, %for.body ]
  %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
  %cmp = icmp ult i32 %i.0, 1000
  br i1 %cmp, label %for.body, label %for.cond.cleanup

for.cond.cleanup:                                 ; preds = %for.cond
  ret i32 %y.0

for.body:                                         ; preds = %for.cond
  %tobool.not = icmp eq i32 %x, 0
  %add = add nsw i32 %y.0, 10
  %spec.select = select i1 %tobool.not, i32 %y.0, i32 %add
  %inc = add nuw nsw i32 %i.0, 1
  br label %for.cond
}

This prevents it from being loop unswitched. In order to get the expected result, this patch hoisted the select instruction, so that indvars pass can rewrite the exit value.

[nikic]

Shouldn't this be handled by adding this fold to InstCombine? https://alive2.llvm.org/ce/z/EXY5JH

[StephenFan]

In D140974#4025888, @nikic wrote:

Shouldn't this be handled by adding this fold to InstCombine? https://alive2.llvm.org/ce/z/EXY5JH

Thanks for your advice! I will try to handle this in InstCombine. But it seems it is not a beneficial transformation if we transform

define i8 @src(i1 %c, i8 %x) {
  %add = add i8 %x, 10
  %select = select i1 %c, i8 %x, i8 %add
  ret i8 %select
}

to

define i8 @tgt(i1 %c, i8 %x) {
  %select = select i1 %c, i8 0, i8 10
  %add = add i8 %select, %x
  ret i8 %add
}

https://godbolt.org/z/zxeojTxqb

[nikic]

In D140974#4026009, @StephenFan wrote:

In D140974#4025888, @nikic wrote:

Shouldn't this be handled by adding this fold to InstCombine? https://alive2.llvm.org/ce/z/EXY5JH

Thanks for your advice! I will try to handle this in InstCombine. But it seems it is not a beneficial transformation if we transform

define i8 @src(i1 %c, i8 %x) {
  %add = add i8 %x, 10
  %select = select i1 %c, i8 %x, i8 %add
  ret i8 %select
}

to

define i8 @tgt(i1 %c, i8 %x) {
  %select = select i1 %c, i8 0, i8 10
  %add = add i8 %select, %x
  ret i8 %add
}

https://godbolt.org/z/zxeojTxqb

This looks beneficial from a middle-end perspective, so we probably need to undo in the backend? @spatel Any thoughts?

[spatel]

We already have that fold?
https://alive2.llvm.org/ce/z/XPZL38

There is already an undo of that or something similar for FP ops on some targets, so it may need some generalization.

[nikic]

In D140974#4026210, @spatel wrote:

We already have that fold?
https://alive2.llvm.org/ce/z/XPZL38

There is already an undo of that or something similar for FP ops on some targets, so it may need some generalization.

Ooops, you're right. I could have sworn I checked that first...

So to return to the original problem here, there may be a phase ordering issue -- can you please add a PhaseOrdering test?

[spatel]

In D140974#4026218, @nikic wrote:

In D140974#4026210, @spatel wrote:

We already have that fold?
https://alive2.llvm.org/ce/z/XPZL38

There is already an undo of that or something similar for FP ops on some targets, so it may need some generalization.

Ooops, you're right. I could have sworn I checked that first...

So to return to the original problem here, there may be a phase ordering issue -- can you please add a PhaseOrdering test?

Aha - there's a limitation for the example with a constant, so it doesn't fold:
https://alive2.llvm.org/ce/z/QfjhJP

I think it's from this (not sure why we have it):
https://github.com/llvm/llvm-project/blob/85960043d594fc12d340ccb66a30861b023ab496/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp#L492

[nikic]

Per above comments, I believe we should 1) remove the InstCombine limitation that @spatel found and 2) undo in the backend if needed. I'd rather not have this kind of special case sitting in LICM.

llvm/lib/Transforms/Scalar/LICM.cpp

for (Instruction &I : llvm::make_early_inc_range(*BB)) {
 
         hoist(*ReciprocalDivisor, DT, CurLoop, CFH.getOrCreateHoistedBlock(BB),
               SafetyInfo, MSSAU, SE, ORE);
         HoistedInstructions.push_back(ReciprocalDivisor);
         Changed = true;
         continue;
       }
 
+      // Try to hoist select instruction that match the following and
+      // %c and %cond is loop invariant:
+      // Loop:
+      //   %x = add i32 %y, %c
+      //   %z = select i1 %cond, i32 %y, i32 %x
+      // ==>
+      // OutOfLoop:
+      //   %a = select i1 %cond, i32 0, i32 %c
+      //   ...
+      // Loop:
+      //   %z = add i32 %y, %a
+      if (I.getOpcode() == Instruction::Select) {
+        using namespace PatternMatch;
+        Instruction *BO = nullptr;
+        Value *B = nullptr;
+        Value *C = I.getOperand(0);
+        // TODO: Support other binary instructions.
+        if (match(I.getOperand(1),
+                  m_OneUse(m_Add(m_Specific(I.getOperand(2)), m_Value(B)))))
+          BO = cast<Instruction>(I.getOperand(1));
+        else if (match(
+                     I.getOperand(2),
+                     m_OneUse(m_Add(m_Specific(I.getOperand(1)), m_Value(B)))))
+          BO = cast<Instruction>(I.getOperand(2));
+
+        if (BO && B && CurLoop->isLoopInvariant(C) &&
+            CurLoop->isLoopInvariant(B)) {
+          BasicBlock *HoistBB = CFH.getOrCreateHoistedBlock(BB);
+          auto *ToAdd =
+              SelectInst::Create(C, ConstantInt::get(I.getType(), 0), B);
+          ToAdd->insertBefore(HoistBB->getTerminator());
+          BO->setOperand(1, ToAdd);
+          I.replaceAllUsesWith(BO);
+          I.eraseFromParent();
+          if (SE) {
+            SE->forgetValue(BO);
+            SE->forgetValue(&I);
+          }
+        }
+      }
+
       auto IsInvariantStart = [&](Instruction &I) {
         using namespace PatternMatch;
         return I.use_empty() &&
                match(&I, m_Intrinsic<Intrinsic::invariant_start>());
       };
       auto MustExecuteWithoutWritesBefore = [&](Instruction &I) {
         return SafetyInfo->isGuaranteedToExecute(I, DT, CurLoop) &&
                SafetyInfo->doesNotWriteMemoryBefore(I, CurLoop);

llvm/test/Transforms/LICM/select.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=licm -S %s | FileCheck %s
+;
+define dso_local i32 @add_constant(i32 %x) {
+; CHECK-LABEL: @add_constant(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = select i1 [[TOBOOL_NOT]], i32 0, i32 10
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    [[Y_1_LCSSA:%.*]] = phi i32 [ [[ADD:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    ret i32 [[Y_1_LCSSA]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_04:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[Y_03:%.*]] = phi i32 [ 10, [[ENTRY]] ], [ [[ADD]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ADD]] = add nsw i32 [[Y_03]], [[TMP0]]
+; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I_04]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], 1000
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  %tobool.not = icmp eq i32 %x, 0
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret i32 %y.1
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.04 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %y.03 = phi i32 [ 10, %entry ], [ %y.1, %for.body ]
+  %add = add nsw i32 %y.03, 10
+  %y.1 = select i1 %tobool.not, i32 %y.03, i32 %add
+  %inc = add nuw nsw i32 %i.04, 1
+  %exitcond.not = icmp eq i32 %inc, 1000
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+}
+
+define dso_local i32 @add_constant_commute(i32 %x) {
+; CHECK-LABEL: @add_constant_commute(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = select i1 [[TOBOOL_NOT]], i32 0, i32 10
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    [[Y_1_LCSSA:%.*]] = phi i32 [ [[ADD:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    ret i32 [[Y_1_LCSSA]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_04:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[Y_03:%.*]] = phi i32 [ 10, [[ENTRY]] ], [ [[ADD]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ADD]] = add nsw i32 [[Y_03]], [[TMP0]]
+; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I_04]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], 1000
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  %tobool.not = icmp eq i32 %x, 0
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret i32 %y.1
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.04 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %y.03 = phi i32 [ 10, %entry ], [ %y.1, %for.body ]
+  %add = add nsw i32 %y.03, 10
+  %y.1 = select i1 %tobool.not, i32 %add, i32 %y.03
+  %inc = add nuw nsw i32 %i.04, 1
+  %exitcond.not = icmp eq i32 %inc, 1000
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+}
+
+define dso_local i32 @add_loop_invariant_but_non_constant(i32 %x, i32 %xx) {
+; CHECK-LABEL: @add_loop_invariant_but_non_constant(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = select i1 [[TOBOOL_NOT]], i32 0, i32 [[XX:%.*]]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    [[Y_1_LCSSA:%.*]] = phi i32 [ [[ADD:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    ret i32 [[Y_1_LCSSA]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_04:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[Y_03:%.*]] = phi i32 [ 10, [[ENTRY]] ], [ [[ADD]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ADD]] = add nsw i32 [[Y_03]], [[TMP0]]
+; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I_04]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], 1000
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  %tobool.not = icmp eq i32 %x, 0
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret i32 %y.1
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.04 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %y.03 = phi i32 [ 10, %entry ], [ %y.1, %for.body ]
+  %add = add nsw i32 %y.03, %xx
+  %y.1 = select i1 %tobool.not, i32 %add, i32 %y.03
+  %inc = add nuw nsw i32 %i.04, 1
+  %exitcond.not = icmp eq i32 %inc, 1000
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+}
+
+declare void @use(i32)
+
+; Negative test, since %add has multiple uses.
+;
+define dso_local i32 @multiple_uses(i32 %x) {
+; CHECK-LABEL: @multiple_uses(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq i32 [[X:%.*]], 0
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    [[Y_1_LCSSA:%.*]] = phi i32 [ [[Y_1:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    ret i32 [[Y_1_LCSSA]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_04:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[Y_03:%.*]] = phi i32 [ 10, [[ENTRY]] ], [ [[Y_1]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[Y_03]], 10
+; CHECK-NEXT:    [[Y_1]] = select i1 [[TOBOOL_NOT]], i32 [[ADD]], i32 [[Y_03]]
+; CHECK-NEXT:    call void @use(i32 [[ADD]])
+; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I_04]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], 1000
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  %tobool.not = icmp eq i32 %x, 0
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret i32 %y.1
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.04 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %y.03 = phi i32 [ 10, %entry ], [ %y.1, %for.body ]
+  %add = add nsw i32 %y.03, 10
+  %y.1 = select i1 %tobool.not, i32 %add, i32 %y.03
+  call void @use(i32 %add)
+  %inc = add nuw nsw i32 %i.04, 1
+  %exitcond.not = icmp eq i32 %inc, 1000
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+}