AMDGPU: Special case uniformity info for single lane workgroups

[arsenm]

Summary

Constructors/destructors and OpenMP make use of single lane groups
in some cases.

[arsenm]

Reorder

[jhuber6]

Shouldn't this similarly apply to wave syncs and ballots?

[arsenm]

In D151341#4368700, @jhuber6 wrote:

Shouldn't this similarly apply to wave syncs and ballots?

Ballots are forcibly uniform all the time. I guess this works for every operation and don't actually need to look for specific intrinsics

[jhuber6]

In D151341#4368704, @arsenm wrote:

Ballots are forcibly uniform all the time. I guess this works for every operation and don't actually need to look for specific intrinsics

We should be able to remove every convergent attribute under this special case, right.

[yassingh]

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

910–914

Will this also work for GMIR?
eg "%4:_(s32) = G_INTRINSIC intrinsic(@llvm.amdgcn.workitem.id.x)"

llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll

44–50

Same as above, equivalent GMIR test?

[arsenm]

In D151341#4368708, @jhuber6 wrote:

In D151341#4368704, @arsenm wrote:

Ballots are forcibly uniform all the time. I guess this works for every operation and don't actually need to look for specific intrinsics

We should be able to remove every convergent attribute under this special case, right.

Maybe, in the future with convergence tokens we should have a convergence reduction optimization.

I'm thinking a cleaner way to handle this would be to just have the divergence analysis early exit for single lane functions. Divergence cannot exist in the first place

[nhaehnle]

This could be a top-level check in isSourceOfDivergence: if a source of divergence executes with a single lane, it de facto stops being a source of divergence.

But even beyond that: it is reasonably common to have compute shaders that do something like this:

if (threadid == 0) {
  ...
}

So it would be interesting to be able to handle this more generally in UniformityAnalysis, perhaps by having a callback that analyzes conditional branches and returns whether one side of the branch is entered by at most one lane.

Once that's the long-term goal, it seems plausible to me to make isSingleLaneExecution(Function &) a callback that is accessible to UniformityAnalysis, as opposed to querying it from isSourceOfDivergence.

[jhuber6]

In D151341#4380694, @nhaehnle wrote:

But even beyond that: it is reasonably common to have compute shaders that do something like this:

if (threadid == 0) {
  ...
}

So it would be interesting to be able to handle this more generally in UniformityAnalysis, perhaps by having a callback that analyzes conditional branches and returns whether one side of the branch is entered by at most one lane.

Once that's the long-term goal, it seems plausible to me to make isSingleLaneExecution(Function &) a callback that is accessible to UniformityAnalysis, as opposed to querying it from isSourceOfDivergence.

We have a similar check for that explicit case in the OpenMPOpt pass https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/IPO/OpenMPOpt.cpp#L2760.

[nhaehnle]

In D151341#4368708, @jhuber6 wrote:

We should be able to remove every convergent attribute under this special case, right.

I'm not convinced that this is true. Consider:

uint64_t helper() {
  return ballot(true);
}

uint64_t tricky() {
  for (int i = 0; i < 64; ++i)
    if (thread == i)
      return helper();
  unreachable();
}

We can miscompile this as follows:

• Mark helper as being "single lane only"
• Remove convergent attributes based on that
• Inline and eliminate the loop, reducing everything down to:

uint64_t tricky() {
  return ballot(true);
}

[arsenm]

In D151341#4380716, @nhaehnle wrote:

I'm not convinced that this is true. Consider:

uint64_t helper() {
  return ballot(true);
}

uint64_t tricky() {
  for (int i = 0; i < 64; ++i)
    if (thread == i)
      return helper();
  unreachable();
}

We can miscompile this as follows:

I think this program is illegal to begin with. The constraint isn't that helper runs with a single lane, it's that it only executes a workgroup of size 1. If the loop in tricky ever went over 1 iteration, that implies execution with a larger workgroup

[sameerds]

In D151341#4369509, @arsenm wrote:

In D151341#4368708, @jhuber6 wrote:

In D151341#4368704, @arsenm wrote:

Ballots are forcibly uniform all the time. I guess this works for every operation and don't actually need to look for specific intrinsics

We should be able to remove every convergent attribute under this special case, right.

Maybe, in the future with convergence tokens we should have a convergence reduction optimization.

I'm thinking a cleaner way to handle this would be to just have the divergence analysis early exit for single lane functions. Divergence cannot exist in the first place

That sounds right to me. Much simpler and assertive than checking for just a handful of intrinsics. Eventually, this can be an input to deciding whether to remove the convergent attribute, but that's a different task.

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

910–914

The change is in TTI, which is clearly meant for LLVM IR. But from the rest of the comments, if we were to make this change in the UA instead, then we could make it work for both LLVM IR and MIR.

[arsenm]

Skip analysis and make everything uniform

[arsenm]

llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll

44–50

Can't test MIR as-is, the TTI based machine uniformity query is broken as mentioned in the last patch

[arsenm]

53fb907df4723f5267f30fe8da103f91dfb1a175

llvm/lib/Analysis/UniformityAnalysis.cpp

 
 llvm::UniformityInfo UniformityInfoAnalysis::run(Function &F,
                                                  FunctionAnalysisManager &FAM) {
   auto &DT = FAM.getResult<DominatorTreeAnalysis>(F);
   auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
   auto &CI = FAM.getResult<CycleAnalysis>(F);
   UniformityInfo UI{F, DT, CI, &TTI};
   // Skip computation if we can assume everything is uniform.
-  if (TTI.hasBranchDivergence())
+  if (TTI.hasBranchDivergence(&F))
     UI.compute();
 
   return UI;
 }
 
 AnalysisKey UniformityInfoAnalysis::Key;
 
 UniformityInfoPrinterPass::UniformityInfoPrinterPass(raw_ostream &OS)
bool UniformityInfoWrapperPass::runOnFunction(Function &F) {
   auto &targetTransformInfo =
       getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
   m_function = &F;
   m_uniformityInfo =
       UniformityInfo{F, domTree, cycleInfo, &targetTransformInfo};
 
   // Skip computation if we can assume everything is uniform.
-  if (targetTransformInfo.hasBranchDivergence())
+  if (targetTransformInfo.hasBranchDivergence(m_function))
     m_uniformityInfo.compute();
 
   return false;
 }
 
 void UniformityInfoWrapperPass::print(raw_ostream &OS, const Module *) const {
   OS << "UniformityInfo for function '" << m_function->getName() << "':\n";
 }
 
 void UniformityInfoWrapperPass::releaseMemory() {
   m_uniformityInfo = UniformityInfo{};
   m_function = nullptr;
 }

llvm/lib/Target/AMDGPU/AMDGPURegBankSelect.cpp

 //===- AMDGPURegBankSelect.cpp -----------------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Use MachineUniformityAnalysis as the primary basis for making SGPR vs. VGPR
 // register bank selection. Use/def analysis as in the default RegBankSelect can
 // be useful in narrower circumstances (e.g. choosing AGPR vs. VGPR for gfx908).
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPURegBankSelect.h"
 #include "AMDGPU.h"
+#include "GCNSubtarget.h"
 #include "llvm/CodeGen/MachineUniformityAnalysis.h"
 #include "llvm/InitializePasses.h"
 
 #define DEBUG_TYPE "regbankselect"
 
 using namespace llvm;
 
 AMDGPURegBankSelect::AMDGPURegBankSelect(Mode RunningMode)
