[NVPTX] allow address space inference for volatile loads/stores.

[tra]

Summary

If particular target supports volatile memory access operations, we can
avoid casting to generic AS. Currently it's only enabled in NVPTX for
loads and stores that access global & shared AS.

This improves performance in situations when we need to exchange data via shared memory
and need to preserve memory access order using 'volatile'.

[jlebar]

llvm/include/llvm/Analysis/TargetTransformInfo.h

493 ↗

(On Diff #119407)

s/particular/the particular/ or perhaps better /the given/

495 ↗

(On Diff #119407)

I'm not sure this second sentence helps much -- it feels sort of apropos of nothing to me here.

Maybe we should say something to cover the fact that if this returns false, the instruction may still have a volatile variant -- we might simply not have implemented this TTI function for that target. :)

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

712 ↗

(On Diff #119407)

Suggest adding a comment, since this is tricky. Maybe "If the memory instruction is volatile, return true only if the target allows the memory instruction to be volatile in the new address space."

llvm/test/CodeGen/NVPTX/ld-st-addrrspace.py

59 ↗

(On Diff #119407)

I'm not sure it's clearer, but ftr we could get rid of the three maps by doing

for (llvm_type, ptx_type), (llvm_addrspace, ptx_addrspace), volatile in product(
  [("i8", "u8"), ...], [(0, ""), (1, ".global"), ...], [True, False]):

At least I'd recommend getting rid of get_as_id, since it's inconsistent with how we do the other two maps.

73 ↗

(On Diff #119407)

Maybe this would be blearer if we moved it above the params dict?

Also, spaces after the commas. :) There *are* pyformatters out there, e.g. https://github.com/myint/pyformat (first hit when randomly searching), which looks like it's actually just a wrapper around https://pypi.python.org/pypi/autopep8.

76 ↗

(On Diff #119407)

This test_params vs. params difference looks left over from the previous python script. Here it seems unnecessary?

[tra]

Addressed Justin's comments.

Added addrspacecast to generic AS in the tests to make sure InferAddressSpace does its job.

[tra]

llvm/include/llvm/Analysis/TargetTransformInfo.h

495 ↗

(On Diff #119407)

Rephrased. PTAL.

llvm/test/CodeGen/NVPTX/ld-st-addrrspace.py

59 ↗

(On Diff #119407)

Fixed as_id.
Folding three maps into a single statement is a bit too much python-fu, IMO.

[jlebar]

llvm/include/llvm/Analysis/TargetTransformInfo.h

493 ↗

(On Diff #119945)

a volatile variant

[tra]

Grammar fix.

llvm/trunk/include/llvm/Analysis/TargetTransformInfo.h

public:
 
   /// Return true if the target has a unified operation to calculate division
   /// and remainder. If so, the additional implicit multiplication and
   /// subtraction required to calculate a remainder from division are free. This
   /// can enable more aggressive transformations for division and remainder than
   /// would typically be allowed using throughput or size cost models.
   bool hasDivRemOp(Type *DataType, bool IsSigned) const;
 
+  /// Return true if the given instruction (assumed to be a memory access
+  /// instruction) has a volatile variant. If that's the case then we can avoid
+  /// addrspacecast to generic AS for volatile loads/stores. Default
+  /// implementation returns false, which prevents address space inference for
+  /// volatile loads/stores.
+  bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) const;
+
   /// Return true if target doesn't mind addresses in vectors.
   bool prefersVectorizedAddressing() const;
 
   /// \brief Return the cost of the scaling factor used in the addressing
   /// mode represented by AM for this target, for a load/store
   /// of the specified type.
   /// If the AM is supported, the return value must be >= 0.
   /// If the AM is not supported, it returns a negative value.
virtual bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV,
                                      Instruction *I) = 0;
   virtual bool isLSRCostLess(TargetTransformInfo::LSRCost &C1,
                              TargetTransformInfo::LSRCost &C2) = 0;
   virtual bool isLegalMaskedStore(Type *DataType) = 0;
   virtual bool isLegalMaskedLoad(Type *DataType) = 0;
   virtual bool isLegalMaskedScatter(Type *DataType) = 0;
   virtual bool isLegalMaskedGather(Type *DataType) = 0;
   virtual bool hasDivRemOp(Type *DataType, bool IsSigned) = 0;
+  virtual bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) = 0;
   virtual bool prefersVectorizedAddressing() = 0;
   virtual int getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
                                    int64_t BaseOffset, bool HasBaseReg,
                                    int64_t Scale, unsigned AddrSpace) = 0;
   virtual bool LSRWithInstrQueries() = 0;
   virtual bool isTruncateFree(Type *Ty1, Type *Ty2) = 0;
   virtual bool isProfitableToHoist(Instruction *I) = 0;
   virtual bool isTypeLegal(Type *Ty) = 0;
bool isLegalMaskedScatter(Type *DataType) override {
     return Impl.isLegalMaskedScatter(DataType);
   }
   bool isLegalMaskedGather(Type *DataType) override {
     return Impl.isLegalMaskedGather(DataType);
   }
   bool hasDivRemOp(Type *DataType, bool IsSigned) override {
     return Impl.hasDivRemOp(DataType, IsSigned);
   }
+  bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) override {
+    return Impl.hasVolatileVariant(I, AddrSpace);
+  }
   bool prefersVectorizedAddressing() override {
     return Impl.prefersVectorizedAddressing();
   }
   int getScalingFactorCost(Type *Ty, GlobalValue *BaseGV, int64_t BaseOffset,
                            bool HasBaseReg, int64_t Scale,
                            unsigned AddrSpace) override {
     return Impl.getScalingFactorCost(Ty, BaseGV, BaseOffset, HasBaseReg,
                                      Scale, AddrSpace);

llvm/trunk/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   bool isLegalMaskedLoad(Type *DataType) { return false; }
 
   bool isLegalMaskedScatter(Type *DataType) { return false; }
 
   bool isLegalMaskedGather(Type *DataType) { return false; }
 
   bool hasDivRemOp(Type *DataType, bool IsSigned) { return false; }
 
+  bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) { return false; }
+
   bool prefersVectorizedAddressing() { return true; }
 
   int getScalingFactorCost(Type *Ty, GlobalValue *BaseGV, int64_t BaseOffset,
                            bool HasBaseReg, int64_t Scale, unsigned AddrSpace) {
     // Guess that all legal addressing mode are free.
     if (isLegalAddressingMode(Ty, BaseGV, BaseOffset, HasBaseReg,
                               Scale, AddrSpace))
       return 0;

llvm/trunk/lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::isLegalMaskedScatter(Type *DataType) const {
   return TTIImpl->isLegalMaskedScatter(DataType);
 }
 
 bool TargetTransformInfo::hasDivRemOp(Type *DataType, bool IsSigned) const {
   return TTIImpl->hasDivRemOp(DataType, IsSigned);
 }
 
+bool TargetTransformInfo::hasVolatileVariant(Instruction *I,
+                                             unsigned AddrSpace) const {
+  return TTIImpl->hasVolatileVariant(I, AddrSpace);
+}
+
 bool TargetTransformInfo::prefersVectorizedAddressing() const {
   return TTIImpl->prefersVectorizedAddressing();
 }
 
 int TargetTransformInfo::getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
                                               int64_t BaseOffset,
                                               bool HasBaseReg,
                                               int64_t Scale,

llvm/trunk/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

int getArithmeticInstrCost(
       TTI::OperandValueKind Opd1Info = TTI::OK_AnyValue,
       TTI::OperandValueKind Opd2Info = TTI::OK_AnyValue,
       TTI::OperandValueProperties Opd1PropInfo = TTI::OP_None,
       TTI::OperandValueProperties Opd2PropInfo = TTI::OP_None,
       ArrayRef<const Value *> Args = ArrayRef<const Value *>());
 
   void getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
                                TTI::UnrollingPreferences &UP);
+  bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) {
+    // Volatile loads/stores are only supported for shared and global address
+    // spaces, or for generic AS that maps to them.
+    if (!(AddrSpace == llvm::ADDRESS_SPACE_GENERIC ||
+          AddrSpace == llvm::ADDRESS_SPACE_GLOBAL ||
+          AddrSpace == llvm::ADDRESS_SPACE_SHARED))
+      return false;
+
+    switch(I->getOpcode()){
+    default:
+      return false;
+    case Instruction::Load:
+    case Instruction::Store:
+      return true;
+    }
+  }
 };
 
 } // end namespace llvm
 
 #endif

