[AggressiveInstCombine] Add arithmetic shift right instr to `TruncInstCombine` DAG

[anton-afanasyev]

Summary

Add `ashr` instruction to the DAG post-dominated by `trunc`, allowing                                         
`TruncInstCombine` to reduce bitwidth of expressions containing                                                          
these instructions.                                                                               
                                                                                                                          
We should be shifting by less than the target bitwidth.                                            
Also it is sufficient to require that all truncated bits                                                       
of the value-to-be-shifted are sign bits (all zeros or ones) and                                                         
one sign bit is left untruncated: https://alive2.llvm.org/ce/z/Ajo2__                                           
                                                                                                        
Part of https://reviews.llvm.org/D107766

[lebedev.ri]

I'm pretty sure we should count the number of sign bits here, much like how we count the number of leading zeros:
https://godbolt.org/z/vGsf664Gf => https://alive2.llvm.org/ce/z/gIc43E

[anton-afanasyev]

In D108355#2954130, @lebedev.ri wrote:

I'm pretty sure we should count the number of sign bits here, much like how we count the number of leading zeros:
https://godbolt.org/z/vGsf664Gf => https://alive2.llvm.org/ce/z/gIc43E

Ok, I've extended both lshr and ashr to use countMinSignBits(), unifying both cases, but we have to computeKnownBits() one more time here.

[anton-afanasyev]

Use countMinSignBits()

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

293–296

These seem like a separate change

411–412

[Why] do we have to do this?
Doesn't seem like something this transform should worry about?

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

297–307

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

293–296

Ok, I'm to split this

297–307

I've intentionally merged this two cases, since ashr and lshr processing doesn't differ. What really matters is positivity/negativity of their operand. This is symmetrical cases (see updated summary).

For instance, for your change we still have to replace ashr with lshr if sign bits were zeros and with ashr if sign bits were ones. Why not to do the same for lshr?

411–412

If sign bits were ones we have to replace original shift with ashr, lshr doesn't work (see @lshr_negative_operand_but_short() test).

[anton-afanasyev]

Split unrelated changes

[lebedev.ri]

This patch does not apply for me.
Please precommit the tests, and rebase the patch so that it applies to main.

[lebedev.ri]

Also, let's split trunc_shifts.ll into trunc_{shl,lshr,ashr}.ll.

[anton-afanasyev]

Rebase after splitted and precommited tests

[anton-afanasyev]

In D108355#2955549, @lebedev.ri wrote:

This patch does not apply for me.
Please precommit the tests, and rebase the patch so that it applies to main.

In D108355#2955569, @lebedev.ri wrote:

Also, let's split trunc_shifts.ll into trunc_{shl,lshr,ashr}.ll.

Ok, done

[lebedev.ri]

In D108355#2956676, @anton-afanasyev wrote:

Rebase after splitted and precommited tests

In D108355#2956680, @anton-afanasyev wrote:

In D108355#2955549, @lebedev.ri wrote:

This patch does not apply for me.
Please precommit the tests, and rebase the patch so that it applies to main.

In D108355#2955569, @lebedev.ri wrote:

Also, let's split trunc_shifts.ll into trunc_{shl,lshr,ashr}.ll.

Ok, done

Thank you!

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

297–307

Ehrm, these two inline comments are separate.
Can you explain why we should ask for known *bits* even for ashr,
even though we then ask for known *sign* bits?

411–412

I see. Please explain all that in a code comment.

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

297–307

(note that there's a typo in my diff, it should be
MinBitWidth = std::max(MinBitWidth, MinSignedBits); not
MinBitWidth = std::max(MinBitWidth, NumSignBits);)

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

411–412

But, this brings another question - should this assert that the sign bit is known?

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

411–412

Actually, let me backtrack this.
This patch adds support for ashr truncation.
It shouldn't also suddenly subtly change reasoning for lshr.
Can we please not do that?

[anton-afanasyev]