bool AMDGPURegBankSelect::runOnMachineFunction(MachineFunction &MF) {
   const Function &F = MF.getFunction();
   Mode SaveOptMode = OptMode;
   if (F.hasOptNone())
     OptMode = Mode::Fast;
   init(MF);
 
   assert(checkFunctionIsLegal(MF));
 
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   MachineCycleInfo &CycleInfo =
       getAnalysis<MachineCycleInfoWrapperPass>().getCycleInfo();
   MachineDominatorTree &DomTree = getAnalysis<MachineDominatorTree>();
 
-  // TODO: Check for single lane execution.
   MachineUniformityInfo Uniformity =
-      computeMachineUniformityInfo(MF, CycleInfo, DomTree.getBase(), true);
+      computeMachineUniformityInfo(MF, CycleInfo, DomTree.getBase(),
+                                   !ST.isSingleLaneExecution(F));
   (void)Uniformity; // TODO: Use this
 
   assignRegisterBanks(MF);
 
   OptMode = SaveOptMode;
   return false;
 }

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h

public:
   /// \returns Maximum number of waves per execution unit supported by the
   /// subtarget without any kind of limitation.
   unsigned getMaxWavesPerEU() const { return MaxWavesPerEU; }
 
   /// Return the maximum workitem ID value in the function, for the given (0, 1,
   /// 2) dimension.
   unsigned getMaxWorkitemID(const Function &Kernel, unsigned Dimension) const;
 
+  /// Return true if only a single workitem can be active in a wave.
+  bool isSingleLaneExecution(const Function &Kernel) const;
+
   /// Creates value range metadata on an workitemid.* intrinsic call or load.
   bool makeLIDRangeMetadata(Instruction *I) const;
 
   /// \returns Number of bytes of arguments that are passed to a shader or
   /// kernel in addition to the explicit ones declared for the function.
   unsigned getImplicitArgNumBytes(const Function &F) const;
   uint64_t getExplicitKernArgSize(const Function &F, Align &MaxAlign) const;
   unsigned getKernArgSegmentSize(const Function &F, Align &MaxAlign) const;

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

 unsigned AMDGPUSubtarget::getMaxWorkitemID(const Function &Kernel,
                                            unsigned Dimension) const {
   unsigned ReqdSize = getReqdWorkGroupSize(Kernel, Dimension);
   if (ReqdSize != std::numeric_limits<unsigned>::max())
     return ReqdSize - 1;
   return getFlatWorkGroupSizes(Kernel).second - 1;
 }
 