llvm/trunk/lib/Transforms/Scalar/InferAddressSpaces.cpp

private:
   void inferAddressSpaces(ArrayRef<WeakTrackingVH> Postorder,
                           ValueToAddrSpaceMapTy *InferredAddrSpace) const;
 
   bool isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const;
 
   // Changes the flat address expressions in function F to point to specific
   // address spaces if InferredAddrSpace says so. Postorder is the postorder of
   // all flat expressions in the use-def graph of function F.
-  bool
-  rewriteWithNewAddressSpaces(ArrayRef<WeakTrackingVH> Postorder,
-                              const ValueToAddrSpaceMapTy &InferredAddrSpace,
-                              Function *F) const;
+  bool rewriteWithNewAddressSpaces(
+      const TargetTransformInfo &TTI, ArrayRef<WeakTrackingVH> Postorder,
+      const ValueToAddrSpaceMapTy &InferredAddrSpace, Function *F) const;
 
   void appendsFlatAddressExpressionToPostorderStack(
     Value *V, std::vector<std::pair<Value *, bool>> &PostorderStack,
     DenseSet<Value *> &Visited) const;
 
   bool rewriteIntrinsicOperands(IntrinsicInst *II,
                                 Value *OldV, Value *NewV) const;
   void collectRewritableIntrinsicOperands(
bool InferAddressSpaces::runOnFunction(Function &F) {
 
   // Runs a data-flow analysis to refine the address spaces of every expression
   // in Postorder.
   ValueToAddrSpaceMapTy InferredAddrSpace;
   inferAddressSpaces(Postorder, &InferredAddrSpace);
 
   // Changes the address spaces of the flat address expressions who are inferred
   // to point to a specific address space.
-  return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace, &F);
+  return rewriteWithNewAddressSpaces(TTI, Postorder, InferredAddrSpace, &F);
 }
 
 // Constants need to be tracked through RAUW to handle cases with nested
 // constant expressions, so wrap values in WeakTrackingVH.
 void InferAddressSpaces::inferAddressSpaces(
     ArrayRef<WeakTrackingVH> Postorder,
     ValueToAddrSpaceMapTy *InferredAddrSpace) const {
   SetVector<Value *> Worklist(Postorder.begin(), Postorder.end());
Optional<unsigned> InferAddressSpaces::updateAddressSpace(
   assert(OldAS != FlatAddrSpace);
   if (OldAS == NewAS)
     return None;
   return NewAS;
 }
 
 /// \p returns true if \p U is the pointer operand of a memory instruction with
 /// a single pointer operand that can have its address space changed by simply
-/// mutating the use to a new value.
-static bool isSimplePointerUseValidToReplace(Use &U) {
+/// mutating the use to a new value. If the memory instruction is volatile,
+/// return true only if the target allows the memory instruction to be volatile
+/// in the new address space.
+static bool isSimplePointerUseValidToReplace(const TargetTransformInfo &TTI,
+                                             Use &U, unsigned AddrSpace) {
   User *Inst = U.getUser();
   unsigned OpNo = U.getOperandNo();
+  bool VolatileIsAllowed = false;
+  if (auto *I = dyn_cast<Instruction>(Inst))
+    VolatileIsAllowed = TTI.hasVolatileVariant(I, AddrSpace);
 
   if (auto *LI = dyn_cast<LoadInst>(Inst))
-    return OpNo == LoadInst::getPointerOperandIndex() && !LI->isVolatile();
+    return OpNo == LoadInst::getPointerOperandIndex() &&
+           (VolatileIsAllowed || !LI->isVolatile());
 
   if (auto *SI = dyn_cast<StoreInst>(Inst))
-    return OpNo == StoreInst::getPointerOperandIndex() && !SI->isVolatile();
+    return OpNo == StoreInst::getPointerOperandIndex() &&
+           (VolatileIsAllowed || !SI->isVolatile());
 
   if (auto *RMW = dyn_cast<AtomicRMWInst>(Inst))
-    return OpNo == AtomicRMWInst::getPointerOperandIndex() && !RMW->isVolatile();
+    return OpNo == AtomicRMWInst::getPointerOperandIndex() &&
+           (VolatileIsAllowed || !RMW->isVolatile());
 
   if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Inst)) {
     return OpNo == AtomicCmpXchgInst::getPointerOperandIndex() &&
-           !CmpX->isVolatile();
+           (VolatileIsAllowed || !CmpX->isVolatile());
   }
 
   return false;
 }
 
 /// Update memory intrinsic uses that require more complex processing than
 /// simple memory instructions. Thse require re-mangling and may have multiple
 /// pointer operands.
