[AMDGPU] Allow scalar loads after barrier

[rampitec]

Summary

Currently we cannot convert a vector load into scalar if there
is dominating barrier or fence. It is considered a clobbering
memory access to prevent memory operations reordering. While
reordering is not possible the actual memory is not being clobbered
by a barrier or fence and we can still use a scalar load for a
uniform pointer.

The solution is not to bail on a first clobbering access but
traverse MemorySSA to the root excluding barriers and fences.

[arsenm]

Can we just mark the barriers as inaccessiblememonly instead

[rampitec]

In D118419#3277932, @arsenm wrote:

Can we just mark the barriers as inaccessiblememonly instead

It will not help MemorySSA but likely will allow unwanted reordering around barriers. Plus there are fences anyway along with any barrier.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

80

Don't need extra new line

93

Don't need extra newline

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

80

It is not extra. This is the actual output on my console:

Checking clobbering of:   %i2 = load i32, i32 addrspace(1)* %i1, align 4, !tbaa !15
  Def:   fence syncscope("workgroup") acquire
  Def:   tail call void @llvm.amdgcn.s.barrier() #3
  Def:   fence syncscope("workgroup") release
      -> no clobber

93

Ditto.

[arsenm]

My gut feeling is this is OK, but I'm also not convinced this is entirely sound. Can you add a test where this does not happen if there is an atomic load?

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

75–129

This needs a comment explaining why this loop is OK

80

Weird that the Instruction and MachineInstr behavior is different

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

75–129

Will add, but in a nutshell like with any SSA we can go up the tree to the root. Along that route we shall see all defs of the memory location in this case. Provided that phis are carefully procesded, which can create loops in a worst case. What we are looking for in that traversal is a store invalidating scalar cache practically.

Some sort of a confirmation GVN does a similar thing. In fact even the Visited set shall not be needed, we will eventually come to the root and exhaust the list, but that will be a bit slower. A testcase for that is memory_phi tests.

Anyhow, I am running PSDB on this to double check.

80

It is.

[rampitec]

• Added atomic load test
• Added comment

PSDB passed.

[rampitec]

Further study showed that a barrier or fence not only defines the load of interest but also clobbering any other load dominating them (for the same reason). The change adds aliasing disambiguation while walking up the tree to only consider defs affecting original memory location.

• Use getClobberingMemoryAccess() in the loop with the memory location of the original load to skip any store which may not alias.
• Added more tests and generated opt metadata checks as well.

[rampitec]

Added -amdgpu-aa to the opt run line in the test to do LDS disambiguation.

[vpykhtin]

Looks almost good

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

92

This can be moved to the bottom of the loop, so you don't need to push first MA and just use it through the first iteration of the loop. The loop would be do-while then. In a case of lucky you won't hit any push.

110

It may worth to make a function like isReallyAClobber just to deduplicate code a bit

[rampitec]

Added helper isReallyAClobber.

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

92

There are a lot of 'continue' here. I need to pop_back the worklist somewhere. Then this is SmallVector and some stack is already preallocated and just initialized.

[vpykhtin]