Make step back, leaving lshr untouched and counting sign bits for ashr only.
The motivation is such that for natural cases unsigned values are
zexted and lshred whereas signed values are sexted and ashred.
So it is naturally to grasp only these two cases. We miss here cases like truncation of ashr(zext(*)),
but it is transformed for AIC by preceding IC to lshr(zext(*)). Also remove unrelated tests.

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

297–307

Ok, thanks, used CountMinSignBits() for ashr.

411–412

Ok, I've changed that lines.

[anton-afanasyev]

Update comment

[anton-afanasyev]

Fix tests

[lebedev.ri]

This comment was removed by lebedev.ri.

[lebedev.ri]

LGTM, thank you.
@spatel ?

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

411–412

Can we please not do that?

Sure, I've eventually done it that way: lshr is transformed as before, ashr transformation works for numbers treated as signed ones.

[anton-afanasyev]

Rebase against dd19f342fa21

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

405–408

FWIW i think we can generalize this like that in the future.

[anton-afanasyev]

Update setting exactness

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

405–408

Thanks, done.

[lebedev.ri]

(i don't have any further comments here, but it might be good to wait on @spatel)

[spatel]

Sorry - I missed the earlier ping.
LGTM.

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

289–291

if (I->isShift()) {

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

289–291

Thanks, done

[anton-afanasyev]

Use isShift()

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

static void getRelevantOperands(Instruction *I, SmallVectorImpl<Value *> &Ops) {
   case Instruction::Add:
   case Instruction::Sub:
   case Instruction::Mul:
   case Instruction::And:
   case Instruction::Or:
   case Instruction::Xor:
   case Instruction::Shl:
   case Instruction::LShr:
+  case Instruction::AShr:
     Ops.push_back(I->getOperand(0));
     Ops.push_back(I->getOperand(1));
     break;
   case Instruction::Select:
     Ops.push_back(I->getOperand(1));
     Ops.push_back(I->getOperand(2));
     break;
   default:
while (!Worklist.empty()) {
     case Instruction::Add:
     case Instruction::Sub:
     case Instruction::Mul:
     case Instruction::And:
     case Instruction::Or:
     case Instruction::Xor:
     case Instruction::Shl:
     case Instruction::LShr:
+    case Instruction::AShr:
     case Instruction::Select: {
       SmallVector<Value *, 2> Operands;
       getRelevantOperands(I, Operands);
       append_range(Worklist, Operands);
       break;
     }
     default:
       // TODO: Can handle more cases here:
       // 1. shufflevector, extractelement, insertelement
       // 2. udiv, urem
-      // 3. ashr
-      // 4. phi node(and loop handling)
+      // 3. phi node(and loop handling)
       // ...
       return false;
     }
   }
   return true;
 }
 
 unsigned TruncInstCombine::getMinBitWidth() {
for (auto *U : I->users())
           DesiredBitWidth = ExtInstBitWidth;
         }
   }
 
   unsigned OrigBitWidth =
       CurrentTruncInst->getOperand(0)->getType()->getScalarSizeInBits();
 
   // Initialize MinBitWidth for shift instructions with the minimum number