static Value::use_iterator skipToNextUser(Value::use_iterator I,
 
   while (I != End && I->getUser() == CurUser)
     ++I;
 
   return I;
 }
 
 bool InferAddressSpaces::rewriteWithNewAddressSpaces(
-    ArrayRef<WeakTrackingVH> Postorder,
+    const TargetTransformInfo &TTI, ArrayRef<WeakTrackingVH> Postorder,
     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) {
for (const WeakTrackingVH &WVH : Postorder) {
     Value::use_iterator I, E, Next;
     for (I = V->use_begin(), E = V->use_end(); I != E; ) {
       Use &U = *I;
 
       // Some users may see the same pointer operand in multiple operands. Skip
       // to the next instruction.
       I = skipToNextUser(I, E);
 
-      if (isSimplePointerUseValidToReplace(U)) {
+      if (isSimplePointerUseValidToReplace(
+              TTI, U, V->getType()->getPointerAddressSpace())) {
         // If V is used as the pointer operand of a compatible memory operation,
         // sets the pointer operand to NewV. This replacement does not change
         // the element type, so the resultant load/store is still valid.
         U.set(NewV);
         continue;
       }
 
       User *CurUser = U.getUser();

llvm/trunk/test/CodeGen/NVPTX/ld-st-addrrspace.py

+# This test generates all variants of load/store instructions and verifies that
+# LLVM generates correct PTX for them.
+
+# RUN: python %s > %t.ll
+# RUN: llc < %t.ll -march=nvptx64 -mcpu=sm_30 | FileCheck -check-prefixes=CHECK,CHECK_P64 %t.ll
+# RUN: llc < %t.ll -march=nvptx -mcpu=sm_30 | FileCheck -check-prefixes=CHECK,CHECK_P32 %t.ll
+
+from itertools import product
+from string import Template
+
+llvm_type_to_ptx_type = {
+    "i8": "u8",
+    "i16": "u16",
+    "i32": "u32",
+    "i64": "u64",
+    "half": "b16",
+    "<2 x half>": "b32",
+    "float": "f32",
+    "double": "f64"
+}
+
+llvm_type_to_ptx_reg = {
+    "i8": "r",
+    "i16": "r",
+    "i32": "r",
+    "i64": "rd",
+    "half": "h",
+    "<2 x half>": "hh",
+    "float": "f",
+    "double": "fd"
+}
+
+addrspace_id = {
+    "": 0,
+    ".global": 1,
+    ".shared": 3,
+    ".const": 4,
+    ".local": 5,
+    ".param": 101
+}
+
+
+def gen_load_tests():
+  load_template = """
+define ${type} @ld${_volatile}${_space}.${ptx_type}(${type} addrspace(${asid})* %ptr) {
+; CHECK_P32: ld${_volatile}${_volatile_as}.${ptx_type} %${ptx_reg}{{[0-9]+}}, [%r{{[0-9]+}}]
+; CHECK_P64: ld${_volatile}${_volatile_as}.${ptx_type} %${ptx_reg}{{[0-9]+}}, [%rd{{[0-9]+}}]
+; CHECK: ret
+  %p = ${generic_ptr}
+  %a = load ${volatile} ${type}, ${type}* %p
+  ret ${type} %a
+}
+"""
+  for op_type, volatile, space in product(
+      ["i8", "i16", "i32", "i64", "half", "float", "double", "<2 x half>"],
+      [True, False],  # volatile
+      ["", ".shared", ".global", ".const", ".local", ".param"]):
+
+    # Volatile is only supported for global, shared and generic.
+    if volatile and not space in ["", ".global", ".shared"]:
+      continue
+
+    # Volatile is only supported for global, shared and generic.
+    # All other volatile accesses are done in generic AS.
+    if volatile and not space in ["", ".global", ".shared"]:
+      volatile_as = ""
+    else:
+      volatile_as = space
+
+    params = {
+        "type": op_type,
+        "volatile": "volatile" if volatile else "",
+        "_volatile": ".volatile" if volatile else "",
+        "_volatile_as": volatile_as,
+        "_space": space,
+        "ptx_reg": llvm_type_to_ptx_reg[op_type],
+        "ptx_type": llvm_type_to_ptx_type[op_type],
+        "asid": addrspace_id[space],
+    }
+
+    # LLVM does not accept "addrspacecast Type* addrspace(0) to Type*", so we
+    # need to avoid it for generic pointer tests.
+    if space:
+      generic_ptr_template = ("addrspacecast ${type} addrspace(${asid})* %ptr "
+                              "to ${type}*")
+    else:
+      generic_ptr_template = "select i1 true, ${type}* %ptr, ${type}* %ptr"
+    params["generic_ptr"] = Template(generic_ptr_template).substitute(params)
+
+    print(Template(load_template).substitute(params))
+
+
+def main():
+  gen_load_tests()
+
+
+main()