[InstCombine] Fix matchRotate bug when one operand is a ConstantExpr shift

[AndrewScheidecker]

Summary

This bug seems to be harmless in release builds, but will cause an error in UBSAN builds or an assertion failure in debug builds.

When it gets to this opcode comparison, it assumes both of the operands are BinaryOperators, but the prior m_LogicalShift will also match a ConstantExpr. The cast<BinaryOperator> will assert in a debug build, or reading an invalid value for BinaryOp from memory with ((BinaryOperator*)constantExpr)->getOpcode() will cause an error in a UBSAN build.

The test I added will fail without this change in debug/UBSAN builds, but not in release.

[spatel]

Not sure if we can even succeed to form a funnel shift with a constant expression here? It's safer to give up immediately if we don't have binops:

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index 7c195daa3e7..aaa883a7037 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -1819,14 +1819,18 @@ static Instruction *matchRotate(Instruction &Or) {
 
   // First, find an or'd pair of opposite shifts with the same shifted operand:
   // or (lshr ShVal, ShAmt0), (shl ShVal, ShAmt1)
-  Value *Or0 = Or.getOperand(0), *Or1 = Or.getOperand(1);
+  BinaryOperator *Or0, *Or1;
+  if (!match(Or.getOperand(0), m_BinOp(Or0)) ||
+      !match(Or.getOperand(1), m_BinOp(Or1)))
+    return nullptr;
+
   Value *ShVal, *ShAmt0, *ShAmt1;
   if (!match(Or0, m_OneUse(m_LogicalShift(m_Value(ShVal), m_Value(ShAmt0)))) ||
       !match(Or1, m_OneUse(m_LogicalShift(m_Specific(ShVal), m_Value(ShAmt1)))))
     return nullptr;
 
-  auto ShiftOpcode0 = cast<BinaryOperator>(Or0)->getOpcode();
-  auto ShiftOpcode1 = cast<BinaryOperator>(Or1)->getOpcode();
+  BinaryOperator::BinaryOps ShiftOpcode0 = Or0->getOpcode();
+  BinaryOperator::BinaryOps ShiftOpcode1 = Or1->getOpcode();
   if (ShiftOpcode0 == ShiftOpcode1)
     return nullptr;

[AndrewScheidecker]

Good point that we can discard the ConstantExpr case entirely.

[spatel]

LGTM - let me know if I should commit on your behalf.

[AndrewScheidecker]

LGTM - let me know if I should commit on your behalf.

Yes, please do.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

static Instruction *matchRotate(Instruction &Or) {
   // First, find an or'd pair of opposite shifts with the same shifted operand:
   // or (lshr ShVal, ShAmt0), (shl ShVal, ShAmt1)
   Value *Or0 = Or.getOperand(0), *Or1 = Or.getOperand(1);
   Value *ShVal, *ShAmt0, *ShAmt1;
   if (!match(Or0, m_OneUse(m_LogicalShift(m_Value(ShVal), m_Value(ShAmt0)))) ||
       !match(Or1, m_OneUse(m_LogicalShift(m_Specific(ShVal), m_Value(ShAmt1)))))
     return nullptr;
 
-  auto ShiftOpcode0 = cast<BinaryOperator>(Or0)->getOpcode();
-  auto ShiftOpcode1 = cast<BinaryOperator>(Or1)->getOpcode();
+  auto getBinaryOperatorOrConstantExprOpcode = [](Value *V) -> unsigned {
+    if (auto TheBinaryOperator = dyn_cast<BinaryOperator>(V))
+      return unsigned(TheBinaryOperator->getOpcode());
+    else if (auto TheConstExpr = dyn_cast<ConstantExpr>(V))
+      return TheConstExpr->getOpcode();
+    else
+      llvm_unreachable(
+          "Expected Value to be either a BinaryOperator or ConstantExpr");
+  };
+
+  auto ShiftOpcode0 = getBinaryOperatorOrConstantExprOpcode(Or0);
+  auto ShiftOpcode1 = getBinaryOperatorOrConstantExprOpcode(Or1);
   if (ShiftOpcode0 == ShiftOpcode1)
     return nullptr;
 
   // Match the shift amount operands for a rotate pattern. This always matches
   // a subtraction on the R operand.
   auto matchShiftAmount = [](Value *L, Value *R, unsigned Width) -> Value * {
     // The shift amount may be masked with negation:
     // (shl ShVal, (X & (Width - 1))) | (lshr ShVal, ((-X) & (Width - 1)))

test/Transforms/InstCombine/rotate.ll

;
   %sub = sub i24 24, %shamt
   %shr = lshr i24 %x, %sub
   %shl = shl i24 %x, %shamt
   %or = or i24 %shl, %shr
   %r = select i1 %cmp, i24 %x, i24 %or
   ret i24 %r
 }
 
+; Test that the transform doesn't crash when there's an "or" with a ConstantExpr operand.
+
+@external_global = external global i8
+
+define i32 @rotl_constant_expr(i32 %shamt) {
+; CHECK-LABEL: @rotl_constant_expr(
+; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 ptrtoint (i8* @external_global to i32), [[SHAMT:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = or i32 [[SHR]], shl (i32 ptrtoint (i8* @external_global to i32), i32 11)
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %shr = lshr i32 ptrtoint (i8* @external_global to i32), %shamt
+  %r = or i32 %shr, shl (i32 ptrtoint (i8* @external_global to i32), i32 11)
+  ret i32 %r
+}
 No newline at end of file