NVPTX: Refactor NVPTXInferAddressSpaces to check TTI

[arsenm]

Summary

Add a new TTI hook for getting the generic address space value.

I considered a more refined TTI hook that could check specific address space combinations, but I don't really have a use for it to justify the increase in complexity.

[arsenm]

Add missing const

[jlebar]

include/llvm/Analysis/TargetTransformInfo.h

230

Nit, can we rephrase so the first sentence (a) says what the function does and (b) isn't a fragment?

230

s/ID value/ID/?

233

Should we say "legal but undesirable"? At which point should we say what we really mean, which is "slower than"?

236

Maybe I'm biased because I'm coming from the CUDA world, but "undesirable" doesn't really seem to capture what's going on as well as 'generic'., and indeed the comment uses "generic" to good effect.

What are you trying to get across with the name, exactly?

236

Maybe this should return optional<unsigned>?

lib/Target/NVPTX/NVPTXInferAddressSpaces.cpp

444

If we used optional<unsigned> we wouldn't need this hack, either.

[arsenm]

include/llvm/Analysis/TargetTransformInfo.h

236

I was trying to avoid a second meaning of the term generic address space. LLVM already unfortunately clobbers the meaning of "generic" address space to mean addrspace(0) and all the assumptions it makes about it. This is generic/flat in a GPU world sense

[jlebar]

I like the new name and comment a lot. Do you think we should rename "GenericAddrSpace"?

I'm still not sold on returning -1 rather than using an optional<unsigned>, though. The meaning of

GenericAddrSpace = TTI.getFlatAddressSpace();
if (GenericAddrSpace == ADDRESS_SPACE_UNINITIALIZED)

is highly nonobvious, relying on the fact that ADDRESS_SPACE_UNINITIALIZED is the same as the default return value of getFlatAddressSpace(). Whereas if we used optional, it would become:

optional<unsigned> MaybeFlatAS = TTI.getFlatAddressSpace();
if (!MaybeFlatAS) return;
FlatAS = *MaybeFlatAS;

which I think pretty directly conveys what we're getting at?

[arsenm]

Renaming

[arsenm]

In D29164#658488, @jlebar wrote:

I like the new name and comment a lot. Do you think we should rename "GenericAddrSpace"?

I'm still not sold on returning -1 rather than using an optional<unsigned>, though. The meaning of

GenericAddrSpace = TTI.getFlatAddressSpace();
if (GenericAddrSpace == ADDRESS_SPACE_UNINITIALIZED)

is highly nonobvious, relying on the fact that ADDRESS_SPACE_UNINITIALIZED is the same as the default return value of getFlatAddressSpace(). Whereas if we used optional, it would become:

optional<unsigned> MaybeFlatAS = TTI.getFlatAddressSpace();
if (!MaybeFlatAS) return;
FlatAS = *MaybeFlatAS;

which I think pretty directly conveys what we're getting at?

I fixed the name of ADDRESS_SPACE_UNINITIALIZED to be UnknownAddressSpace in the cleanup commit which I think is more clear. I'd rather not deviate from the API/convention the other address space TTI functions already use (e.g. LSR defines the same constant and passes it to isLegalAddressingMode sometimes)

include/llvm/Analysis/TargetTransformInfo.h

236

Seems like overkill, and other places already use the -1 convention. Address spaces are only allowed to be 24-bits, so we don't need to spare every bit.

[jlebar]

I fixed the name of ADDRESS_SPACE_UNINITIALIZED to be UnknownAddressSpace in the cleanup commit which I think is more clear. I'd rather not deviate from the API/convention the other address space TTI functions already use (e.g. LSR defines the same constant and passes it to isLegalAddressingMode sometimes)

The documentation as written doesn't ever say that we return -1 if the target doesn't have a flat AS. In fact it's not even explicit whether all targets have a valid "flat" AS -- maybe we just return 0 for "normal" targets.

