1 files changed, 54 insertions, 6 deletions

diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index 71fa795ec294a..6e24f03c4cfd6 100644
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -792,13 +792,36 @@ Constant *llvm::ConstantFoldExtractElementInstruction(Constant *Val,
   if (isa<UndefValue>(Val) || isa<UndefValue>(Idx))
     return UndefValue::get(Val->getType()->getVectorElementType());
 
-  if (ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx)) {
-    // ee({w,x,y,z}, wrong_value) -> undef
-    if (CIdx->uge(Val->getType()->getVectorNumElements()))
-      return UndefValue::get(Val->getType()->getVectorElementType());
-    return Val->getAggregateElement(CIdx->getZExtValue());
+  auto *CIdx = dyn_cast<ConstantInt>(Idx);
+  if (!CIdx)
+    return nullptr;
+
+  // ee({w,x,y,z}, wrong_value) -> undef
+  if (CIdx->uge(Val->getType()->getVectorNumElements()))
+    return UndefValue::get(Val->getType()->getVectorElementType());
+
+  // ee (gep (ptr, idx0, ...), idx) -> gep (ee (ptr, idx), ee (idx0, idx), ...)
+  if (auto *CE = dyn_cast<ConstantExpr>(Val)) {
+    if (CE->getOpcode() == Instruction::GetElementPtr) {
+      SmallVector<Constant *, 8> Ops;
+      Ops.reserve(CE->getNumOperands());
+      for (unsigned i = 0, e = CE->getNumOperands(); i != e; ++i) {
+        Constant *Op = CE->getOperand(i);
+        if (Op->getType()->isVectorTy()) {
+          Constant *ScalarOp = ConstantExpr::getExtractElement(Op, Idx);
+          if (!ScalarOp)
+            return  nullptr;
+          Ops.push_back(ScalarOp);
+        } else
+          Ops.push_back(Op);
+      }
+      return CE->getWithOperands(Ops, CE->getType()->getVectorElementType(),
+                                 false,
+                                 Ops[0]->getType()->getPointerElementType());
+    }
   }
-  return nullptr;
+
+  return Val->getAggregateElement(CIdx);
 }
 
 Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
@@ -810,6 +833,12 @@ Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
   ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx);
   if (!CIdx) return nullptr;
 
+  // Do not iterate on scalable vector. The num of elements is unknown at
+  // compile-time.
+  VectorType *ValTy = cast<VectorType>(Val->getType());
+  if (ValTy->isScalable())
+    return nullptr;
+
   unsigned NumElts = Val->getType()->getVectorNumElements();
   if (CIdx->uge(NumElts))
     return UndefValue::get(Val->getType());
@@ -844,6 +873,12 @@ Constant *llvm::ConstantFoldShuffleVectorInstruction(Constant *V1,
   // Don't break the bitcode reader hack.
   if (isa<ConstantExpr>(Mask)) return nullptr;
 
+  // Do not iterate on scalable vector. The num of elements is unknown at
+  // compile-time.
+  VectorType *ValTy = cast<VectorType>(V1->getType());
+  if (ValTy->isScalable())
+    return nullptr;
+
   unsigned SrcNumElts = V1->getType()->getVectorNumElements();
 
   // Loop over the shuffle mask, evaluating each element.
@@ -965,6 +1000,19 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
                                               Constant *C2) {
   assert(Instruction::isBinaryOp(Opcode) && "Non-binary instruction detected");
 
+  // Simplify BinOps with their identity values first. They are no-ops and we
+  // can always return the other value, including undef or poison values.
+  // FIXME: remove unnecessary duplicated identity patterns below.
+  // FIXME: Use AllowRHSConstant with getBinOpIdentity to handle additional ops,
+  //        like X << 0 = X.
+  Constant *Identity = ConstantExpr::getBinOpIdentity(Opcode, C1->getType());
+  if (Identity) {
+    if (C1 == Identity)
+      return C2;
+    if (C2 == Identity)
+      return C1;
+  }
+
   // Handle scalar UndefValue. Vectors are always evaluated per element.
   bool HasScalarUndef = !C1->getType()->isVectorTy() &&
                         (isa<UndefValue>(C1) || isa<UndefValue>(C2));