LGTM.

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

 //===-- AMDGPUAnnotateUniformValues.cpp - ---------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 /// This pass adds amdgpu.uniform metadata to IR values so this information
 /// can be used during instruction selection.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "Utils/AMDGPUBaseInfo.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/LegacyDivergenceAnalysis.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/InitializePasses.h"
 
 #define DEBUG_TYPE "amdgpu-annotate-uniform"
 
 using namespace llvm;
 
 namespace {
 

 
 static void setUniformMetadata(Instruction *I) {
   I->setMetadata("amdgpu.uniform", MDNode::get(I->getContext(), {}));
 }
 static void setNoClobberMetadata(Instruction *I) {
   I->setMetadata("amdgpu.noclobber", MDNode::get(I->getContext(), {}));
 }
 
-bool AMDGPUAnnotateUniformValues::isClobberedInFunction(LoadInst * Load) {
-  const MemoryAccess *MA = MSSA->getWalker()->getClobberingMemoryAccess(Load);
-  return !MSSA->isLiveOnEntryDef(MA);
+bool AMDGPUAnnotateUniformValues::isClobberedInFunction(LoadInst *Load) {
+  MemorySSAWalker *Walker = MSSA->getSkipSelfWalker();
+  SmallVector<MemoryAccess *> WorkList{Walker->getClobberingMemoryAccess(Load)};
+  SmallSet<MemoryAccess *, 8> Visited;
+
+  LLVM_DEBUG(dbgs() << "Checking clobbering of: " << *Load << '\n');
+
+  // Start with a nearest dominating clobbering access, it will be either
+  // live on entry (nothing to do, load is not clobbered), MemoryDef, or
+  // MemoryPhi if several MemoryDefs can define this memory state. In that
+  // case add all Defs to WorkList and continue going up and checking all
+  // the definitions of this memory location until the root. When all the
+  // defs are exhausted and came to the entry state we have no clobber.
+  // Along the scan ignore barriers and fences which are considered clobbers
+  // by the MemorySSA, but not really writing anything into the memory.
+  while (!WorkList.empty()) {
+    MemoryAccess *MA = WorkList.pop_back_val();
+    if (!Visited.insert(MA).second)
+      continue;
+
+    if (MSSA->isLiveOnEntryDef(MA))
+      continue;
+
+    if (MemoryDef *Def = dyn_cast<MemoryDef>(MA)) {
+      Instruction *DefInst = Def->getMemoryInst();
+
+      LLVM_DEBUG(dbgs() << "  Def: " << *DefInst << '\n');
+
+      if (isa<FenceInst>(DefInst)) {
+        WorkList.push_back(Def->getDefiningAccess());
+        continue;
+      }
+
+      if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(DefInst)) {
+        switch (II->getIntrinsicID()) {
+        case Intrinsic::amdgcn_s_barrier:
+        case Intrinsic::amdgcn_wave_barrier:
+          WorkList.push_back(Def->getDefiningAccess());
+          continue;
+        default:
+          break;
+        }
+      }
+
+      LLVM_DEBUG(dbgs() << "      -> load is clobbered\n");
+      return true;
+    }
+
+    const MemoryPhi *Phi = cast<MemoryPhi>(MA);
+    for (auto &Use : Phi->incoming_values())
+      WorkList.push_back(cast<MemoryAccess>(&Use));
+  }
+
+  LLVM_DEBUG(dbgs() << "      -> no clobber\n");
+  return false;
 }
 
 void AMDGPUAnnotateUniformValues::visitBranchInst(BranchInst &I) {
   if (DA->isUniform(&I))
     setUniformMetadata(&I);
 }
 
 void AMDGPUAnnotateUniformValues::visitLoadInst(LoadInst &I) {

llvm/test/CodeGen/AMDGPU/noclobber-barrier.ll

+; RUN: llc -march=amdgcn -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -check-prefix=GCN %s
+
+; Check that barrier or fence in between of loads is not considered a clobber
+; for the purpose of converting vector loads into scalar.
+
+; GCN-LABEL: {{^}}simple_barrier:
+; GCN: s_load_dword s
+; GCN: s_waitcnt lgkmcnt(0)
+; GCN: s_barrier
+; GCN: s_waitcnt lgkmcnt(0)
+; GCN: ; wave barrier
+; GCN-NOT: global_load_dword
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @simple_barrier(i32 addrspace(1)* %arg) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  tail call void @llvm.amdgcn.wave.barrier()
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 1
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %i
+  %i4 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 2
+  store i32 %i3, i32 addrspace(1)* %i4, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}memory_phi_no_clobber:
+; GCN: s_load_dword s
+; GCN: s_waitcnt lgkmcnt(0)
+; GCN: s_waitcnt lgkmcnt(0)
+; GCN: s_barrier
+; GCN-NOT: global_load_dword
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @memory_phi_no_clobber(i32 addrspace(1)* %arg) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  br i1 undef, label %if.then, label %if.else
+
+if.then:
+  tail call void @llvm.amdgcn.s.barrier()
+  br label %if.end
+
+if.else:
+  fence syncscope("workgroup") release
+  br label %if.end
+
+if.end:
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 1
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %i
+  %i4 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 2
+  store i32 %i3, i32 addrspace(1)* %i4, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}memory_phi_clobber1:
+; GCN: s_load_dword s
+; GCN: s_barrier
+; GCN: global_store_dword
+; GCN: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @memory_phi_clobber1(i32 addrspace(1)* %arg) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  br i1 undef, label %if.then, label %if.else
+
+if.then:
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 3
+  store i32 1, i32 addrspace(1)* %gep, align 4
+  br label %if.end
+
+if.else:
+  tail call void @llvm.amdgcn.s.barrier()
+  br label %if.end
+
+if.end:
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 1
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %i
+  %i4 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 2
+  store i32 %i3, i32 addrspace(1)* %i4, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}memory_phi_clobber2:
+; GCN-DAG: s_load_dword s
+; GCN-DAG: global_store_dword
+; GCN: s_barrier
+; GCN: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @memory_phi_clobber2(i32 addrspace(1)* %arg) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  br i1 undef, label %if.then, label %if.else
+
+if.then:
+  tail call void @llvm.amdgcn.s.barrier()
+  br label %if.end
+
+if.else:
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 3
+  store i32 1, i32 addrspace(1)* %gep, align 4
+  br label %if.end
+
+if.end:
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 1
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %i
+  %i4 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 2
+  store i32 %i3, i32 addrspace(1)* %i4, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}clobbering_loop:
+; GCN: s_load_dword s
+; GCN: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @clobbering_loop(i32 addrspace(1)* %arg, i1 %cc) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  br label %while.cond
+
+while.cond:
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 1
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %i
+  %i4 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 2
+  store i32 %i3, i32 addrspace(1)* %i4, align 4
+  tail call void @llvm.amdgcn.wave.barrier()
+  br i1 %cc, label %while.cond, label %end
+
+end:
+  ret void
+}
+
+; GCN-LABEL: {{^}}no_clobbering_loop:
+; GCN: s_load_dword s
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @no_clobbering_loop(i32 addrspace(1)* noalias %arg, i32 addrspace(1)* noalias %out, i32 %n) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  br label %while.cond
+
+while.cond:
+  %c = phi i32 [ 0, %bb ], [ %inc, %while.cond ]
+  %acc = phi i32 [ %i, %bb ], [ %i3, %while.cond ]
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i32 %c
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %acc
+  tail call void @llvm.amdgcn.wave.barrier()
+  %inc = add nuw nsw i32 %c, 1
+  %cc = icmp eq i32 %inc, %n
+  br i1 %cc, label %while.cond, label %end
+
+end:
+  store i32 %i3, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}clobber_by_atomic_load:
+; GCN: s_load_dword s
+; GCN: global_load_dword {{.*}} glc
+; GCN: global_load_dword
+; GCN: global_store_dword
+define amdgpu_kernel void @clobber_by_atomic_load(i32 addrspace(1)* %arg) {
+bb:
+  %i = load i32, i32 addrspace(1)* %arg, align 4
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 2
+  %val = load atomic i32, i32 addrspace(1)* %gep  seq_cst, align 4
+  %i1 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 3
+  %i2 = load i32, i32 addrspace(1)* %i1, align 4
+  %i3 = add i32 %i2, %i
+  %i4 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 4
+  store i32 %i3, i32 addrspace(1)* %i4, align 4
+  ret void
+}
+
+declare void @llvm.amdgcn.s.barrier()
+declare void @llvm.amdgcn.wave.barrier()