If you don't want to return an optional, can you please fix the comments? In addition, I think we should use an explicit "-1" literal when we check the return value of getFlatAddressSpace. The UNINITIALIZED constant used in the pass does not mean the same thing as the -1 returned by getFlatAS, and conflating those two harms understanding.

I also think we should change the name back from "unknown AS" to "uninitialized AS". In the pass, we have a lattice of states, where "uninitialized" is at the top, the specific address spaces are in the middle, and "I don't know what the heck this thing is" is at the bottom. Things start in the "uninitialized" state, and they get to the "I don't know" state when we realize that they might be more than one thing, or when we have no idea where they came from.

"Unknown" fits the meaning of this "I don't know" state (it's a "known unknown"), so it's ambiguous whether "unknown AS" is the top or bottom of the lattice. "Uninitialized" is unambiguously the top, so I think it's a better name.

[arsenm]

In D29164#660049, @jlebar wrote:

I fixed the name of ADDRESS_SPACE_UNINITIALIZED to be UnknownAddressSpace in the cleanup commit which I think is more clear. I'd rather not deviate from the API/convention the other address space TTI functions already use (e.g. LSR defines the same constant and passes it to isLegalAddressingMode sometimes)

The documentation as written doesn't ever say that we return -1 if the target doesn't have a flat AS. In fact it's not even explicit whether all targets have a valid "flat" AS -- maybe we just return 0 for "normal" targets.

Then the pass will run and do something nonsensical. I think this would break people using addrspacecast for GC etc.

[arsenm]

Update documentation

[jlebar]

lgtm, thanks for being patient with me.

[arsenm]

r293563 + some function renames to match the new flat terminology

include/llvm/Analysis/TargetTransformInfo.h

public:
   /// \brief Returns whether V is a source of divergence.
   ///
   /// This function provides the target-dependent information for
   /// the target-independent DivergenceAnalysis. DivergenceAnalysis first
   /// builds the dependency graph, and then runs the reachability algorithm
   /// starting with the sources of divergence.
   bool isSourceOfDivergence(const Value *V) const;
 
+  /// For targets with different pointer representations which can be converted
+  /// with the addrspacecast instruction. Returns the address space ID value for
+  /// a target's 'generic' address space. The 'generic' address space in this
+  /// context means a pointer where access of a memory location through a
+  /// pointer with this address space is expected to be undesirable compared to
+  /// the same memory location accessed through a pointer with a different
+  /// address space.
+  unsigned getUndesirableAddressSpace() const;
+
   /// \brief Test whether calls to a function lower to actual program function
   /// calls.
   ///
   /// The idea is to test whether the program is likely to require a 'call'
   /// instruction or equivalent in order to call the given function.
   ///
   /// FIXME: It's not clear that this is a good or useful query API. Client's
   /// should probably move to simpler cost metrics using the above.
public:
   virtual unsigned getInliningThresholdMultiplier() = 0;
   virtual int getIntrinsicCost(Intrinsic::ID IID, Type *RetTy,
                                ArrayRef<Type *> ParamTys) = 0;
   virtual int getIntrinsicCost(Intrinsic::ID IID, Type *RetTy,
                                ArrayRef<const Value *> Arguments) = 0;
   virtual int getUserCost(const User *U) = 0;
   virtual bool hasBranchDivergence() = 0;
   virtual bool isSourceOfDivergence(const Value *V) = 0;
+  virtual unsigned getUndesirableAddressSpace() = 0;
   virtual bool isLoweredToCall(const Function *F) = 0;
   virtual void getUnrollingPreferences(Loop *L, UnrollingPreferences &UP) = 0;
   virtual bool isLegalAddImmediate(int64_t Imm) = 0;
   virtual bool isLegalICmpImmediate(int64_t Imm) = 0;
   virtual bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV,
                                      int64_t BaseOffset, bool HasBaseReg,
                                      int64_t Scale,
                                      unsigned AddrSpace) = 0;
