[InstCombine] Identify partial bitreverses

[jmolloy]

Summary

It could be beneficial to emit a bitreverse intrinsic even when some of the result of that intrinsic need to be masked out. For example, assume we have a bitreverse but just without some bits:

int4_t b = 0;
if (a & 1) b |= 8;
if (a & 2) b |= 4;
// (a & 4) -> (b |= 2) is missing
if (a & 8) b |= 1;

Here, we can simply perform the bitreverse and mask: (res &= ~2).

This patch provides support for this. Obviously, doing this unconditionally would not be a good idea. A single bit move that happens to be correct for a bitreversal will cause a bitreverse (which may not be cheap on the target) and mask to be emitted. Therefore in this patch we heuristically say that if a bitreverse would require more than half its bits masked out not to bother. Perhaps this should be a target hook, depending on whether bitreverse is cheap or not?

[jmolloy]

Optimistic Christmas Season Ping!

[sanjoy]

I have a couple of (possibly naive) questions.

I'm not familiar with this code, so while I can help review it, I think the final LGTM should come from someone else more familiar with this code.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1584

Should this be "false if not"?

1765

I don't understand the logic here. If BitProvenance[i] == i then I suppose that means the ith bit of the output is the ith bit of V. How does that bit end up in the instruction you return?

1782

Will teaching CollectBitParts about the bswap and bitreverse intrinsics have the same effect. If so, that sounds slightly better.

1790–1799

Might want to use Builder->CreateCall here, so that the call instruction gets enqueued in the worklist.

[jmolloy]

Hi Sanjoy,

Thanks for the review! New version coming up.

James

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1584

I honestly don't know. The problem is that would be a double negative ("if not unsuccessful").

The best thing to do here is for me to rewrite this function to be sane and return true on success and false on failure. I'll do this in a dedicated followup.

1765

In a partial bitreverse/bswap, some of the bits are going to be the same as the original value. We're essentially saying that the result is a OR-ed mask of a reversed value and the original value:

uint8_t a;
uint8_t b = bitreverse(a);
uint8_t c = a & M | b & ~M

This can come from something like:

uint8_t dest = src; // In a pure bitreverse, dest would be initialized to zero.
if (src & 1<<0) dest |= 1<<7;
if (src & 1<<1) dest |= 1<<6;
...

1782

Indeed, that sounds much better. My new patch does this.

[sanjoy]

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1765

In a partial bitreverse/bswap, some of the bits are going to be the same as the original value. We're essentially saying that the result is a OR-ed mask of a reversed value and the original value:

Right, and I don't see where you're creating that OR -- all I see is the bitreverse(a) & ~M part. If I understand the scheme correctly and am not misreading the code, you'll have to keep two masks -- IgnoreMask (set if BitProvenance[i] == -1) and IdentityMask (set if BitProvenance[i] == i), and then return bitreverse(a) & ~IgnoreMask | a & IdentityMask.

(I assume BitProvenance[i] == -1 means "bit i is always 0 in the output"?)

[sanjoy]

