[InstCombine] Teach removeBitcastsFromLoadStoreOnMinMax not to change the size of a store.

[craig.topper]

Summary

We can change the type as long as we don't change the size.

Fixes PR44306

[lebedev.ri]

Sounds good to me.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

558–559

This was even noted here...

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

for (const auto &MDPair : MD) {
     }
   }
 
   return NewStore;
 }
 
 /// Returns true if instruction represent minmax pattern like:
 ///   select ((cmp load V1, load V2), V1, V2).
-static bool isMinMaxWithLoads(Value *V) {
+static bool isMinMaxWithLoads(Value *V, Type *&LoadTy) {
   assert(V->getType()->isPointerTy() && "Expected pointer type.");
   // Ignore possible ty* to ixx* bitcast.
   V = peekThroughBitcast(V);
   // Check that select is select ((cmp load V1, load V2), V1, V2) - minmax
   // pattern.
   CmpInst::Predicate Pred;
   Instruction *L1;
   Instruction *L2;
   Value *LHS;
   Value *RHS;
   if (!match(V, m_Select(m_Cmp(Pred, m_Instruction(L1), m_Instruction(L2)),
                          m_Value(LHS), m_Value(RHS))))
     return false;
+  LoadTy = L1->getType();
   return (match(L1, m_Load(m_Specific(LHS))) &&
           match(L2, m_Load(m_Specific(RHS)))) ||
          (match(L1, m_Load(m_Specific(RHS))) &&
           match(L2, m_Load(m_Specific(LHS))));
 }
 
 /// Combine loads to match the type of their uses' value after looking
 /// through intervening bitcasts.
 ///
 /// The core idea here is that if the result of a load is used in an operation,
 /// we should load the type most conducive to that operation. For example, when
 /// loading an integer and converting that immediately to a pointer, we should
 /// instead directly load a pointer.
 ///
 /// However, this routine must never change the width of a load or the number of
 /// loads as that would introduce a semantic change. This combine is expected to
 /// be a semantic no-op which just allows loads to more closely model the types
 /// of their consuming operations.
 ///
 /// Currently, we also refuse to change the precise type used for an atomic load
 /// or a volatile load. This is debatable, and might be reasonable to change
 /// later. However, it is risky in case some backend or other part of LLVM is
 /// relying on the exact type loaded to select appropriate atomic operations.
 static Instruction *combineLoadToOperationType(InstCombiner &IC, LoadInst &LI) {
static Instruction *combineLoadToOperationType(InstCombiner &IC, LoadInst &LI) {
   const DataLayout &DL = IC.getDataLayout();
 
   // Try to canonicalize loads which are only ever stored to operate over
   // integers instead of any other type. We only do this when the loaded type
   // is sized and has a size exactly the same as its store size and the store
   // size is a legal integer type.
   // Do not perform canonicalization if minmax pattern is found (to avoid
   // infinite loop).
+  Type *Dummy;
   if (!Ty->isIntegerTy() && Ty->isSized() &&
       !(Ty->isVectorTy() && Ty->getVectorIsScalable()) &&
       DL.isLegalInteger(DL.getTypeStoreSizeInBits(Ty)) &&
       DL.typeSizeEqualsStoreSize(Ty) &&
       !DL.isNonIntegralPointerType(Ty) &&
       !isMinMaxWithLoads(
-          peekThroughBitcast(LI.getPointerOperand(), /*OneUseOnly=*/true))) {
+          peekThroughBitcast(LI.getPointerOperand(), /*OneUseOnly=*/true),
+          Dummy)) {
     if (all_of(LI.users(), [&LI](User *U) {
           auto *SI = dyn_cast<StoreInst>(U);
           return SI && SI->getPointerOperand() != &LI &&
                  !SI->getPointerOperand()->isSwiftError();
         })) {
       LoadInst *NewLoad = combineLoadToNewType(
           IC, LI,
           Type::getIntNTy(LI.getContext(), DL.getTypeStoreSizeInBits(Ty)));
if (!match(SI.getPointerOperand(), m_BitCast(m_Value())))
     return false;
   // load? integer?
   Value *LoadAddr;
   if (!match(SI.getValueOperand(), m_Load(m_BitCast(m_Value(LoadAddr)))))
     return false;
   auto *LI = cast<LoadInst>(SI.getValueOperand());
   if (!LI->getType()->isIntegerTy())
     return false;
-  if (!isMinMaxWithLoads(LoadAddr))
+  Type *CmpLoadTy;
+  if (!isMinMaxWithLoads(LoadAddr, CmpLoadTy))
+    return false;
+
+  // Make sure we're not changing the size of the load/store.
+  const auto &DL = IC.getDataLayout();
+  if (DL.getTypeStoreSizeInBits(LI->getType()) !=
+      DL.getTypeStoreSizeInBits(CmpLoadTy))
     return false;
 
   if (!all_of(LI->users(), [LI, LoadAddr](User *U) {
         auto *SI = dyn_cast<StoreInst>(U);
         return SI && SI->getPointerOperand() != LI &&
                peekThroughBitcast(SI->getPointerOperand()) != LoadAddr &&
                !SI->getPointerOperand()->isSwiftError();
       }))
     return false;
 
   IC.Builder.SetInsertPoint(LI);
-  LoadInst *NewLI = combineLoadToNewType(
-      IC, *LI, LoadAddr->getType()->getPointerElementType());
+  LoadInst *NewLI = combineLoadToNewType(IC, *LI, CmpLoadTy);
   // Replace all the stores with stores of the newly loaded value.
   for (auto *UI : LI->users()) {
     auto *USI = cast<StoreInst>(UI);
     IC.Builder.SetInsertPoint(USI);
     combineStoreToNewValue(IC, *USI, NewLI);
   }
   IC.replaceInstUsesWith(*LI, UndefValue::get(LI->getType()));
   IC.eraseInstFromFunction(*LI);

llvm/test/Transforms/InstCombine/PR37526.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine -S < %s | FileCheck %s
 
 define void @PR37526(i32* %pz, i32* %px, i32* %py) {
 ; CHECK-LABEL: @PR37526(
+; CHECK-NEXT:    [[T1:%.*]] = bitcast i32* [[PZ:%.*]] to i64*
 ; CHECK-NEXT:    [[T2:%.*]] = load i32, i32* [[PY:%.*]], align 4
 ; CHECK-NEXT:    [[T3:%.*]] = load i32, i32* [[PX:%.*]], align 4
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[T2]], [[T3]]
-; CHECK-NEXT:    [[R1:%.*]] = select i1 [[CMP]], i32 [[T3]], i32 [[T2]]
-; CHECK-NEXT:    store i32 [[R1]], i32* [[PZ:%.*]], align 4
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i32* [[PX]], i32* [[PY]]
+; CHECK-NEXT:    [[BC:%.*]] = bitcast i32* [[SELECT]] to i64*
+; CHECK-NEXT:    [[R:%.*]] = load i64, i64* [[BC]], align 4
+; CHECK-NEXT:    store i64 [[R]], i64* [[T1]], align 4
 ; CHECK-NEXT:    ret void
 ;
   %t1 = bitcast i32* %pz to i64*
   %t2 = load i32, i32* %py
   %t3 = load i32, i32* %px
   %cmp = icmp slt i32 %t2, %t3
   %select = select i1 %cmp, i32* %px, i32* %py
   %bc = bitcast i32* %select to i64*
   %r = load i64, i64* %bc
   store i64 %r, i64* %t1
   ret void
 }