-  // that is greater than shift amount (i.e. shift amount + 1). For `lshr`
-  // adjust MinBitWidth so that all potentially truncated bits of
-  // the value-to-be-shifted are zeros.
-  // Also normalize MinBitWidth not to be greater than source bitwidth.
+  // that is greater than shift amount (i.e. shift amount + 1).
+  // For `lshr` adjust MinBitWidth so that all potentially truncated
+  // bits of the value-to-be-shifted are zeros.
+  // For `ashr` adjust MinBitWidth so that all potentially truncated
+  // bits of the value-to-be-shifted are sign bits (all zeros or ones)
+  // and even one (first) untruncated bit is sign bit.
+  // Exit early if MinBitWidth is not less than original bitwidth.
   for (auto &Itr : InstInfoMap) {
     Instruction *I = Itr.first;
-    if (I->getOpcode() == Instruction::Shl ||
-        I->getOpcode() == Instruction::LShr) {
+    if (I->isShift()) {
       KnownBits KnownRHS = computeKnownBits(I->getOperand(1), DL);
       unsigned MinBitWidth = KnownRHS.getMaxValue()
                                  .uadd_sat(APInt(OrigBitWidth, 1))
                                  .getLimitedValue(OrigBitWidth);
       if (MinBitWidth == OrigBitWidth)
         return nullptr;
       if (I->getOpcode() == Instruction::LShr) {
         KnownBits KnownLHS = computeKnownBits(I->getOperand(0), DL);
         MinBitWidth =
             std::max(MinBitWidth, KnownLHS.getMaxValue().getActiveBits());
-        if (MinBitWidth >= OrigBitWidth)
-          return nullptr;
-      }
+      }
+      if (I->getOpcode() == Instruction::AShr) {
+        unsigned NumSignBits = ComputeNumSignBits(I->getOperand(0), DL);
+        MinBitWidth = std::max(MinBitWidth, OrigBitWidth - NumSignBits + 1);
+      }
       Itr.second.MinBitWidth = MinBitWidth;
     }
   }
 
   // Calculate minimum allowed bit-width allowed for shrinking the currently
   // visited truncate's operand.
   unsigned MinBitWidth = getMinBitWidth();
 
for (auto &Itr : InstInfoMap) { // Forward
     }
     case Instruction::Add:
     case Instruction::Sub:
     case Instruction::Mul:
     case Instruction::And:
     case Instruction::Or:
     case Instruction::Xor:
     case Instruction::Shl:
-    case Instruction::LShr: {
+    case Instruction::LShr:
+    case Instruction::AShr: {
       Value *LHS = getReducedOperand(I->getOperand(0), SclTy);
       Value *RHS = getReducedOperand(I->getOperand(1), SclTy);
       Res = Builder.CreateBinOp((Instruction::BinaryOps)Opc, LHS, RHS);
       // Preserve `exact` flag since truncation doesn't change exactness
-      if (Opc == Instruction::LShr)
+      if (auto *PEO = dyn_cast<PossiblyExactOperator>(I))
         if (auto *ResI = dyn_cast<Instruction>(Res))
-          ResI->setIsExact(I->isExact());
+          ResI->setIsExact(PEO->isExact());
       break;
     }
     case Instruction::Select: {
       Value *Op0 = I->getOperand(0);
       Value *LHS = getReducedOperand(I->getOperand(1), SclTy);
       Value *RHS = getReducedOperand(I->getOperand(2), SclTy);
       Res = Builder.CreateSelect(Op0, LHS, RHS);
       break;
     }
     default:
       llvm_unreachable("Unhandled instruction");
     }

llvm/test/Transforms/AggressiveInstCombine/trunc_ashr.ll

;
   %zext = zext i16 %x to i32
   %ashr = ashr i32 %zext, 15
   %trunc = trunc i32 %ashr to i16
   ret i16 %trunc
 }
 
 define i16 @ashr_sext_15(i16 %x) {
 ; CHECK-LABEL: @ashr_sext_15(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ASHR:%.*]] = ashr i32 [[SEXT]], 15
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[ASHR]] to i16
-; CHECK-NEXT:    ret i16 [[TRUNC]]
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr i16 [[X:%.*]], 15
+; CHECK-NEXT:    ret i16 [[ASHR]]
 ;
   %sext = sext i16 %x to i32
   %ashr = ashr i32 %sext, 15
   %trunc = trunc i32 %ashr to i16
   ret i16 %trunc
 }
 
 ; Negative test
