[InstCombine] Defend against worst-case exponential execution time

[jmolloy]

Summary

It is possible to create an IR pattern that results in exponential execution
time in CollectBitParts, by making a tree where every OR refers to one other OR via *both*
operands:

for (i = 0; i < N; ++i)
    b |= b >> 1;  // Causes 2^N executions of CollectBitParts!

To defend against this (and because ORs are the only node that cause us to
fork control flow) we keep a counter alive between all invocations of
CollectBitParts, expected to be initialized to bitwidth(b).

[majnemer]

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1775–1777

This seems a bit much if you have i128. What is NumOrsRemaining typically when this optimization kicks in?

[majnemer]

Out of curiosity, is there a major benefit to canonicalizing this in the middle-end vs doing this transform in code-gen prepare?

[jmolloy]

Hi David, Joerg,

This version rewrites the algorithm significantly. In the original algorithm every recursive call was a function of a Value* and multiple other operands. This made memoization impossible as it is not often that the collectBitParts is called with the same Value* and also the same BitMask/OverallLeftShift.

collectBitParts is now just a function of a single Value*, and it is up to the caller to apply any transformations on its result. This means the result can be memoized which means that Joerg's testcase requires ~64 recursions instead of 2^32. The new design makes the code smaller and more concise too.

David, I have implemented this in InstCombine still. I also have a version of this patch implemented in CodeGenPrepare. The CGP patch cannot fold bitreverse(bitreverse(X)) -> X, but perhaps that pattern is less common than the equivalent with bswap.

I am happy to switch to the CGP version if you wish, which would mean I'd revert all of my patches to InstCombineAndOrXor.cpp.

Cheers,

James

[jmolloy]

Hi Joerg, David,

I think this version should tick all the boxes. There are several requirements that ended up with this design;

1. Matching bitreversals is too heavyweight for InstCombine and doesn't really need to be done so early.
2. Bitreversals and byteswaps are very related in their matching logic.
3. We want to implement support for matching more advanced bswap/bitreverse patterns like partial bswaps/bitreverses.
4. Bswaps are best matched early in InstCombine.

The result of these is that a new utility function is created in Transforms/Utils/Local.h that can be configured to search for bswaps, bitreverses or both. InstCombine uses it to find only bswaps, CGP uses it to find only bitreversals.

We can then extend the matching logic in one place only.

[jmolloy]

Hi Joerg,

Thanks for that bug report. It saved me some embarassment when the asan bot would have picked it up post-commit.

The issue was holding a reference into a DenseMap which can be invalidated during a hashtable resize. This is solved (as in other places in LLVM) by explicitly using a std::map.

James

[majnemer]

lib/CodeGen/CodeGenPrepare.cpp

5241–5250 ↗

(On Diff #44450)

Should we bother doing this on targets which can't lower biterverse to anything good?

[jmolloy]

Hi David,

This makes sense. I've updated the diff.

James

[jmolloy]

Hi David,

Does this look OK now? I'm keen to get this committed as it needs to be merged to 3.8 to fix a compiler hang.

Cheers,

James

[majnemer]

LGTM

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

 /// 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.
 ///
+/// It is possible to create an IR pattern that results in exponential execution
+/// time here, by making a tree where every OR refers to one other OR via *both*
+/// operands:
+///
+///   for (i = 0; i < N; ++i)
+///     b |= b >> 1;  // Causes 2^N executions of CollectBitParts!
+///
+/// To defend against this (and because ORs are the only node that cause us to
+/// fork control flow) we keep a counter alive between all invocations of
+/// CollectBitParts, expected to be initialized to bitwidth(b).
 static bool CollectBitParts(Value *V, int OverallLeftShift, APInt BitMask,
                             SmallVectorImpl<Value *> &BitValues,
                             SmallVectorImpl<int> &BitProvenance,
-                            SmallPtrSetImpl<Value*> &OrValues) {
+                            SmallPtrSetImpl<Value*> &OrValues,
+                            unsigned &NumOrsRemaining) {
   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.
if (isa<IntrinsicInst>(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)
+    if (I->getOpcode() == Instruction::Or) {
+      if (NumOrsRemaining-- == 0)
+        return true;
       return CollectBitParts(I->getOperand(0), OverallLeftShift, BitMask,
-                             BitValues, BitProvenance, OrValues) ||
+                             BitValues, BitProvenance, OrValues,
+                             NumOrsRemaining) ||
              CollectBitParts(I->getOperand(1), OverallLeftShift, BitMask,
-                             BitValues, BitProvenance, OrValues);
+                             BitValues, BitProvenance, OrValues,
+                             NumOrsRemaining);
+    }
 
     // If this is a logical shift by a constant, recurse with OverallLeftShift
     // and BitMask adjusted.
     if (I->isLogicalShift() && isa<ConstantInt>(I->getOperand(1))) {
       unsigned ShAmt =
           cast<ConstantInt>(I->getOperand(1))->getLimitedValue(~0U);
       // Ensure the shift amount is defined.
       if (ShAmt > BitValues.size())
if (I->isLogicalShift() && isa<ConstantInt>(I->getOperand(1))) {
       }
 
       if (OverallLeftShift >= (int)BitValues.size())
         return true;
       if (OverallLeftShift <= -(int)BitValues.size())
         return true;
 
       return CollectBitParts(I->getOperand(0), OverallLeftShift, BitMask,
-                             BitValues, BitProvenance, OrValues);
+                             BitValues, BitProvenance, OrValues,
+                             NumOrsRemaining);
     }
 
     // If this is a logical 'and' with a mask that clears bits, clear the
     // corresponding bits in BitMask.
     if (I->getOpcode() == Instruction::And &&
         isa<ConstantInt>(I->getOperand(1))) {
       unsigned NumBits = BitValues.size();
       APInt Bit(I->getType()->getPrimitiveSizeInBits(), 1);
if (I->getOpcode() == Instruction::And &&
           BitMask.clearBit(i);
           continue;
         }
 
         // Otherwise, this bit is kept.
       }
 
       return CollectBitParts(I->getOperand(0), OverallLeftShift, BitMask,
-                             BitValues, BitProvenance, OrValues);
+                             BitValues, BitProvenance, OrValues,
+                             NumOrsRemaining);
     }
   }
 
   // Okay, we got to something that isn't a shift, 'or' or 'and'.  This must be
   // the input value to the bswap/bitreverse. To be part of a bswap or
   // bitreverse we must be demanding a contiguous range of bits from it.
   unsigned InputBitLen = BitMask.countPopulation();
   unsigned InputBitNo = BitMask.countTrailingZeros();