int getIntrinsicCost(Intrinsic::ID IID, Type *RetTy,
                        ArrayRef<const Value *> Arguments) override {
     return Impl.getIntrinsicCost(IID, RetTy, Arguments);
   }
   int getUserCost(const User *U) override { return Impl.getUserCost(U); }
   bool hasBranchDivergence() override { return Impl.hasBranchDivergence(); }
   bool isSourceOfDivergence(const Value *V) override {
     return Impl.isSourceOfDivergence(V);
   }
+
+  unsigned getUndesirableAddressSpace() override {
+    return Impl.getUndesirableAddressSpace();
+  }
+
   bool isLoweredToCall(const Function *F) override {
     return Impl.isLoweredToCall(F);
   }
   void getUnrollingPreferences(Loop *L, UnrollingPreferences &UP) override {
     return Impl.getUnrollingPreferences(L, UP);
   }
   bool isLegalAddImmediate(int64_t Imm) override {
     return Impl.isLegalAddImmediate(Imm);

include/llvm/Analysis/TargetTransformInfoImpl.h

case Intrinsic::coro_subfn_addr:
       return TTI::TCC_Free;
     }
   }
 
   bool hasBranchDivergence() { return false; }
 
   bool isSourceOfDivergence(const Value *V) { return false; }
 
+  unsigned getUndesirableAddressSpace () {
+    return -1;
+  }
+
   bool isLoweredToCall(const Function *F) {
     // FIXME: These should almost certainly not be handled here, and instead
     // handled with the help of TLI or the target itself. This was largely
     // ported from existing analysis heuristics here so that such refactorings
     // can take place in the future.
 
     if (F->isIntrinsic())
       return false;

include/llvm/CodeGen/BasicTTIImpl.h

bool allowsMisalignedMemoryAccesses(LLVMContext &Context,
     EVT E = EVT::getIntegerVT(Context, BitWidth);
     return getTLI()->allowsMisalignedMemoryAccesses(E, AddressSpace, Alignment, Fast);
   }
 
   bool hasBranchDivergence() { return false; }
 
   bool isSourceOfDivergence(const Value *V) { return false; }
 
+  unsigned getUndesirableAddressSpace() {
+    // Return an invalid address space.
+    return -1;
+  }
+
   bool isLegalAddImmediate(int64_t imm) {
     return getTLI()->isLegalAddImmediate(imm);
   }
 
   bool isLegalICmpImmediate(int64_t imm) {
     return getTLI()->isLegalICmpImmediate(imm);
   }
 

lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::hasBranchDivergence() const {
   return TTIImpl->hasBranchDivergence();
 }
 
 bool TargetTransformInfo::isSourceOfDivergence(const Value *V) const {
   return TTIImpl->isSourceOfDivergence(V);
 }
 
+unsigned TargetTransformInfo::getUndesirableAddressSpace() const {
+  return TTIImpl->getUndesirableAddressSpace();
+}
+
 bool TargetTransformInfo::isLoweredToCall(const Function *F) const {
   return TTIImpl->isLoweredToCall(F);
 }
 
 void TargetTransformInfo::getUnrollingPreferences(
     Loop *L, UnrollingPreferences &UP) const {
   return TTIImpl->getUnrollingPreferences(L, UP);
 }

lib/Target/NVPTX/NVPTXInferAddressSpaces.cpp

 // Finally, it fixes the undef in %y' so that
 //   %y' = phi float addrspace(3)* [ %input, %y2' ]
 //
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "nvptx-infer-addrspace"
 
 #include "NVPTX.h"
-#include "MCTargetDesc/NVPTXBaseInfo.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 
 using namespace llvm;
 
 namespace {
 const unsigned ADDRESS_SPACE_UNINITIALIZED = (unsigned)-1;
 
 using ValueToAddrSpaceMapTy = DenseMap<const Value *, unsigned>;
 
 /// \brief NVPTXInferAddressSpaces
 class NVPTXInferAddressSpaces: public FunctionPass {
+  /// Target specific address space which uses of should be replaced if
+  /// possible.
+  unsigned GenericAddrSpace;
+
 public:
   static char ID;
 
   NVPTXInferAddressSpaces() : FunctionPass(ID) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
   }
 
   bool runOnFunction(Function &F) override;
 
 private:
   // Returns the new address space of V if updated; otherwise, returns None.
   Optional<unsigned>
   updateAddressSpace(const Value &V,
                      const ValueToAddrSpaceMapTy &InferredAddrSpace);
 
   // Tries to infer the specific address space of each address expression in
   // Postorder.
   void inferAddressSpaces(const std::vector<Value *> &Postorder,
                           ValueToAddrSpaceMapTy *InferredAddrSpace);
 
   // Changes the generic address expressions in function F to point to specific
   // address spaces if InferredAddrSpace says so. Postorder is the postorder of
   // all generic address expressions in the use-def graph of function F.
   bool
   rewriteWithNewAddressSpaces(const std::vector<Value *> &Postorder,
                               const ValueToAddrSpaceMapTy &InferredAddrSpace,
-                              Function *F);
+                              Function *F) const;
+
+  void appendsGenericAddressExpressionToPostorderStack(
+    Value *V, std::vector<std::pair<Value *, bool>> *PostorderStack,
+    DenseSet<Value *> *Visited) const;
+
+  std::vector<Value *> collectGenericAddressExpressions(Function &F) const;
+  Value *cloneValueWithNewAddressSpace(
+    Value *V, unsigned NewAddrSpace,
+    const ValueToValueMapTy &ValueWithNewAddrSpace,
+    SmallVectorImpl<const Use *> *UndefUsesToFix) const;
+  unsigned joinAddressSpaces(unsigned AS1, unsigned AS2) const;
 };
 } // end anonymous namespace
 
 char NVPTXInferAddressSpaces::ID = 0;
 
 namespace llvm {
 void initializeNVPTXInferAddressSpacesPass(PassRegistry &);
 }