+bool AMDGPUSubtarget::isSingleLaneExecution(const Function &Func) const {
+  for (int I = 0; I < 3; ++I) {
+    if (getMaxWorkitemID(Func, I) > 0)
+      return false;
+  }
+
+  return true;
+}
+
 bool AMDGPUSubtarget::makeLIDRangeMetadata(Instruction *I) const {
   Function *Kernel = I->getParent()->getParent();
   unsigned MinSize = 0;
   unsigned MaxSize = getFlatWorkGroupSizes(*Kernel).second;
   bool IdQuery = false;
 
   // If reqd_work_group_size is present it narrows value down.
   if (auto *CI = dyn_cast<CallInst>(I)) {

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

: BaseT(TM, F.getParent()->getDataLayout()),
       TLI(ST->getTargetLowering()), CommonTTI(TM, F),
       IsGraphics(AMDGPU::isGraphics(F.getCallingConv())) {
   SIModeRegisterDefaults Mode(F);
   HasFP32Denormals = Mode.allFP32Denormals();
   HasFP64FP16Denormals = Mode.allFP64FP16Denormals();
 }
 
 bool GCNTTIImpl::hasBranchDivergence(const Function *F) const {
-  return true;
+  return !F || !ST->isSingleLaneExecution(*F);
 }
 
 unsigned GCNTTIImpl::getNumberOfRegisters(unsigned RCID) const {
   // NB: RCID is not an RCID. In fact it is 0 or 1 for scalar or vector
   // registers. See getRegisterClassForType for the implementation.
   // In this case vector registers are not vector in terms of
   // VGPRs, but those which can hold multiple values.
 
bool GCNTTIImpl::isSourceOfDivergence(const Value *V) const {
 
   // Atomics are divergent because they are executed sequentially: when an
   // atomic operation refers to the same address in each thread, then each
   // thread after the first sees the value written by the previous thread as
   // original value.
   if (isa<AtomicRMWInst>(V) || isa<AtomicCmpXchgInst>(V))
     return true;
 
   if (const IntrinsicInst *Intrinsic = dyn_cast<IntrinsicInst>(V)) {
     if (Intrinsic->getIntrinsicID() == Intrinsic::read_register)
       return isReadRegisterSourceOfDivergence(Intrinsic);
 
     return AMDGPU::isIntrinsicSourceOfDivergence(Intrinsic->getIntrinsicID());
   }
 
   // Assume all function calls are a source of divergence.
   if (const CallInst *CI = dyn_cast<CallInst>(V)) {
     if (CI->isInlineAsm())
       return isInlineAsmSourceOfDivergence(CI);
     return true;
   }

llvm/test/Analysis/UniformityAnalysis/AMDGPU/always_uniform.ll

define void @asm_mixed_sgpr_vgpr(i32 %divergent) {
   %asm = call { i32, i32 } asm "; def $0, $1, $2","=s,=v,v"(i32 %divergent)
   %sgpr = extractvalue { i32, i32 } %asm, 0
   %vgpr = extractvalue { i32, i32 } %asm, 1
   store i32 %sgpr, ptr addrspace(1) undef
   store i32 %vgpr, ptr addrspace(1) undef
   ret void
 }
 
+; CHECK-LABEL: for function 'single_lane_func_arguments':
+; CHECK-NOT: DIVERGENT
+define void @single_lane_func_arguments(i32 %i32, i1 %i1) #2 {
+ ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x() #0
 declare i32 @llvm.amdgcn.readfirstlane(i32) #0
 declare i64 @llvm.amdgcn.icmp.i32(i32, i32, i32) #1
 declare i64 @llvm.amdgcn.fcmp.i32(float, float, i32) #1
 declare i64 @llvm.amdgcn.ballot.i32(i1) #1
 
 attributes #0 = { nounwind readnone }
 attributes #1 = { nounwind readnone convergent }
+attributes #2 = { "amdgpu-flat-work-group-size"="1,1" }

llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll

 
 ; CHECK: DIVERGENT:  %mbcnt.hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 0)
 define amdgpu_kernel void @mbcnt_hi() #1 {
   %mbcnt.hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 0)
   store volatile i32 %mbcnt.hi, ptr addrspace(1) undef
   ret void
 }
 
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_singlethreaded':
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_singlethreaded() #2 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_y_singlethreaded':
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_y_singlethreaded() #2 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.y()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_z_singlethreaded':
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_z_singlethreaded() #2 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.y()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_singlethreaded_md':
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_singlethreaded_md() !reqd_work_group_size !0 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_y_singlethreaded_md':
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_y_singlethreaded_md() !reqd_work_group_size !0 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.y()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_z_singlethreaded_md':
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_z_singlethreaded_md() !reqd_work_group_size !0 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.y()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_not_singlethreaded_dimx':
+; CHECK: DIVERGENT:  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+define amdgpu_kernel void @workitem_id_x_not_singlethreaded_dimx() !reqd_work_group_size !1 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_not_singlethreaded_dimy':
+; CHECK: DIVERGENT:  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+define amdgpu_kernel void @workitem_id_x_not_singlethreaded_dimy() !reqd_work_group_size !2 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_not_singlethreaded_dimz':
+; CHECK: DIVERGENT:  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+define amdgpu_kernel void @workitem_id_x_not_singlethreaded_dimz() !reqd_work_group_size !3 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  store volatile i32 %id.x, ptr addrspace(1) undef
+  ret void
+}
+
 attributes #0 = { nounwind readnone }
 attributes #1 = { nounwind }
+attributes #2 = { "amdgpu-flat-work-group-size"="1,1" }
+
+!0 = !{i32 1, i32 1, i32 1}
+!1 = !{i32 2, i32 1, i32 1}
+!2 = !{i32 1, i32 2, i32 1}
+!3 = !{i32 1, i32 1, i32 2}