[InstCombine] Simplify shift-by-sext to shift-by-zext - ignore use count on sext

[lebedev.ri]

Summary

In D68103 the InstCombine learned that shift-by-sext is simply a shift-by-zext.
But the transform is limited to single-use sext.
That is too limiting. We should handle cases with extra uses.
Here it is proposed to just always do the transformation,
on the basis that zext is unarguably better for any further analysis.

There are other alternatives

• Simply ensure that all uses of this sext are shifts - O(N^2) in case if there are non-shift uses
• Add more greedy version of the fold into AggressiveInstCombine - checking that all uses can be converted to zext - could be as simple as again simply checking that they are all shifts. Will not have O(N^2) problem, but will have usability problem - AggressiveInstCombine only runs in -O3, not even in -O2?

So i suspect this might be the least of the evils..

[lebedev.ri]

Also perform sudo-CSE to avoid creating numerous zexts and leaving them up to the next time EarlyCSE.
This obviously doesn't help with pre-existing zexts..

[lebedev.ri]

https://lists.llvm.org/pipermail/llvm-dev/2019-September/135494.html

[lebedev.ri]

@efriedma any thoughts on this?

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

Instruction *replaceInstUsesWith(Instruction &I, Value *V) {
 
     LLVM_DEBUG(dbgs() << "IC: Replacing " << I << "\n"
                       << "    with " << *V << '\n');
 
     I.replaceAllUsesWith(V);
     return &I;
   }
 
