InferAddressSpaces: Support atomics

[jlebar]

lib/Transforms/Scalar/InferAddressSpaces.cpp

243–244

auto*?

249

Would this be better as

auto push = [&](Value* Ptr) {
  appendsGenericAddressExpressionToPostorderStack(
    Ptr, &PostorderStack, &Visited);
};
if (LoadInst *LI = ...)
  push(LI->getPointerOperand());

etc? As is I have to read closely to see whether or not there's anything interesting going on in these calls.

538

Can we use a different word than "compatible"? It's not clear obvious what is the indirect object. A brief comment might also be helpful here.

[arsenm]

Address comments

[jlebar]

lib/Transforms/Scalar/InferAddressSpaces.cpp

236

Nit, this is a variable, so it should not use the function naming convention, counterintuitive as that may be.

[arsenm]

r293584

lib/Transforms/Scalar/InferAddressSpaces.cpp

 // ordered in postorder.
 std::vector<Value *>
 InferAddressSpaces::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;
+
+  auto pushPtrOperand = [&](Value *Ptr) {
+    appendsGenericAddressExpressionToPostorderStack(
+      Ptr, &PostorderStack, &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)) {
-    if (isa<LoadInst>(I)) {
-      appendsGenericAddressExpressionToPostorderStack(
-        I.getOperand(0), &PostorderStack, &Visited);
-    } else if (isa<StoreInst>(I)) {
-      appendsGenericAddressExpressionToPostorderStack(
-        I.getOperand(1), &PostorderStack, &Visited);
-    }
+    if (auto *LI = dyn_cast<LoadInst>(&I))
+      pushPtrOperand(LI->getPointerOperand());
+    else if (auto *SI = dyn_cast<StoreInst>(&I))
+      pushPtrOperand(SI->getPointerOperand());
+    else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
+      pushPtrOperand(RMW->getPointerOperand());
+    else if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(&I))
+      pushPtrOperand(CmpX->getPointerOperand());
+
+    // TODO: Support intrinsics
   }
 
   std::vector<Value *> Postorder; // The resultant postorder.
   while (!PostorderStack.empty()) {
     // If the operands of the expression on the top are already explored,
     // adds that expression to the resultant postorder.
     if (PostorderStack.back().second) {
       Postorder.push_back(PostorderStack.back().first);
Optional<unsigned> InferAddressSpaces::updateAddressSpace(
 
   unsigned OldAS = InferredAddrSpace.lookup(&V);
   assert(OldAS != GenericAddrSpace);
   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) {
+  User *Inst = U.getUser();
+  unsigned OpNo = U.getOperandNo();
+
+  if (auto *LI = dyn_cast<LoadInst>(Inst))
+    return OpNo == LoadInst::getPointerOperandIndex() && !LI->isVolatile();
+
+  if (auto *SI = dyn_cast<StoreInst>(Inst))
+    return OpNo == StoreInst::getPointerOperandIndex() && !SI->isVolatile();
+
+  if (auto *RMW = dyn_cast<AtomicRMWInst>(Inst))
+    return OpNo == AtomicRMWInst::getPointerOperandIndex() && !RMW->isVolatile();
+
+  if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Inst)) {
+    return OpNo == AtomicCmpXchgInst::getPointerOperandIndex() &&
+           !CmpX->isVolatile();
+  }
+
+  return false;
+}
+
 bool InferAddressSpaces::rewriteWithNewAddressSpaces(
   const std::vector<Value *> &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;
for (Value *V : Postorder) {
     SmallVector<Use *, 4> Uses;
     for (Use &U : V->uses())
       Uses.push_back(&U);
 
     DEBUG(dbgs() << "Replacing the uses of " << *V
                  << "\n  with\n  " << *NewV << '\n');
 
     for (Use *U : Uses) {
-      LoadInst *LI = dyn_cast<LoadInst>(U->getUser());
-      StoreInst *SI = dyn_cast<StoreInst>(U->getUser());
-
-      if ((LI && !LI->isVolatile()) ||
-          (SI && !SI->isVolatile() &&
-           U->getOperandNo() == StoreInst::getPointerOperandIndex())) {
-        // If V is used as the pointer operand of a load/store, sets the pointer
-        // operand to NewV. This replacement does not change the element type,
-        // so the resultant load/store is still valid.
+      if (isSimplePointerUseValidToReplace(*U)) {
+        // 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);
       } else if (isa<Instruction>(U->getUser())) {
         // Otherwise, replaces the use with generic(NewV).
         // TODO: Some optimization opportunities are missed. For example, in
         //   %0 = icmp eq float* %p, %q
         // if both p and q are inferred to be shared, we can rewrite %0 as
         //   %0 = icmp eq float addrspace(3)* %new_p, %new_q
         // instead of currently

test/Transforms/InferAddressSpaces/AMDGPU/basic.ll

 ; CHECK: %cast = addrspacecast i32 addrspace(1)* %input to i32 addrspace(4)*
 ; CHECK-NEXT: store i32 addrspace(4)* %cast, i32 addrspace(4)* addrspace(1)* %output, align 4
 define void @store_addrspacecast_ptr_value(i32 addrspace(1)* nocapture %input, i32 addrspace(4)* addrspace(1)* nocapture %output) #0 {
   %cast = addrspacecast i32 addrspace(1)* %input to i32 addrspace(4)*
   store i32 addrspace(4)* %cast, i32 addrspace(4)* addrspace(1)* %output, align 4
   ret void
 }
 
+; CHECK-LABEL: @atomicrmw_add_global_to_flat(
+; CHECK-NEXT: %ret = atomicrmw add i32 addrspace(1)* %global.ptr, i32 %y seq_cst
+define i32 @atomicrmw_add_global_to_flat(i32 addrspace(1)* %global.ptr, i32 %y) #0 {
+  %cast = addrspacecast i32 addrspace(1)* %global.ptr to i32 addrspace(4)*
+  %ret = atomicrmw add i32 addrspace(4)* %cast, i32 %y seq_cst
+  ret i32 %ret
+}
+
+; CHECK-LABEL: @atomicrmw_add_group_to_flat(
+; CHECK-NEXT: %ret = atomicrmw add i32 addrspace(3)* %group.ptr, i32 %y seq_cst
+define i32 @atomicrmw_add_group_to_flat(i32 addrspace(3)* %group.ptr, i32 %y) #0 {
+  %cast = addrspacecast i32 addrspace(3)* %group.ptr to i32 addrspace(4)*
+  %ret = atomicrmw add i32 addrspace(4)* %cast, i32 %y seq_cst
+  ret i32 %ret
+}
+
+; CHECK-LABEL: @cmpxchg_global_to_flat(
+; CHECK: %ret = cmpxchg i32 addrspace(1)* %global.ptr, i32 %cmp, i32 %val seq_cst monotonic
+define { i32, i1 } @cmpxchg_global_to_flat(i32 addrspace(1)* %global.ptr, i32 %cmp, i32 %val) #0 {
+  %cast = addrspacecast i32 addrspace(1)* %global.ptr to i32 addrspace(4)*
+  %ret = cmpxchg i32 addrspace(4)* %cast, i32 %cmp, i32 %val seq_cst monotonic
+  ret { i32, i1 } %ret
+}
+
+; CHECK-LABEL: @cmpxchg_group_to_flat(
+; CHECK: %ret = cmpxchg i32 addrspace(3)* %group.ptr, i32 %cmp, i32 %val seq_cst monotonic
+define { i32, i1 } @cmpxchg_group_to_flat(i32 addrspace(3)* %group.ptr, i32 %cmp, i32 %val) #0 {
+  %cast = addrspacecast i32 addrspace(3)* %group.ptr to i32 addrspace(4)*
+  %ret = cmpxchg i32 addrspace(4)* %cast, i32 %cmp, i32 %val seq_cst monotonic
+  ret { i32, i1 } %ret
+}
+
+; Not pointer operand
+; CHECK-LABEL: @cmpxchg_group_to_flat_wrong_operand(
+; CHECK: %cast.cmp = addrspacecast i32 addrspace(3)* %cmp.ptr to i32 addrspace(4)*
+; CHECK: %ret = cmpxchg i32 addrspace(4)* addrspace(3)* %cas.ptr, i32 addrspace(4)* %cast.cmp, i32 addrspace(4)* %val seq_cst monotonic
+define { i32 addrspace(4)*, i1 } @cmpxchg_group_to_flat_wrong_operand(i32 addrspace(4)* addrspace(3)* %cas.ptr, i32 addrspace(3)* %cmp.ptr, i32 addrspace(4)* %val) #0 {
+  %cast.cmp = addrspacecast i32 addrspace(3)* %cmp.ptr to i32 addrspace(4)*
+  %ret = cmpxchg i32 addrspace(4)* addrspace(3)* %cas.ptr, i32 addrspace(4)* %cast.cmp, i32 addrspace(4)* %val seq_cst monotonic
+  ret { i32 addrspace(4)*, i1 } %ret
+}
+
 attributes #0 = { nounwind }

test/Transforms/InferAddressSpaces/AMDGPU/volatile.ll

 define void @volatile_store_flat_to_private(i32* nocapture %input, i32* nocapture %output) #0 {
   %tmp0 = addrspacecast i32* %input to i32 addrspace(4)*
   %tmp1 = addrspacecast i32* %output to i32 addrspace(4)*
   %val = load i32, i32 addrspace(4)* %tmp0, align 4
   store volatile i32 %val, i32 addrspace(4)* %tmp1, align 4
   ret void
 }
 
+; CHECK-LABEL: @volatile_atomicrmw_add_group_to_flat(
+; CHECK: addrspacecast i32 addrspace(3)* %group.ptr to i32 addrspace(4)*
+; CHECK: atomicrmw volatile add i32 addrspace(4)*
+define i32 @volatile_atomicrmw_add_group_to_flat(i32 addrspace(3)* %group.ptr, i32 %y) #0 {
+  %cast = addrspacecast i32 addrspace(3)* %group.ptr to i32 addrspace(4)*
+  %ret = atomicrmw volatile add i32 addrspace(4)* %cast, i32 %y seq_cst
+  ret i32 %ret
+}
+
+; CHECK-LABEL: @volatile_atomicrmw_add_global_to_flat(
+; CHECK: addrspacecast i32 addrspace(1)* %global.ptr to i32 addrspace(4)*
+; CHECK: %ret = atomicrmw volatile add i32 addrspace(4)*
+define i32 @volatile_atomicrmw_add_global_to_flat(i32 addrspace(1)* %global.ptr, i32 %y) #0 {
+  %cast = addrspacecast i32 addrspace(1)* %global.ptr to i32 addrspace(4)*
+  %ret = atomicrmw volatile add i32 addrspace(4)* %cast, i32 %y seq_cst
+  ret i32 %ret
+}
+
+; CHECK-LABEL: @volatile_cmpxchg_global_to_flat(
+; CHECK: addrspacecast i32 addrspace(1)* %global.ptr to i32 addrspace(4)*
+; CHECK: cmpxchg volatile i32 addrspace(4)*
+define { i32, i1 } @volatile_cmpxchg_global_to_flat(i32 addrspace(1)* %global.ptr, i32 %cmp, i32 %val) #0 {
+  %cast = addrspacecast i32 addrspace(1)* %global.ptr to i32 addrspace(4)*
+  %ret = cmpxchg volatile i32 addrspace(4)* %cast, i32 %cmp, i32 %val seq_cst monotonic
+  ret { i32, i1 } %ret
+}
+
+; CHECK-LABEL: @volatile_cmpxchg_group_to_flat(
+; CHECK: addrspacecast i32 addrspace(3)* %group.ptr to i32 addrspace(4)*
+; CHECK: cmpxchg volatile i32 addrspace(4)*
+define { i32, i1 } @volatile_cmpxchg_group_to_flat(i32 addrspace(3)* %group.ptr, i32 %cmp, i32 %val) #0 {
+  %cast = addrspacecast i32 addrspace(3)* %group.ptr to i32 addrspace(4)*
+  %ret = cmpxchg volatile i32 addrspace(4)* %cast, i32 %cmp, i32 %val seq_cst monotonic
+  ret { i32, i1 } %ret
+}
+
 attributes #0 = { nounwind }
 No newline at end of file