Getting this out of my queue for now. If my last comment was incorrect (and I'm misreading the code or something), please let me know.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

 /// validate that all `value`s are identical and that the bitnumber to bitnumber
 /// mapping is correct for a bswap or bitreverse.
 ///
 /// For example, if the current subexpression if "(shl i32 %X, 24)" then we know
 /// that the expression deposits the low byte of %X into the high byte of the
 /// result and that all other bits are zero. This expression is accepted,
 /// BitValues[24-31] are set to %X and BitProvenance[24-31] are set to [0-7].
 ///
 /// This function returns true if the match was unsuccessful and false if so.
 /// On entry to the function the "OverallLeftShift" is a signed integer value
 /// indicating the number of bits that the subexpression is later shifted.  For
 /// example, if the expression is later right shifted by 16 bits, the
 /// OverallLeftShift value would be -16 on entry.  This is used to specify which
 /// bits of BitValues are actually being set.
 ///
 /// Similarly, BitMask is a bitmask where a bit is clear if its corresponding
 /// bit is masked to zero by a user.  For example, in (X & 255), X will be
 /// processed with a bytemask of 255. BitMask is always in the local
 /// (OverallLeftShift) coordinate space.
 ///
 static bool CollectBitParts(Value *V, int OverallLeftShift, APInt BitMask,
                             SmallVectorImpl<Value *> &BitValues,
                             SmallVectorImpl<int> &BitProvenance) {
   if (Instruction *I = dyn_cast<Instruction>(V)) {
+    // If this is a bitreverse intrinsic, it can obviously be part of a
+    // bitreverse/bswap.
+    // FIXME: Do the same for Intrinsic::bswap!
+    if (isa<IntrinsicInst>(I) &&
+        cast<IntrinsicInst>(I)->getIntrinsicID() == Intrinsic::bitreverse) {
+      auto *V = I->getOperand(0);
+      // We don't have to demand a contiguous range of bits from a bitreverse.
+      // We know all bits will be in the correct order.
+      auto BW = BitMask.getBitWidth();
+      unsigned DestBitNo = OverallLeftShift;
+      for (unsigned I = 0; I < BW - OverallLeftShift; ++I)
+        if (BitMask[I] && BitValues[DestBitNo + I] &&
+            BitValues[DestBitNo + I] != V)
+          return true;
+
+      for (unsigned I = 0; I < BitMask.getBitWidth() - OverallLeftShift; ++I)
+        if (BitMask[I]) {
+          BitProvenance[DestBitNo + I] = BW - I - 1;
+          BitValues[DestBitNo + I] = V;
+        }
+      return false;
+    }
+    
     // If this is an or instruction, it may be an inner node of the bswap.
     if (I->getOpcode() == Instruction::Or)
       return CollectBitParts(I->getOperand(0), OverallLeftShift, BitMask,
                              BitValues, BitProvenance) ||
              CollectBitParts(I->getOperand(1), OverallLeftShift, BitMask,
                              BitValues, BitProvenance);
 
     // If this is a logical shift by a constant, recurse with OverallLeftShift
Instruction *InstCombiner::MatchBSwapOrBitReverse(BinaryOperator &I) {
   SmallVector<int, 8> BitProvenance(BW, -1);
   
   // Try to find all the pieces corresponding to the bswap.
   APInt BitMask = APInt::getAllOnesValue(BitValues.size());
   if (CollectBitParts(&I, 0, BitMask, BitValues, BitProvenance))
     return nullptr;
 
   // Check to see if all of the bits come from the same value.
-  Value *V = BitValues[0];
-  if (!V) return nullptr;  // Didn't find a bit?  Must be zero.
-
-  if (!std::all_of(BitValues.begin(), BitValues.end(),
-                   [&](const Value *X) { return X == V; }))
+  Value *V = nullptr;
+  for (unsigned I = 0, E = BitValues.size(); I != E; ++I) {
+    if (BitValues[I] && !V)
+      V = BitValues[I];
+    else if (BitProvenance[I] != -1 && BitValues[I] != V)
+      return nullptr;
+  }
+  if (!V)
     return nullptr;
   
   // Now, is the bit permutation correct for a bswap or a bitreverse? We can
   // only byteswap values with an even number of bytes.
-  bool OKForBSwap = BW % 16 == 0, OKForBitReverse = true;;
-  for (unsigned i = 0, e = BitValues.size(); i != e; ++i) {
+  bool OKForBSwap = BW % 16 == 0, OKForBitReverse = true;
+  APInt IgnoreMask(BW, 0);
+  for (int i = 0, e = BitValues.size(); i != e; ++i) {
+    if (BitProvenance[i] == -1 || BitProvenance[i] == i) {
+      IgnoreMask.setBit(i);
+      continue;
+    }
     OKForBSwap &= bitTransformIsCorrectForBSwap(BitProvenance[i], i, BW);
     OKForBitReverse &=
         bitTransformIsCorrectForBitReverse(BitProvenance[i], i, BW);
   }
 
   Intrinsic::ID Intrin;
   if (OKForBSwap)
     Intrin = Intrinsic::bswap;
   else if (OKForBitReverse)
     Intrin = Intrinsic::bitreverse;
   else
     return nullptr;
-  
+  // We know we can create a bswap or bitreverse out of this. But is it worth
+  // it? Heuristically let's say it's worth it if we're setting at least half
+  // of the bits of the resulting value from the bitreverse/bswap.
+  if (IgnoreMask.countPopulation() > BW / 2)
+    return nullptr;
+  
   Module *M = I.getParent()->getParent()->getParent();
   Function *F = Intrinsic::getDeclaration(M, Intrin, ITy);
-  return CallInst::Create(F, V);
+  auto *CI = CallInst::Create(F, V);
+
+  if (!IgnoreMask.getBoolValue())
+    return CI;
+
+  auto *And = BinaryOperator::Create(Instruction::And, CI,
+                                     ConstantInt::get(ITy, ~IgnoreMask));
+  Worklist.Add(CI);
+  CI->insertBefore(&I);
+  return And;
 }
 
 /// We have an expression of the form (A&C)|(B&D).  Check if A is (cond?-1:0)
 /// and either B or D is ~(cond?-1,0) or (cond?0,-1), then we can simplify this
 /// expression to "cond ? C : D or B".
 static Instruction *MatchSelectFromAndOr(Value *A, Value *B,
                                          Value *C, Value *D) {
   // If A is not a select of -1/0, this cannot match.

test/Transforms/InstCombine/bitreverse-partial.ll

+; RUN: opt -instcombine -S < %s | FileCheck %s
+
+; CHECK-LABEL: @partial_i4
+; CHECK: %[[x:.*]] = call i4 @llvm.bitreverse.i4(i4 %a)
+; CHECK: %[[y:.*]] = and i4 %[[x]], -2
+; CHECK: ret i4 %[[y]]
+define i4 @partial_i4(i4 %a) #0 {
+entry:
+  %and = shl i4 %a, 3
+  %0 = and i4 %and, 8
+  %and1 = shl i4 %a, 1
+  %1 = and i4 %and1, 4
+  %2 = or i4 %1, %0
+  %and6 = lshr i4 %a, 1
+  %3 = and i4 %and6, 2
+  %b.1 = or i4 %2, %3
+  ret i4 %b.1
+}
+
+; CHECK-LABEL: @partial_i4_or_2
+; CHECK: %[[x:.*]] = call i4 @llvm.bitreverse.i4(i4 %a)
+; CHECK: %[[y:.*]] = and i4 %[[x]], -4
+; CHECK: %[[z:.*]] = or i4 %[[y]], 2
+; CHECK: ret i4 %[[z]]
+define i4 @partial_i4_or_2(i4 %a) #0 {
+entry:
+  %and = shl i4 %a, 3
+  %0 = and i4 %and, 8
+  %and1 = shl i4 %a, 1
+  %1 = and i4 %and1, 4
+  %2 = or i4 %1, %0
+  %and6 = lshr i4 %a, 1
+  %3 = and i4 %and6, 2
+  %b.1 = or i4 %2, %3
+  %z = or i4 %b.1, 2
+  ret i4 %z
+}