+  /// A combiner-aware RAUW-like routine.
+  ///
+  /// This method is to be used when an instruction is found to be partially
+  /// dead, where selected uses can be adjusted to use another preexisting
+  /// expression. Here we add *all* uses of I to the worklist,
+  /// conditionally replace all uses of I with the new value, then return I,
+  /// so that the inst combiner will know that I was modified.
+  Instruction *
+  replaceInstUsesWithIf(Instruction &I, Value *V,
+                        llvm::function_ref<bool(Use &U)> ShouldReplace) {
+    // If there are no uses to replace, then we return nullptr to indicate that
+    // no changes were made to the program.
+    if (I.use_empty())
+      return nullptr;
+
+    Worklist.AddUsersToWorkList(I); // Add *all* instrs to worklist.
+
+    // If we are replacing the instruction with itself, this must be in a
+    // segment of unreachable code, so just clobber the instruction.
+    if (&I == V)
+      V = UndefValue::get(I.getType());
+
+    LLVM_DEBUG(dbgs() << "IC: Conditionally replacing " << I << "\n"
+                      << "    with " << *V << '\n');
+
+    I.replaceUsesWithIf(V, ShouldReplace);
+    return &I;
+  }
+
   /// Creates a result tuple for an overflow intrinsic \p II with a given
   /// \p Result and a constant \p Overflow value.
   Instruction *CreateOverflowTuple(IntrinsicInst *II, Value *Result,
                                    Constant *Overflow) {
     Constant *V[] = {UndefValue::get(Result->getType()), Overflow};
     StructType *ST = cast<StructType>(II->getType());
     Constant *Struct = ConstantStruct::get(ST, V);
     return InsertValueInst::Create(Struct, Result, 0);

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

dropRedundantMaskingOfLeftShiftInput(BinaryOperator *OuterShift,
   Builder.Insert(NewShift);
   return BinaryOperator::Create(Instruction::And, NewShift, NewMask);
 }
 
 Instruction *InstCombiner::commonShiftTransforms(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   assert(Op0->getType() == Op1->getType());
 
-  // If the shift amount is a one-use `sext`, we can demote it to `zext`.
+  // If the shift amount is a `sext`, we can demote it to `zext`.
+  // We don't limit this to single-use `sext`, or ensure that we can change
+  // all uses of `sext` to `zext`, since we believe that `zext` is *much*
+  // better for further analysis. We also perform sudo-CSE to avoid producing
+  // multiple `zext` per each eligible user and obscruing the fact that their
+  // shift amounts are actually the same Value. Though that does not help if
+  // there was a pre-existing identical `zext`, that one is left in place.
   Value *Y;
-  if (match(Op1, m_OneUse(m_SExt(m_Value(Y))))) {
-    Value *NewExt = Builder.CreateZExt(Y, I.getType(), Op1->getName());
-    return BinaryOperator::Create(I.getOpcode(), Op0, NewExt);
+  if (isa<Instruction>(Op1) && match(Op1, m_SExt(m_Value(Y)))) {
+    BuilderTy::InsertPointGuard Guard(Builder);
+    Instruction *OldExt = cast<Instruction>(Op1);
+    Builder.SetInsertPoint(OldExt);
+    Value *NewExt = Builder.CreateZExt(Y, I.getType(), Y->getName() + ".zext");
+    replaceInstUsesWithIf(*OldExt, NewExt, [](Use &U) {
+      return match(U.getUser(), m_Shift(m_Value(), m_Value()));
+    });
+    return &I;
   }
 
   // See if we can fold away this shift.
   if (SimplifyDemandedInstructionBits(I))
     return &I;
 
   // Try to fold constant and into select arguments.
   if (isa<Constant>(Op0))

llvm/test/Transforms/InstCombine/shift-by-signext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -instcombine -S | FileCheck %s
 
 ; If we have a shift by sign-extended value, we can replace sign-extension
 ; with zero-extension.
 
 define i32 @t0_shl(i32 %x, i8 %shamt) {
 ; CHECK-LABEL: @t0_shl(
-; CHECK-NEXT:    [[SHAMT_WIDE1:%.*]] = zext i8 [[SHAMT:%.*]] to i32
-; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[SHAMT_WIDE1]]
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shamt_wide = sext i8 %shamt to i32
   %r = shl i32 %x, %shamt_wide
   ret i32 %r
 }
 define i32 @t1_lshr(i32 %x, i8 %shamt) {
 ; CHECK-LABEL: @t1_lshr(
-; CHECK-NEXT:    [[SHAMT_WIDE1:%.*]] = zext i8 [[SHAMT:%.*]] to i32
-; CHECK-NEXT:    [[R:%.*]] = lshr i32 [[X:%.*]], [[SHAMT_WIDE1]]
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[R:%.*]] = lshr i32 [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shamt_wide = sext i8 %shamt to i32
   %r = lshr i32 %x, %shamt_wide
   ret i32 %r
 }
 define i32 @t2_ashr(i32 %x, i8 %shamt) {
 ; CHECK-LABEL: @t2_ashr(
-; CHECK-NEXT:    [[SHAMT_WIDE1:%.*]] = zext i8 [[SHAMT:%.*]] to i32
-; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[X:%.*]], [[SHAMT_WIDE1]]
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shamt_wide = sext i8 %shamt to i32
   %r = ashr i32 %x, %shamt_wide
   ret i32 %r
 }
 
 define <2 x i32> @t3_vec_shl(<2 x i32> %x, <2 x i8> %shamt) {
 ; CHECK-LABEL: @t3_vec_shl(
-; CHECK-NEXT:    [[SHAMT_WIDE1:%.*]] = zext <2 x i8> [[SHAMT:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = shl <2 x i32> [[X:%.*]], [[SHAMT_WIDE1]]
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext <2 x i8> [[SHAMT:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i32> [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %shamt_wide = sext <2 x i8> %shamt to <2 x i32>
   %r = shl <2 x i32> %x, %shamt_wide
   ret <2 x i32> %r
 }
 define <2 x i32> @t4_vec_lshr(<2 x i32> %x, <2 x i8> %shamt) {
 ; CHECK-LABEL: @t4_vec_lshr(
-; CHECK-NEXT:    [[SHAMT_WIDE1:%.*]] = zext <2 x i8> [[SHAMT:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i32> [[X:%.*]], [[SHAMT_WIDE1]]
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext <2 x i8> [[SHAMT:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i32> [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %shamt_wide = sext <2 x i8> %shamt to <2 x i32>
   %r = lshr <2 x i32> %x, %shamt_wide
   ret <2 x i32> %r
 }
 define <2 x i32> @t5_vec_ashr(<2 x i32> %x, <2 x i8> %shamt) {
 ; CHECK-LABEL: @t5_vec_ashr(
-; CHECK-NEXT:    [[SHAMT_WIDE1:%.*]] = zext <2 x i8> [[SHAMT:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = ashr <2 x i32> [[X:%.*]], [[SHAMT_WIDE1]]
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext <2 x i8> [[SHAMT:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[R:%.*]] = ashr <2 x i32> [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %shamt_wide = sext <2 x i8> %shamt to <2 x i32>
   %r = ashr <2 x i32> %x, %shamt_wide
   ret <2 x i32> %r
 }
 
 define i32 @t6_twoshifts(i32 %x, i8 %shamt) {
 ; CHECK-LABEL: @t6_twoshifts(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[SHAMT_WIDE:%.*]] = sext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT:%.*]] to i32
 ; CHECK-NEXT:    br label [[WORK:%.*]]
 ; CHECK:       work:
 ; CHECK-NEXT:    br label [[END:%.*]]
 ; CHECK:       end:
-; CHECK-NEXT:    [[N0:%.*]] = shl i32 [[X:%.*]], [[SHAMT_WIDE]]
-; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[N0]], [[SHAMT_WIDE]]
+; CHECK-NEXT:    [[N0:%.*]] = shl i32 [[X:%.*]], [[SHAMT_ZEXT]]
+; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[N0]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
 bb:
   %shamt_wide = sext i8 %shamt to i32
   br label %work
 work:
   %n0 = shl i32 %x, %shamt_wide
   %r = ashr i32 %n0, %shamt_wide
   br label %end
 end:
   ret i32 %r
 }
 
+define i32 @t6_twoshifts_2(i32 %x, i8 %shamt) {
+; CHECK-LABEL: @t6_twoshifts_2(
+; CHECK-NEXT:    [[SHAMT_WIDE_UNSIGNED:%.*]] = zext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT]] to i32
+; CHECK-NEXT:    [[N0:%.*]] = shl i32 [[X:%.*]], [[SHAMT_WIDE_UNSIGNED]]
+; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[N0]], [[SHAMT_ZEXT]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %shamt_wide_unsigned = zext i8 %shamt to i32
+  %shamt_wide_signed = sext i8 %shamt to i32
+  %n0 = shl i32 %x, %shamt_wide_unsigned
+  %r = ashr i32 %n0, %shamt_wide_signed
+  ret i32 %r
+}
+
 ; This is not valid for funnel shifts in general
 declare i7 @llvm.fshl.i7(i7 %a, i7 %b, i7 %c)
 declare i7 @llvm.fshr.i7(i7 %a, i7 %b, i7 %c)
 define i7 @n7_fshl(i7 %x, i7 %y, i6 %shamt) {
 ; CHECK-LABEL: @n7_fshl(
 ; CHECK-NEXT:    [[SHAMT_WIDE:%.*]] = sext i6 [[SHAMT:%.*]] to i7
 ; CHECK-NEXT:    [[R:%.*]] = call i7 @llvm.fshl.i7(i7 [[X:%.*]], i7 [[Y:%.*]], i7 [[SHAMT_WIDE]])
 ; CHECK-NEXT:    ret i7 [[R]]
;
   %shamt_wide = sext i6 %shamt to i8
   %r = call i8 @llvm.fshr.i8(i8 %x, i8 %y, i8 %shamt_wide)
   ret i8 %r
 }
 
 declare void @use32(i32)
 define i32 @n11_extrause(i32 %x, i8 %shamt) {
 ; CHECK-LABEL: @n11_extrause(
-; CHECK-NEXT:    [[SHAMT_WIDE:%.*]] = sext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[SHAMT_WIDE:%.*]] = sext i8 [[SHAMT]] to i32
 ; CHECK-NEXT:    call void @use32(i32 [[SHAMT_WIDE]])
-; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[SHAMT_WIDE]]
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[X:%.*]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shamt_wide = sext i8 %shamt to i32
   call void @use32(i32 %shamt_wide)
   %r = shl i32 %x, %shamt_wide
   ret i32 %r
 }
 define i32 @n12_twoshifts_and_extrause(i32 %x, i8 %shamt) {
 ; CHECK-LABEL: @n12_twoshifts_and_extrause(
-; CHECK-NEXT:    [[SHAMT_WIDE:%.*]] = sext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[SHAMT_ZEXT:%.*]] = zext i8 [[SHAMT:%.*]] to i32
+; CHECK-NEXT:    [[SHAMT_WIDE:%.*]] = sext i8 [[SHAMT]] to i32
 ; CHECK-NEXT:    br label [[WORK:%.*]]
 ; CHECK:       work:
 ; CHECK-NEXT:    br label [[END:%.*]]
 ; CHECK:       end:
-; CHECK-NEXT:    [[N0:%.*]] = shl i32 [[X:%.*]], [[SHAMT_WIDE]]
-; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[N0]], [[SHAMT_WIDE]]
+; CHECK-NEXT:    [[N0:%.*]] = shl i32 [[X:%.*]], [[SHAMT_ZEXT]]
+; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[N0]], [[SHAMT_ZEXT]]
 ; CHECK-NEXT:    call void @use32(i32 [[SHAMT_WIDE]])
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %shamt_wide = sext i8 %shamt to i32
   br label %work
 work:
   %n0 = shl i32 %x, %shamt_wide
   %r = ashr i32 %n0, %shamt_wide
   br label %end
 end:
   call void @use32(i32 %shamt_wide)
   ret i32 %r
 }

llvm/test/Transforms/PhaseOrdering/two-shifts-by-sext.ll

 ; int two_shifts_by_same_sext_with_extra_use(int val, char len) {
 ;   int wide_len = len;
 ;   use_int32(wide_len);
 ;   return (val << wide_len) >> wide_len;
 ; }
 
 define i32 @two_shifts_by_sext(i32 %val, i8 signext %len) {
 ; CHECK-LABEL: @two_shifts_by_sext(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[LEN:%.*]] to i32
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[CONV]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[CONV]]
+; CHECK-NEXT:    [[LEN_ZEXT:%.*]] = zext i8 [[LEN:%.*]] to i32
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[LEN_ZEXT]]
+; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[LEN_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %val.addr = alloca i32, align 4
   %len.addr = alloca i8, align 1
   store i32 %val, i32* %val.addr, align 4
   store i8 %len, i8* %len.addr, align 1
   %val.reloaded = load i32, i32* %val.addr, align 4
   %len.reloaded.0 = load i8, i8* %len.addr, align 1
   %conv = sext i8 %len.reloaded.0 to i32
   %shl = shl i32 %val.reloaded, %conv
   %len.reloaded.1 = load i8, i8* %len.addr, align 1
   %conv1 = sext i8 %len.reloaded.1 to i32
   %shr = ashr i32 %shl, %conv1
   ret i32 %shr
 }
 
 define i32 @two_shifts_by_same_sext(i32 %val, i8 signext %len) {
 ; CHECK-LABEL: @two_shifts_by_same_sext(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[LEN:%.*]] to i32
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[CONV]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[CONV]]
+; CHECK-NEXT:    [[LEN_ZEXT:%.*]] = zext i8 [[LEN:%.*]] to i32
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[LEN_ZEXT]]
+; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[LEN_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %val.addr = alloca i32, align 4
   %len.addr = alloca i8, align 1
   %wide_len = alloca i32, align 4
   store i32 %val, i32* %val.addr, align 4
   store i8 %len, i8* %len.addr, align 1
   %len.reloaded.0 = load i8, i8* %len.addr, align 1
   %conv = sext i8 %len.reloaded.0 to i32
   store i32 %conv, i32* %wide_len, align 4
   %val.reloaded = load i32, i32* %val.addr, align 4
   %len.reloaded.1 = load i32, i32* %wide_len, align 4
   %shl = shl i32 %val.reloaded, %len.reloaded.1
   %len.reloaded.2 = load i32, i32* %wide_len, align 4
   %shr = ashr i32 %shl, %len.reloaded.2
   ret i32 %shr
 }
 
 define i32 @two_shifts_by_sext_with_extra_use(i32 %val, i8 signext %len) {
 ; CHECK-LABEL: @two_shifts_by_sext_with_extra_use(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[LEN:%.*]] to i32
+; CHECK-NEXT:    [[LEN_ZEXT:%.*]] = zext i8 [[LEN:%.*]] to i32
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[LEN]] to i32
 ; CHECK-NEXT:    tail call void @use_int32(i32 [[CONV]])
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[CONV]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[CONV]]
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[LEN_ZEXT]]
+; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[LEN_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %val.addr = alloca i32, align 4
   %len.addr = alloca i8, align 1
   store i32 %val, i32* %val.addr, align 4
   store i8 %len, i8* %len.addr, align 1
   %len.reloaded.0 = load i8, i8* %len.addr, align 1
   %conv = sext i8 %len.reloaded.0 to i32
   call void @use_int32(i32 %conv)
   %val.reloaded = load i32, i32* %val.addr, align 4
   %len.reloaded.1 = load i8, i8* %len.addr, align 1
   %conv1 = sext i8 %len.reloaded.1 to i32
   %shl = shl i32 %val.reloaded, %conv1
   %len.reloaded.2 = load i8, i8* %len.addr, align 1
   %conv2 = sext i8 %len.reloaded.2 to i32
   %shr = ashr i32 %shl, %conv2
   ret i32 %shr
 }
 
 declare void @use_int32(i32)
 
 define i32 @two_shifts_by_same_sext_with_extra_use(i32 %val, i8 signext %len) {
 ; CHECK-LABEL: @two_shifts_by_same_sext_with_extra_use(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[LEN:%.*]] to i32
+; CHECK-NEXT:    [[LEN_ZEXT:%.*]] = zext i8 [[LEN:%.*]] to i32
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[LEN]] to i32
 ; CHECK-NEXT:    tail call void @use_int32(i32 [[CONV]])
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[CONV]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[CONV]]
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL:%.*]], [[LEN_ZEXT]]
+; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[SHL]], [[LEN_ZEXT]]
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %val.addr = alloca i32, align 4
   %len.addr = alloca i8, align 1
   %wide_len = alloca i32, align 4
   store i32 %val, i32* %val.addr, align 4
   store i8 %len, i8* %len.addr, align 1
   %len.reloaded.0 = load i8, i8* %len.addr, align 1