Instruction *InstCombiner::MatchBSwapOrBitReverse(BinaryOperator &I) {
   /// We keep track of which bit (BitProvenance) inside which value (BitValues)
   /// defines each bit in the result.
   SmallVector<Value *, 8> BitValues(BW, nullptr);
   SmallVector<int, 8> BitProvenance(BW, -1);
   
   // Try to find all the pieces corresponding to the bswap.
   APInt BitMask = APInt::getAllOnesValue(BitValues.size());
   SmallPtrSet<Value*, 2> OrValues;
-  if (CollectBitParts(&I, 0, BitMask, BitValues, BitProvenance, OrValues))
+  unsigned NumOrsRemaining = BitMask.getBitWidth();
+  if (CollectBitParts(&I, 0, BitMask, BitValues, BitProvenance, OrValues,
+                      NumOrsRemaining))
     return nullptr;
 
   // Check to see if all of the bits come from the same value.
   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)

test/Transforms/InstCombine/bitreverse-hang.ll

+; RUN: opt < %s -loop-unroll -instcombine -S | FileCheck %s
+
+; This test is a worst-case scenario for bitreversal/byteswap detection.
+; After loop unrolling (the unrolled loop is unreadably large so it has been kept
+; rolled here), we have a binary tree of OR operands (as bitreversal detection
+; looks straight through shifts):
+;
+;  OR
+;  | \
+;  |  LSHR
+;  | /
+;  OR
+;  | \
+;  |  LSHR
+;  | /
+;  OR
+;
+; This results in exponential runtime. The loop here is 32 iterations which will
+; totally hang if we don't deal with this case cleverly.
+
+@b = common global i32 0, align 4
+
+; CHECK: define i32 @fn1
+define i32 @fn1() #0 {
+entry:
+  %b.promoted = load i32, i32* @b, align 4, !tbaa !2
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %or4 = phi i32 [ %b.promoted, %entry ], [ %or, %for.body ]
+  %i.03 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %shr = lshr i32 %or4, 1
+  %or = or i32 %shr, %or4
+  %inc = add nuw nsw i32 %i.03, 1
+  %exitcond = icmp eq i32 %inc, 32
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  store i32 %or, i32* @b, align 4, !tbaa !2
+  ret i32 undef
+}
+
+attributes #0 = { norecurse nounwind ssp uwtable "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="core2" "target-features"="+cx16,+fxsr,+mmx,+sse,+sse2,+sse3,+ssse3" "unsafe-fp-math"="false" "use-soft-float"="false" }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"PIC Level", i32 2}
+!1 = !{!"clang version 3.8.0 (http://llvm.org/git/clang.git eb70f4e9cc9a4dc3dd57b032fb858d56b4b64a0e)"}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"int", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}