case Instruction::GetElementPtr:
     return {Op.getOperand(0)};
   default:
     llvm_unreachable("Unexpected instruction type.");
   }
 }
 
 // If V is an unvisited generic address expression, appends V to PostorderStack
 // and marks it as visited.
-static void appendsGenericAddressExpressionToPostorderStack(
+void NVPTXInferAddressSpaces::appendsGenericAddressExpressionToPostorderStack(
     Value *V, std::vector<std::pair<Value *, bool>> *PostorderStack,
-    DenseSet<Value *> *Visited) {
+    DenseSet<Value *> *Visited) const {
   assert(V->getType()->isPointerTy());
   if (isAddressExpression(*V) &&
-      V->getType()->getPointerAddressSpace() ==
-          AddressSpace::ADDRESS_SPACE_GENERIC) {
+      V->getType()->getPointerAddressSpace() == GenericAddrSpace) {
     if (Visited->insert(V).second)
       PostorderStack->push_back(std::make_pair(V, false));
   }
 }
 
 // Returns all generic address expressions in function F. The elements are
 // ordered in postorder.
-static std::vector<Value *> collectGenericAddressExpressions(Function &F) {
+std::vector<Value *>
+NVPTXInferAddressSpaces::collectGenericAddressExpressions(Function &F) const {
   // This function implements a non-recursive postorder traversal of a partial
   // use-def graph of function F.
   std::vector<std::pair<Value*, bool>> PostorderStack;
   // The set of visited expressions.
   DenseSet<Value*> Visited;
   // We only explore address expressions that are reachable from loads and
   // stores for now because we aim at generating faster loads and stores.
   for (Instruction &I : instructions(F)) {
static Value *cloneConstantExprWithNewAddressSpace(
   return CE->getWithOperands(NewOperands, TargetType);
 }
 
 // Returns a clone of the value `V`, with its operands replaced as specified in
 // ValueWithNewAddrSpace. This function is called on every generic address
 // expression whose address space needs to be modified, in postorder.
 //
 // See cloneInstructionWithNewAddressSpace for the meaning of UndefUsesToFix.
-static Value *
-cloneValueWithNewAddressSpace(Value *V, unsigned NewAddrSpace,
-                              const ValueToValueMapTy &ValueWithNewAddrSpace,
-                              SmallVectorImpl<const Use *> *UndefUsesToFix) {
+Value *NVPTXInferAddressSpaces::cloneValueWithNewAddressSpace(
+  Value *V, unsigned NewAddrSpace,
+  SmallVectorImpl<const Use *> *UndefUsesToFix) const {
   // All values in Postorder are generic address expressions.
   assert(isAddressExpression(*V) &&
-         V->getType()->getPointerAddressSpace() ==
-             AddressSpace::ADDRESS_SPACE_GENERIC);
+         V->getType()->getPointerAddressSpace() == GenericAddrSpace);
 
   if (Instruction *I = dyn_cast<Instruction>(V)) {
     Value *NewV = cloneInstructionWithNewAddressSpace(
         I, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix);
     if (Instruction *NewI = dyn_cast<Instruction>(NewV)) {
       if (NewI->getParent() == nullptr) {
         NewI->insertBefore(I);
         NewI->takeName(I);
       }
     }
     return NewV;
   }
 
   return cloneConstantExprWithNewAddressSpace(
       cast<ConstantExpr>(V), NewAddrSpace, ValueWithNewAddrSpace);
 }
 
 // Defines the join operation on the address space lattice (see the file header
 // comments).
-static unsigned joinAddressSpaces(unsigned AS1, unsigned AS2) {
-  if (AS1 == AddressSpace::ADDRESS_SPACE_GENERIC ||
-      AS2 == AddressSpace::ADDRESS_SPACE_GENERIC)
-    return AddressSpace::ADDRESS_SPACE_GENERIC;
+unsigned NVPTXInferAddressSpaces::joinAddressSpaces(unsigned AS1,
+                                                    unsigned AS2) const {
+  if (AS1 == GenericAddrSpace || AS2 == GenericAddrSpace)
+    return GenericAddrSpace;
 
   if (AS1 == ADDRESS_SPACE_UNINITIALIZED)
     return AS2;
   if (AS2 == ADDRESS_SPACE_UNINITIALIZED)
     return AS1;
 
   // The join of two different specific address spaces is generic.
-  return AS1 == AS2 ? AS1 : (unsigned)AddressSpace::ADDRESS_SPACE_GENERIC;
+  return (AS1 == AS2) ? AS1 : GenericAddrSpace;
 }
 
 bool NVPTXInferAddressSpaces::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
+  const TargetTransformInfo &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  GenericAddrSpace = TTI.getUndesirableAddressSpace();
+  if (GenericAddrSpace == ADDRESS_SPACE_UNINITIALIZED)
+    return false;
+
   // Collects all generic address expressions in postorder.
   std::vector<Value *> Postorder = collectGenericAddressExpressions(F);
 
   // Runs a data-flow analysis to refine the address spaces of every expression
   // in Postorder.
   ValueToAddrSpaceMapTy InferredAddrSpace;
   inferAddressSpaces(Postorder, &InferredAddrSpace);
 
for (Value *User : V->users()) {
       // Our algorithm only updates the address spaces of generic address
       // expressions, which are those in InferredAddrSpace.
       if (Pos == InferredAddrSpace->end())
         continue;
 
       // Function updateAddressSpace moves the address space down a lattice
       // path. Therefore, nothing to do if User is already inferred as
       // generic (the bottom element in the lattice).
-      if (Pos->second == AddressSpace::ADDRESS_SPACE_GENERIC)
+      if (Pos->second == GenericAddrSpace)
         continue;
 
       Worklist.insert(User);
     }
   }
 }
 
 Optional<unsigned> NVPTXInferAddressSpaces::updateAddressSpace(
     const Value &V, const ValueToAddrSpaceMapTy &InferredAddrSpace) {
   assert(InferredAddrSpace.count(&V));
 
   // The new inferred address space equals the join of the address spaces
   // of all its pointer operands.
   unsigned NewAS = ADDRESS_SPACE_UNINITIALIZED;
   for (Value *PtrOperand : getPointerOperands(V)) {
     unsigned OperandAS;
     if (InferredAddrSpace.count(PtrOperand))
       OperandAS = InferredAddrSpace.lookup(PtrOperand);
     else
       OperandAS = PtrOperand->getType()->getPointerAddressSpace();
     NewAS = joinAddressSpaces(NewAS, OperandAS);
     // join(generic, *) = generic. So we can break if NewAS is already generic.
-    if (NewAS == AddressSpace::ADDRESS_SPACE_GENERIC)
+    if (NewAS == GenericAddrSpace)
       break;
   }
 
   unsigned OldAS = InferredAddrSpace.lookup(&V);
-  assert(OldAS != AddressSpace::ADDRESS_SPACE_GENERIC);
+  assert(OldAS != GenericAddrSpace);
   if (OldAS == NewAS)
     return None;
   return NewAS;
 }
 
 bool NVPTXInferAddressSpaces::rewriteWithNewAddressSpaces(
     const std::vector<Value *> &Postorder,
-    const ValueToAddrSpaceMapTy &InferredAddrSpace, Function *F) {
+    const ValueToAddrSpaceMapTy &InferredAddrSpace, Function *F) const {
   // For each address expression to be modified, creates a clone of it with its
   // pointer operands converted to the new address space. Since the pointer
   // operands are converted, the clone is naturally in the new address space by
   // construction.
   ValueToValueMapTy ValueWithNewAddrSpace;
   SmallVector<const Use *, 32> UndefUsesToFix;
   for (Value* V : Postorder) {
     unsigned NewAddrSpace = InferredAddrSpace.lookup(V);

lib/Target/NVPTX/NVPTXTargetTransformInfo.h

public:
   explicit NVPTXTTIImpl(const NVPTXTargetMachine *TM, const Function &F)
       : BaseT(TM, F.getParent()->getDataLayout()), ST(TM->getSubtargetImpl()),
         TLI(ST->getTargetLowering()) {}
 
   bool hasBranchDivergence() { return true; }
 
   bool isSourceOfDivergence(const Value *V);
 
+  unsigned getUndesirableAddressSpace() const  {
+    return AddressSpace::ADDRESS_SPACE_GENERIC;
+  }
+
   // Increase the inlining cost threshold by a factor of 5, reflecting that
   // calls are particularly expensive in NVPTX.
   unsigned getInliningThresholdMultiplier() { return 5; }
 
   int getArithmeticInstrCost(
       unsigned Opcode, Type *Ty,
       TTI::OperandValueKind Opd1Info = TTI::OK_AnyValue,
       TTI::OperandValueKind Opd2Info = TTI::OK_AnyValue,

test/CodeGen/NVPTX/access-non-generic.ll

 ; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s --check-prefix PTX
 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefix PTX
-; RUN: opt < %s -S -nvptx-infer-addrspace | FileCheck %s --check-prefix IR
+; RUN: opt -mtriple=nvptx-- < %s -S -nvptx-infer-addrspace | FileCheck %s --check-prefix IR
+; RUN: opt -mtriple=nvptx64-- < %s -S -nvptx-infer-addrspace | FileCheck %s --check-prefix IR
 
 @array = internal addrspace(3) global [10 x float] zeroinitializer, align 4
 @scalar = internal addrspace(3) global float 0.000000e+00, align 4
 @generic_scalar = internal global float 0.000000e+00, align 4
 
 define float @ld_from_shared() {
   %1 = addrspacecast float* @generic_scalar to float addrspace(3)*
   %2 = load float, float addrspace(3)* %1