[GlobalOpt] Fix the load types when OptimizeGlobalAddressOfMalloc

Currently, in OptimizeGlobalAddressOfMalloc, the transformation for global loads assumes that they have the same Type. With the support of ConstantExpr (https://reviews.llvm.org/D106589), this may not be true any more (as seen in the test case), and we miss the code to handle this, This is to fix that.

Reviewed By: efriedma

Differential Revision: https://reviews.llvm.org/D107397

Author: scui-ibm

llvm/lib/Transforms/IPO/GlobalOpt.cpp

@@ -967,9 +967,8 @@ OptimizeGlobalAddressOfMalloc(GlobalVariable *GV, CallInst *CI, Type *AllocTy,
     }
   }
 
-  Constant *RepValue = NewGV;
-  if (NewGV->getType() != GV->getValueType())
-    RepValue = ConstantExpr::getBitCast(RepValue, GV->getValueType());
+  SmallPtrSet<Constant *, 1> RepValues;
+  RepValues.insert(NewGV);
 
   // If there is a comparison against null, we will insert a global bool to
   // keep track of whether the global was initialized yet or not.
@@ -1001,7 +1000,9 @@ OptimizeGlobalAddressOfMalloc(GlobalVariable *GV, CallInst *CI, Type *AllocTy,
       Use &LoadUse = *LI->use_begin();
       ICmpInst *ICI = dyn_cast<ICmpInst>(LoadUse.getUser());
       if (!ICI) {
-        LoadUse = RepValue;
+        auto *CE = ConstantExpr::getBitCast(NewGV, LI->getType());
+        RepValues.insert(CE);
+        LoadUse.set(CE);
         continue;
       }
 
@@ -1047,9 +1048,8 @@ OptimizeGlobalAddressOfMalloc(GlobalVariable *GV, CallInst *CI, Type *AllocTy,
   // To further other optimizations, loop over all users of NewGV and try to
   // constant prop them.  This will promote GEP instructions with constant
   // indices into GEP constant-exprs, which will allow global-opt to hack on it.
-  ConstantPropUsersOf(NewGV, DL, TLI);
-  if (RepValue != NewGV)
-    ConstantPropUsersOf(RepValue, DL, TLI);
+  for (auto *CE : RepValues)
+    ConstantPropUsersOf(CE, DL, TLI);
 
   return NewGV;
 }

llvm/test/Transforms/GlobalOpt/2021-08-03-StoreOnceLoadMultiCasts.ll

@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -globalopt -S < %s | FileCheck %s
+; RUN: opt -passes=globalopt -S < %s | FileCheck %s
+
+@g = internal global i32* null, align 8
+
+define signext i32 @f() local_unnamed_addr {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @f1()
+; CHECK-NEXT:    store i32 1, i32* @g.body, align 4
+; CHECK-NEXT:    call void @f1()
+; CHECK-NEXT:    store i8 2, i8* bitcast (i32* @g.body to i8*), align 4
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  %call = call i8* @malloc(i64 4)
+  %b = bitcast i8* %call to i32*
+  store i32* %b, i32** @g, align 8
+  call void @f1()
+  %0 = load i32*, i32** @g, align 8
+  store i32 1, i32* %0, align 4
+  call void @f1()
+  %1 = load i8*, i8** bitcast (i32** @g to i8**), align 8
+  store i8 2, i8* %1, align 4
+  ret i32 1
+}
+
+define signext i32 @main() {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CALL:%.*]] = call signext i32 @f()
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* @g.body, align 4
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %call = call signext i32 @f()
+  %0 = load i32*, i32** @g, align 8
+  %1 = load i32, i32* %0, align 4
+  ret i32 %1
+}
+
+declare noalias align 16 i8* @malloc(i64)
+declare void @f1()