[CodeGenPrepare] Avoid out-of-bounds shift

[bjope]

Summary

AddressingModeMatcher::matchOperationAddr may attempt to shift a
variable by the same amount of steps as found in the IR in a SHL
instruction. This was done without considering that there could be
undefined behavior in the IR, so the shift performed when compiling
could end up having undefined behavior as well.

This patch avoid UB in the codegenprepare by making sure that we
limit the shift amount used, in a similar way as already being done
in CodeGenPrepare::optimizeLoadExt.

[bjope]

Fix typo (there was an extra "Scale =" that wasn't supposed to be there).

[spatel]

This seems fine, but check my understanding: there's no difference in the test output with this fix if opt is not built with ubsan enabled?

Could/should we just exit if we detect malformed IR?

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index c888adeafca5..294a37689013 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -4551,9 +4551,12 @@ bool AddressingModeMatcher::matchOperationAddr(User *AddrInst, unsigned Opcode,
     if (!RHS || RHS->getBitWidth() > 64)
       return false;
     int64_t Scale = RHS->getSExtValue();
-    if (Opcode == Instruction::Shl)
+    if (Opcode == Instruction::Shl) {
+      // Bail out if the IR is not well-defined (overshift is poison).
+      if (RHS->getZExtValue() > 63)
+        return false;
       Scale = 1LL << Scale;
-
+    }
     return matchScaledValue(AddrInst->getOperand(0), Scale, Depth);
   }
   case Instruction::GetElementPtr: {

[bjope]

In D118602#3293795, @spatel wrote:

This seems fine, but check my understanding: there's no difference in the test output with this fix if opt is not built with ubsan enabled?

Yes. Unfortunately the test case (as-is) doen't show any difference. I'm not sure exactly how to modify the test case to achive that.

Could/should we just exit if we detect malformed IR?

Yes, that is another option. I kind of copied how this was dealt with on line 6491 (another place where CGP is analyzing Shl), but bailing out when finding an OOB shift amount seems safe here. Is that preferred?

[spatel]

In D118602#3293833, @bjope wrote:

Yes, that is another option. I kind of copied how this was dealt with on line 6491 (another place where CGP is analyzing Shl), but bailing out when finding an OOB shift amount seems safe here. Is that preferred?

I don't think it makes a difference either way. Since there's precedence for this code pattern, LGTM.

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool AddressingModeMatcher::matchOperationAddr(User *AddrInst, unsigned Opcode,
   //break;
   case Instruction::Mul:
   case Instruction::Shl: {
     // Can only handle X*C and X << C.
     AddrMode.InBounds = false;
     ConstantInt *RHS = dyn_cast<ConstantInt>(AddrInst->getOperand(1));
     if (!RHS || RHS->getBitWidth() > 64)
       return false;
-    int64_t Scale = RHS->getSExtValue();
-    if (Opcode == Instruction::Shl)
-      Scale = 1LL << Scale;
+    int64_t Scale = Opcode == Instruction::Shl
+                        ? 1LL << RHS->getLimitedValue(RHS->getBitWidth() - 1)
+                        : RHS->getSExtValue();
 
     return matchScaledValue(AddrInst->getOperand(0), Scale, Depth);
   }
   case Instruction::GetElementPtr: {
     // Scan the GEP.  We check it if it contains constant offsets and at most
     // one variable offset.
     int VariableOperand = -1;
     unsigned VariableScale = 0;

llvm/test/CodeGen/X86/codegen-prepare-oob-shl.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -mtriple i686-unknown-unknown -codegenprepare -S | FileCheck %s
+
+target datalayout = "e-p:8:8"
+
+; The shl has UB (shift count oob). This used to result in undefined behavior
+; in codegenprepare when AddressingModeMatcher::matchOperationAddr tried to
+; shift a variable by that amount during compilation. Intent with the test
+; case is to verify that this compiles without complaints even if opt is built
+; with ubsan enabled.
+define dso_local void @main(i32 %a, [3 x { i8, i8 }*]* %p) {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], -1229216766
+; CHECK-NEXT:    [[ARRAYIDX926:%.*]] = getelementptr inbounds [3 x { i8, i8 }*], [3 x { i8, i8 }*]* [[P:%.*]], i32 0, i32 [[SHL]]
+; CHECK-NEXT:    [[L0:%.*]] = load { i8, i8 }*, { i8, i8 }** [[ARRAYIDX926]], align 1
+; CHECK-NEXT:    ret void
+;
+  %shl = shl i32 %a, -1229216766
+  %arrayidx926 = getelementptr inbounds [3 x { i8, i8 }*], [3 x { i8, i8 }*]* %p, i32 0, i32 %shl
+  %l0 = load { i8, i8 }*, { i8, i8 }** %arrayidx926, align 1
+  ret void
+}