;
   %a = add i32 %s, %z
   %s2 = ashr i32 %a, 2
   %t = trunc i32 %s2 to i16
   ret i16 %t
 }
 
 define i16 @ashr_var_bounded_shift_amount(i8 %x, i8 %amt) {
 ; CHECK-LABEL: @ashr_var_bounded_shift_amount(
-; CHECK-NEXT:    [[Z:%.*]] = zext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ZA:%.*]] = zext i8 [[AMT:%.*]] to i32
-; CHECK-NEXT:    [[ZA2:%.*]] = and i32 [[ZA]], 15
-; CHECK-NEXT:    [[S:%.*]] = ashr i32 [[Z]], [[ZA2]]
-; CHECK-NEXT:    [[A:%.*]] = add i32 [[S]], [[Z]]
-; CHECK-NEXT:    [[S2:%.*]] = ashr i32 [[A]], 2
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[S2]] to i16
-; CHECK-NEXT:    ret i16 [[T]]
+; CHECK-NEXT:    [[Z:%.*]] = zext i8 [[X:%.*]] to i16
+; CHECK-NEXT:    [[ZA:%.*]] = zext i8 [[AMT:%.*]] to i16
+; CHECK-NEXT:    [[ZA2:%.*]] = and i16 [[ZA]], 15
+; CHECK-NEXT:    [[S:%.*]] = ashr i16 [[Z]], [[ZA2]]
+; CHECK-NEXT:    [[A:%.*]] = add i16 [[S]], [[Z]]
+; CHECK-NEXT:    [[S2:%.*]] = ashr i16 [[A]], 2
+; CHECK-NEXT:    ret i16 [[S2]]
 ;
   %z = zext i8 %x to i32
   %za = zext i8 %amt to i32
   %za2 = and i32 %za, 15
   %s = ashr i32 %z, %za2
   %a = add i32 %s, %z
   %s2 = ashr i32 %a, 2
   %t = trunc i32 %s2 to i16
;
   %and = and i64 %sext, 4294967295
   %shl = ashr i64 %zext, %and
   %trunc = trunc i64 %shl to i32
   ret i32 %trunc
 }
 
 define void @ashr_big_dag(i16* %a, i8 %b, i8 %c) {
 ; CHECK-LABEL: @ashr_big_dag(
-; CHECK-NEXT:    [[ZEXT1:%.*]] = zext i8 [[B:%.*]] to i32
-; CHECK-NEXT:    [[ZEXT2:%.*]] = zext i8 [[C:%.*]] to i32
-; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[ZEXT1]], [[ZEXT2]]
-; CHECK-NEXT:    [[SFT1:%.*]] = and i32 [[ADD1]], 15
-; CHECK-NEXT:    [[SHR1:%.*]] = ashr i32 [[ADD1]], [[SFT1]]
-; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[ADD1]], [[SHR1]]
-; CHECK-NEXT:    [[SFT2:%.*]] = and i32 [[ADD2]], 7
-; CHECK-NEXT:    [[SHR2:%.*]] = ashr i32 [[ADD2]], [[SFT2]]
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[SHR2]] to i16
-; CHECK-NEXT:    store i16 [[TRUNC]], i16* [[A:%.*]], align 2
+; CHECK-NEXT:    [[ZEXT1:%.*]] = zext i8 [[B:%.*]] to i16
+; CHECK-NEXT:    [[ZEXT2:%.*]] = zext i8 [[C:%.*]] to i16
+; CHECK-NEXT:    [[ADD1:%.*]] = add i16 [[ZEXT1]], [[ZEXT2]]
+; CHECK-NEXT:    [[SFT1:%.*]] = and i16 [[ADD1]], 15
+; CHECK-NEXT:    [[SHR1:%.*]] = ashr i16 [[ADD1]], [[SFT1]]
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[ADD1]], [[SHR1]]
+; CHECK-NEXT:    [[SFT2:%.*]] = and i16 [[ADD2]], 7
+; CHECK-NEXT:    [[SHR2:%.*]] = ashr i16 [[ADD2]], [[SFT2]]
+; CHECK-NEXT:    store i16 [[SHR2]], i16* [[A:%.*]], align 2
 ; CHECK-NEXT:    ret void
 ;
   %zext1 = zext i8 %b to i32
   %zext2 = zext i8 %c to i32
   %add1 = add i32 %zext1, %zext2
   %sft1 = and i32 %add1, 15
   %shr1 = ashr i32 %add1, %sft1
   %add2 = add i32 %add1, %shr1
