[GlobalOpt][FIX] Do not embed initializers into AS!=0 globals

[jdoerfert]

Summary

Not all address spaces support initializers for globals and we can
therefore not set them without checking if they are allowed. This
patch adds a hook into TTI to check if an AS allows non-undef
initializers. We disable it for all but address space 0 by default,
NVPTX and AMDGPU targets allow all but address space 3.

[jdoerfert]

clang format

[tra]

For NVPTX, it's the __shared__ AS that's the odd one. Static initializers do work for __device__ and __constant__ globals.
I think this may be something we should plumb via TTI with AS(0) being the only AS allowed by default.

[jdoerfert]

Introduce TTI hook

[tra]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

While we should not have globals with ADDRESS_SPACE_LOCAL or ADDRESS_SPACE_PARAM, it may be worth to return false for them, too.

llvm/test/Transforms/GlobalOpt/address_space_initializer.ll

2

You may want to add another run for NVPTX with non-default AS and check that we do create the initializer.

[jdoerfert]

Better test coverage, added assertion

[jdoerfert]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

That sounds to me more like an assertion. I'll add one for now, you let me know if that is ok too.

[tra]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

Assertions are for compiler programming errors, not for diagnosing invalid user-provided input.
Considering that one can write IR with a global var in an invalid AS I think assertion is not the right choice here.
An error would be more appropriate, IMO.

[jdoerfert]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

TBH, I think such an input is broken and the verifier should trigger, which it does not: https://godbolt.org/z/xr1v361Wo
Assuming the IR is "valid" we should not misuse this interface with an AS that is not allowed for globals, IMHO.
That said, I'm not even sure what "an error" would look like here.

[tra]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

Right. TTI just does not have enough context to tell whether this particular call indicates a problem. Perhaps the best we can do here is to just provide the answer whether the AS can/can't handle global initializers and leave it up to the callers to deal with overall IR correctness.

In this particular patch, returning false for local/param AS would be valid. It will block the optimization and that would the the correct action regardless of soundness of the IR we're optimizing. Improving target-specific IR checks in the verifier would be nice, but it's beyond the scope of this patch and should be done separately.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

187

Should also disallow REGION_ADDRESS, and PRIVATE_ADDRESS (we don't even allow private globals)

[jdoerfert]

Return false for nonsensical address spaces

[tra]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

Nit: Making it a switch and enumerating all known AS would make it harder to give a wrong answer if/when a new AS is added.

[jdoerfert]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

The AMD enum has dozens of members :( and we already do these checks in other places. I'd leave it like this for now if it's ok.

[tra]

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

52

OK. This is fine, too.

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

109

I'm not sure this should be so strict by default. At the minimum it should probably allow the datalayout's default global address space

[jdoerfert]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

109

I have no strong feelings either way. This seems more cautious and maybe we want targets to overwrite it explicitly? That said, feel free to patch it.

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
   /// they use a flat address space pointer.
   ///
   /// \returns true if the intrinsic was handled.
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const;
 
   bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const;
 
+  /// Return true if globals in this address space can have initializers other
+  /// than `undef`.
+  bool canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const;
+
   unsigned getAssumedAddrSpace(const Value *V) const;
 
   /// Rewrite intrinsic call \p II such that \p OldV will be replaced with \p
   /// NewV, which has a different address space. This should happen for every
   /// operand index that collectFlatAddressOperands returned for the intrinsic.
   /// \returns nullptr if the intrinsic was not handled. Otherwise, returns the
   /// new value (which may be the original \p II with modified operands).
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
public:
   virtual bool hasBranchDivergence() = 0;
   virtual bool useGPUDivergenceAnalysis() = 0;
   virtual bool isSourceOfDivergence(const Value *V) = 0;
   virtual bool isAlwaysUniform(const Value *V) = 0;
   virtual unsigned getFlatAddressSpace() = 0;
   virtual bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                           Intrinsic::ID IID) const = 0;
   virtual bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const = 0;
+  virtual bool
+  canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const = 0;
   virtual unsigned getAssumedAddrSpace(const Value *V) const = 0;
   virtual Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II,
                                                   Value *OldV,
                                                   Value *NewV) const = 0;
   virtual bool isLoweredToCall(const Function *F) = 0;
   virtual void getUnrollingPreferences(Loop *L, ScalarEvolution &,
                                        UnrollingPreferences &UP,
                                        OptimizationRemarkEmitter *ORE) = 0;
bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const override {
     return Impl.collectFlatAddressOperands(OpIndexes, IID);
   }
 
   bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const override {
     return Impl.isNoopAddrSpaceCast(FromAS, ToAS);
   }
 
+  bool
+  canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const override {
+    return Impl.canHaveNonUndefGlobalInitializerInAddressSpace(AS);
+  }
+
   unsigned getAssumedAddrSpace(const Value *V) const override {
     return Impl.getAssumedAddrSpace(V);
   }
 
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const override {
     return Impl.rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
   }

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   unsigned getFlatAddressSpace() const { return -1; }
 
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const {
     return false;
   }
 
   bool isNoopAddrSpaceCast(unsigned, unsigned) const { return false; }
+  bool canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const {
+    return AS == 0;
+  };
 
   unsigned getAssumedAddrSpace(const Value *V) const { return -1; }
 
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const {
     return nullptr;
   }
 

llvm/include/llvm/Transforms/Utils/GlobalStatus.h

 //===- GlobalStatus.h - Compute status info for globals ---------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_GLOBALSTATUS_H
 #define LLVM_TRANSFORMS_UTILS_GLOBALSTATUS_H
 
+#include "llvm/IR/Instructions.h"
 #include "llvm/Support/AtomicOrdering.h"
 
 namespace llvm {
 
 class Constant;
 class Function;
 class Value;
 
enum StoredType {
     NotStored,
 
     /// This global is stored to, but the only thing stored is the constant it
     /// was initialized with. This is only tracked for scalar globals.
     InitializerStored,
 
     /// This global is stored to, but only its initializer and one other value
     /// is ever stored to it.  If this global isStoredOnce, we track the value
-    /// stored to it in StoredOnceValue below.  This is only tracked for scalar
+    /// stored to it via StoredOnceStore below.  This is only tracked for scalar
     /// globals.
     StoredOnce,
 
     /// This global is stored to by multiple values or something else that we
     /// cannot track.
     Stored
   } StoredType = NotStored;
 
   /// If only one value (besides the initializer constant) is ever stored to
-  /// this global, keep track of what value it is.
-  Value *StoredOnceValue = nullptr;
+  /// this global, keep track of what value it is via the store instruction.
+  const StoreInst *StoredOnceStore = nullptr;
+
+  /// If only one value (besides the initializer constant) is ever stored to
+  /// this global return the stored value.
+  Value *getStoredOnceValue() const {
+    return (StoredType == StoredOnce && StoredOnceStore)
+               ? StoredOnceStore->getOperand(0)
+               : nullptr;
+  }
 
   /// These start out null/false.  When the first accessing function is noticed,
   /// it is recorded. When a second different accessing function is noticed,
   /// HasMultipleAccessingFunctions is set to true.
   const Function *AccessingFunction = nullptr;
   bool HasMultipleAccessingFunctions = false;
 
   /// Set to true if this global has a user that is not an instruction (e.g. a

llvm/lib/Analysis/TargetTransformInfo.cpp

bool TargetTransformInfo::collectFlatAddressOperands(
   return TTIImpl->collectFlatAddressOperands(OpIndexes, IID);
 }
 
 bool TargetTransformInfo::isNoopAddrSpaceCast(unsigned FromAS,
                                               unsigned ToAS) const {
   return TTIImpl->isNoopAddrSpaceCast(FromAS, ToAS);
 }
 
+bool TargetTransformInfo::canHaveNonUndefGlobalInitializerInAddressSpace(
+    unsigned AS) const {
+  return TTIImpl->canHaveNonUndefGlobalInitializerInAddressSpace(AS);
+}
+
 unsigned TargetTransformInfo::getAssumedAddrSpace(const Value *V) const {
   return TTIImpl->getAssumedAddrSpace(V);
 }
 
 Value *TargetTransformInfo::rewriteIntrinsicWithAddressSpace(
     IntrinsicInst *II, Value *OldV, Value *NewV) const {
   return TTIImpl->rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
 }

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

unsigned getFlatAddressSpace() const {
     // don't use flat addressing.
     if (IsGraphics)
       return -1;
     return AMDGPUAS::FLAT_ADDRESS;
   }
 
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const;
+
+  bool canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const {
+    return AS != AMDGPUAS::LOCAL_ADDRESS && AS != AMDGPUAS::REGION_ADDRESS &&
+           AS != AMDGPUAS::PRIVATE_ADDRESS;
+  }
+
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const;
 
   bool canSimplifyLegacyMulToMul(const Value *Op0, const Value *Op1,
                                  InstCombiner &IC) const;
   Optional<Instruction *> instCombineIntrinsic(InstCombiner &IC,
                                                IntrinsicInst &II) const;
   Optional<Value *> simplifyDemandedVectorEltsIntrinsic(

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

public:
   bool hasBranchDivergence() { return true; }
 
   bool isSourceOfDivergence(const Value *V);
 
   unsigned getFlatAddressSpace() const {
     return AddressSpace::ADDRESS_SPACE_GENERIC;
   }
 
+  bool canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const {
+    return AS != AddressSpace::ADDRESS_SPACE_SHARED &&
+           AS != AddressSpace::ADDRESS_SPACE_LOCAL && AS != ADDRESS_SPACE_PARAM;
+  }
+
   Optional<Instruction *> instCombineIntrinsic(InstCombiner &IC,
                                                IntrinsicInst &II) const;
 
   // Loads and stores can be vectorized if the alignment is at least as big as
   // the load/store we want to vectorize.
   bool isLegalToVectorizeLoadChain(unsigned ChainSizeInBytes, Align Alignment,
                                    unsigned AddrSpace) const {
     return Alignment >= ChainSizeInBytes;

llvm/lib/Transforms/IPO/GlobalOpt.cpp

for (auto *U : Users) {
     U->destroyConstant();
   }
 }
 
 /// Analyze the specified global variable and optimize
 /// it if possible.  If we make a change, return true.
 static bool
 processInternalGlobal(GlobalVariable *GV, const GlobalStatus &GS,
+                      function_ref<TargetTransformInfo &(Function &)> GetTTI,
                       function_ref<TargetLibraryInfo &(Function &)> GetTLI,
                       function_ref<DominatorTree &(Function &)> LookupDomTree) {
   auto &DL = GV->getParent()->getDataLayout();
   // If this is a first class global and has only one accessing function and
   // this function is non-recursive, we replace the global with a local alloca
   // in this function.
   //
   // NOTE: It doesn't make sense to promote non-single-value types since we
if (GS.StoredType <= GlobalStatus::InitializerStored) {
     // Fall through to the next check; see if we can optimize further.
     ++NumMarked;
   }
   if (!GV->getInitializer()->getType()->isSingleValueType()) {
     const DataLayout &DL = GV->getParent()->getDataLayout();
     if (SRAGlobal(GV, DL))
       return true;
   }
-  if (GS.StoredType == GlobalStatus::StoredOnce && GS.StoredOnceValue) {
+  Value *StoredOnceValue = GS.getStoredOnceValue();
+  if (GS.StoredType == GlobalStatus::StoredOnce && StoredOnceValue) {
     // Avoid speculating constant expressions that might trap (div/rem).
-    auto *SOVConstant = dyn_cast<Constant>(GS.StoredOnceValue);
+    auto *SOVConstant = dyn_cast<Constant>(StoredOnceValue);
     if (SOVConstant && SOVConstant->canTrap())
       return Changed;
 
+    Function &StoreFn =
+        const_cast<Function &>(*GS.StoredOnceStore->getFunction());
     // If the initial value for the global was an undef value, and if only
     // one other value was stored into it, we can just change the
     // initializer to be the stored value, then delete all stores to the
     // global.  This allows us to mark it constant.
+    // This is restricted to address spaces that allow globals to have
+    // initializers. NVPTX, for example, does not support initializers for
+    // shared memory (AS 3).
     if (SOVConstant && SOVConstant->getType() == GV->getValueType() &&
-        isa<UndefValue>(GV->getInitializer())) {
+        isa<UndefValue>(GV->getInitializer()) &&
+        GetTTI(StoreFn).canHaveNonUndefGlobalInitializerInAddressSpace(
+            GV->getType()->getAddressSpace())) {
       // Change the initial value here.
       GV->setInitializer(SOVConstant);
 
       // Clean up any obviously simplifiable users now.
       CleanupConstantGlobalUsers(GV, GV->getInitializer(), DL, GetTLI);
 
       if (GV->use_empty()) {
         LLVM_DEBUG(dbgs() << "   *** Substituting initializer allowed us to "
                           << "simplify all users and delete global!\n");
         GV->eraseFromParent();
         ++NumDeleted;
       }
       ++NumSubstitute;
       return true;
     }
 
     // Try to optimize globals based on the knowledge that only one value
     // (besides its initializer) is ever stored to the global.
-    if (optimizeOnceStoredGlobal(GV, GS.StoredOnceValue, GS.Ordering, DL,
-                                 GetTLI))
+    if (optimizeOnceStoredGlobal(GV, StoredOnceValue, GS.Ordering, DL, GetTLI))
       return true;
 
     // Otherwise, if the global was not a boolean, we can shrink it to be a
     // boolean.
     if (SOVConstant && GS.Ordering == AtomicOrdering::NotAtomic) {
       if (TryToShrinkGlobalToBoolean(GV, SOVConstant)) {
         ++NumShrunkToBool;
         return true;
       }
     }
   }
 
   return Changed;
 }
 
 /// Analyze the specified global variable and optimize it if possible.  If we
 /// make a change, return true.
 static bool
 processGlobal(GlobalValue &GV,
+              function_ref<TargetTransformInfo &(Function &)> GetTTI,
               function_ref<TargetLibraryInfo &(Function &)> GetTLI,
               function_ref<DominatorTree &(Function &)> LookupDomTree) {
   if (GV.getName().startswith("llvm."))
     return false;
 
   GlobalStatus GS;
 
   if (GlobalStatus::analyzeGlobal(&GV, GS))
processGlobal(GlobalValue &GV,
 
   auto *GVar = dyn_cast<GlobalVariable>(&GV);
   if (!GVar)
     return Changed;
 
   if (GVar->isConstant() || !GVar->hasInitializer())
     return Changed;
 
-  return processInternalGlobal(GVar, GS, GetTLI, LookupDomTree) || Changed;
+  return processInternalGlobal(GVar, GS, GetTTI, GetTLI, LookupDomTree) ||
+         Changed;
 }
 
 /// Walk all of the direct calls of the specified function, changing them to
 /// FastCC.
 static void ChangeCalleesToFastCall(Function *F) {
   for (User *U : F->users()) {
     if (isa<BlockAddress>(U))
       continue;
for (Module::iterator FI = M.begin(), E = M.end(); FI != E; ) {
     if (!F->isDeclaration()) {
       if (removeUnreachableBlocks(*F)) {
         auto &DT = LookupDomTree(*F);
         DT.recalculate(*F);
         Changed = true;
       }
     }
 
-    Changed |= processGlobal(*F, GetTLI, LookupDomTree);
+    Changed |= processGlobal(*F, GetTTI, GetTLI, LookupDomTree);
 
     if (!F->hasLocalLinkage())
       continue;
 
     // If we have an inalloca parameter that we can safely remove the
     // inalloca attribute from, do so. This unlocks optimizations that
     // wouldn't be safe in the presence of inalloca.
     // FIXME: We should also hoist alloca affected by this to the entry
if (F->getAttributes().hasAttrSomewhere(Attribute::Nest) &&
       Changed = true;
     }
   }
   return Changed;
 }
 
 static bool
 OptimizeGlobalVars(Module &M,
+                   function_ref<TargetTransformInfo &(Function &)> GetTTI,
                    function_ref<TargetLibraryInfo &(Function &)> GetTLI,
                    function_ref<DominatorTree &(Function &)> LookupDomTree,
                    SmallPtrSetImpl<const Comdat *> &NotDiscardableComdats) {
   bool Changed = false;
 
   for (Module::global_iterator GVI = M.global_begin(), E = M.global_end();
        GVI != E; ) {
     GlobalVariable *GV = &*GVI++;
if (GV->hasInitializer())
           GV->setInitializer(New);
       }
 
     if (deleteIfDead(*GV, NotDiscardableComdats)) {
       Changed = true;
       continue;
     }
 
-    Changed |= processGlobal(*GV, GetTLI, LookupDomTree);
+    Changed |= processGlobal(*GV, GetTTI, GetTLI, LookupDomTree);
   }
   return Changed;
 }
 
 /// Evaluate a piece of a constantexpr store into a global initializer.  This
 /// returns 'Init' modified to reflect 'Val' stored into it.  At this point, the
 /// GEP operands of Addr [0, OpNo) have been stepped into.
 static Constant *EvaluateStoreInto(Constant *Init, Constant *Val,
LocalChange |= OptimizeFunctions(M, GetTLI, GetTTI, GetBFI, LookupDomTree,
                                      NotDiscardableComdats);
 
     // Optimize global_ctors list.
     LocalChange |= optimizeGlobalCtorsList(M, [&](Function *F) {
       return EvaluateStaticConstructor(F, DL, &GetTLI(*F));
     });
 
     // Optimize non-address-taken globals.
-    LocalChange |=
-        OptimizeGlobalVars(M, GetTLI, LookupDomTree, NotDiscardableComdats);
+    LocalChange |= OptimizeGlobalVars(M, GetTTI, GetTLI, LookupDomTree,
+                                      NotDiscardableComdats);
 
     // Resolve aliases, when possible.
     LocalChange |= OptimizeGlobalAliases(M, NotDiscardableComdats);
 
     // Try to remove trivial global destructors if they are not removed
     // already.
     Function *CXAAtExitFn = FindCXAAtExit(M, GetTLI);
     if (CXAAtExitFn)

llvm/lib/Transforms/Utils/GlobalStatus.cpp

if (const ConstantExpr *CE = dyn_cast<ConstantExpr>(UR)) {
               if (GS.StoredType < GlobalStatus::InitializerStored)
                 GS.StoredType = GlobalStatus::InitializerStored;
             } else if (isa<LoadInst>(StoredVal) &&
                        cast<LoadInst>(StoredVal)->getOperand(0) == GV) {
               if (GS.StoredType < GlobalStatus::InitializerStored)
                 GS.StoredType = GlobalStatus::InitializerStored;
             } else if (GS.StoredType < GlobalStatus::StoredOnce) {
               GS.StoredType = GlobalStatus::StoredOnce;
-              GS.StoredOnceValue = StoredVal;
+              GS.StoredOnceStore = SI;
             } else if (GS.StoredType == GlobalStatus::StoredOnce &&
-                       GS.StoredOnceValue == StoredVal) {
+                       GS.getStoredOnceValue() == StoredVal) {
               // noop.
             } else {
               GS.StoredType = GlobalStatus::Stored;
             }
           } else {
             GS.StoredType = GlobalStatus::Stored;
           }
         }

llvm/test/Transforms/GlobalOpt/address_space_initializer.ll

+; RUN: opt -passes=globalopt < %s -S | FileCheck %s
+; RUN: opt -passes=globalopt --mtriple=nvptx64 < %s -S | FileCheck %s --check-prefix=GPU
+; RUN: opt -passes=globalopt --mtriple=amdgcn < %s -S | FileCheck %s --check-prefix=GPU
+
+; Check that we don't try to set a global initializer for non AS(0) globals.
+
+@g0 = internal global i16 undef
+@g1 = internal addrspace(3) global i16 undef
+@g2 = internal addrspace(1) global i16 undef
+; CHECK: internal unnamed_addr constant i16 3
+; CHECK: internal unnamed_addr addrspace(3) global i16 undef
+; CHECK: internal unnamed_addr addrspace(1) global i16 undef
+; GPU: internal unnamed_addr constant i16 3
+; GPU: internal unnamed_addr addrspace(3) global i16 undef
+; GPU: internal unnamed_addr addrspace(1) constant i16 7
+
+define void @a() {
+  store i16 3, i16* @g0, align 8
+  store i16 5, i16* addrspacecast (i16 addrspace(3)* @g1 to i16*), align 8
+  store i16 7, i16* addrspacecast (i16 addrspace(1)* @g2 to i16*), align 8
+  ret void
+}
+
+define i8 @get0() {
+  %bc = bitcast i16* @g0 to i8*
+  %gep = getelementptr i8, i8* %bc, i64 1
+  %r = load i8, i8* %gep
+  ret i8 %r
+}
+define i8 @get1() {
+  %ac = addrspacecast i16 addrspace(3)* @g1 to i16*
+  %bc = bitcast i16* %ac to i8*
+  %gep = getelementptr i8, i8* %bc, i64 1
+  %r = load i8, i8* %gep
+  ret i8 %r
+}
+define i8 @get2() {
+  %ac = addrspacecast i16 addrspace(1)* @g2 to i16*
+  %bc = bitcast i16* %ac to i8*
+  %gep = getelementptr i8, i8* %bc, i64 1
+  %r = load i8, i8* %gep
+  ret i8 %r
+}