;
   %zext2 = zext i16 %ashr to i32
   %ashr2 = ashr i32 %zext2, 2
   %trunc = trunc i32 %ashr2 to i8
   ret i8 %trunc
 }
 
 define <2 x i16> @ashr_vector(<2 x i8> %x) {
 ; CHECK-LABEL: @ashr_vector(
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[ZA:%.*]] = and <2 x i32> [[Z]], <i32 7, i32 8>
-; CHECK-NEXT:    [[S:%.*]] = ashr <2 x i32> [[Z]], [[ZA]]
-; CHECK-NEXT:    [[A:%.*]] = add <2 x i32> [[S]], [[Z]]
-; CHECK-NEXT:    [[S2:%.*]] = ashr <2 x i32> [[A]], <i32 4, i32 5>
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S2]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[T]]
+; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[ZA:%.*]] = and <2 x i16> [[Z]], <i16 7, i16 8>
+; CHECK-NEXT:    [[S:%.*]] = ashr <2 x i16> [[Z]], [[ZA]]
+; CHECK-NEXT:    [[A:%.*]] = add <2 x i16> [[S]], [[Z]]
+; CHECK-NEXT:    [[S2:%.*]] = ashr <2 x i16> [[A]], <i16 4, i16 5>
+; CHECK-NEXT:    ret <2 x i16> [[S2]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %za = and <2 x i32> %z, <i32 7, i32 8>
   %s = ashr <2 x i32> %z, %za
   %a = add <2 x i32> %s, %z
   %s2 = ashr <2 x i32> %a, <i32 4, i32 5>
   %t = trunc <2 x i32> %s2 to <2 x i16>
   ret <2 x i16> %t
;
   %a = add <2 x i32> %s, %z
   %s2 = ashr <2 x i32> %a, <i32 16, i32 5>
   %t = trunc <2 x i32> %s2 to <2 x i16>
   ret <2 x i16> %t
 }
 
 define i16 @ashr_exact(i16 %x) {
 ; CHECK-LABEL: @ashr_exact(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[ZEXT]], 32767
-; CHECK-NEXT:    [[ASHR:%.*]] = ashr exact i32 [[AND]], 15
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[ASHR]] to i16
-; CHECK-NEXT:    ret i16 [[TRUNC]]
+; CHECK-NEXT:    [[AND:%.*]] = and i16 [[X:%.*]], 32767
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr exact i16 [[AND]], 15
+; CHECK-NEXT:    ret i16 [[ASHR]]
 ;
   %zext = zext i16 %x to i32
   %and = and i32 %zext, 32767
   %ashr = ashr exact i32 %and, 15
   %trunc = trunc i32 %ashr to i16
   ret i16 %trunc
 }
 
;
   %xor = xor i32 -1, %zext
   %lshr2 = ashr i32 %xor, 2
   %trunc = trunc i32 %lshr2 to i16
   ret i16 %trunc
 }
 
 define i16 @ashr_negative_operand_but_short(i16 %x) {
 ; CHECK-LABEL: @ashr_negative_operand_but_short(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[ZEXT]], 32767
-; CHECK-NEXT:    [[XOR:%.*]] = xor i32 -1, [[AND]]
-; CHECK-NEXT:    [[LSHR2:%.*]] = ashr i32 [[XOR]], 2
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[LSHR2]] to i16
-; CHECK-NEXT:    ret i16 [[TRUNC]]
+; CHECK-NEXT:    [[AND:%.*]] = and i16 [[X:%.*]], 32767
+; CHECK-NEXT:    [[XOR:%.*]] = xor i16 -1, [[AND]]
+; CHECK-NEXT:    [[LSHR2:%.*]] = ashr i16 [[XOR]], 2
+; CHECK-NEXT:    ret i16 [[LSHR2]]
 ;
   %zext = zext i16 %x to i32
   %and = and i32 %zext, 32767
   %xor = xor i32 -1, %and
   %lshr2 = ashr i32 %xor, 2
   %trunc = trunc i32 %lshr2 to i16
   ret i16 %trunc
 }