Refactor uniform Y/Z helper, add tests

Author: krzysz00

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

@@ -246,6 +246,24 @@ AMDGPUSubtarget::getReqdWorkGroupSize(const Function &Kernel,
   return std::nullopt;
 }
 
+bool AMDGPUSubtarget::hasWavefrontsEvenlySplittingXDim(
+    const Function &F, bool RequiresUniformYZ) const {
+  auto *Node = F.getMetadata("reqd_work_group_size");
+  if (!Node || Node->getNumOperands() != 3)
+    return false;
+  unsigned XLen =
+      mdconst::extract<ConstantInt>(Node->getOperand(0))->getZExtValue();
+  unsigned YLen =
+      mdconst::extract<ConstantInt>(Node->getOperand(1))->getZExtValue();
+  unsigned ZLen =
+      mdconst::extract<ConstantInt>(Node->getOperand(2))->getZExtValue();
+
+  bool Is1D = YLen <= 1 && ZLen <= 1;
+  bool IsXLargeEnough =
+      isPowerOf2_32(XLen) && (!RequiresUniformYZ || XLen >= getWavefrontSize());
+  return Is1D || IsXLargeEnough;
+}
+
 bool AMDGPUSubtarget::isMesaKernel(const Function &F) const {
   return isMesa3DOS() && !AMDGPU::isShader(F.getCallingConv());
 }

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h

@@ -106,6 +106,20 @@ class AMDGPUSubtarget {
   std::optional<unsigned> getReqdWorkGroupSize(const Function &F,
                                                unsigned Dim) const;
 
+  /// \returns true if \p F will execute in a manner that leaves the X
+  /// dimensions of the workitem ID evenly tiling wavefronts - that is, if X /
+  /// wavefrontsize is uniform. This is true if either the Y and Z block
+  /// dimensions are known to always be 1 or if the X dimension will always be a
+  /// power of 2. If \p RequireUniformYZ is true, it also ensures that the Y and
+  /// Z workitem IDs will be uniform (so, while a (32, 2, 1) launch with
+  /// wavesize64 would ordinarily pass this test, it won't with
+  /// \pRequiresUniformYZ).
+  ///
+  /// This information is currently only gathered from the !reqd_work_group_size
+  /// metadata on \p F, but this may be improved in the future.
+  bool hasWavefrontsEvenlySplittingXDim(const Function &F,
+                                        bool REquiresUniformYZ = false) const;
+
   /// \returns Subtarget's default pair of minimum/maximum number of waves per
   /// execution unit for function \p F, or minimum/maximum number of waves per
   /// execution unit explicitly requested using "amdgpu-waves-per-eu" attribute

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

@@ -1006,18 +1006,12 @@ bool GCNTTIImpl::isSourceOfDivergence(const Value *V) const {
     }
     case Intrinsic::amdgcn_workitem_id_y:
     case Intrinsic::amdgcn_workitem_id_z: {
-      // If the X dimension is guaranteed to launch with a size that is a power
-      // of 2
-      // >= the wavefront size, then the Y and Z dimensions are uniform.
-      // Similarly, if the dimension has size 1, it is also uniform.
       const Function *F = Intrinsic->getFunction();
-      std::optional<unsigned> ReqdXDimSize = ST->getReqdWorkGroupSize(*F, 0);
-      if (ReqdXDimSize && isPowerOf2_32(*ReqdXDimSize) &&
-          *ReqdXDimSize >= ST->getWavefrontSize())
-        return false;
+      bool HasUniformYZ =
+          ST->hasWavefrontsEvenlySplittingXDim(*F, /*RequitezUniformYZ=*/true);
       std::optional<unsigned> ThisDimSize = ST->getReqdWorkGroupSize(
           *F, IID == Intrinsic::amdgcn_workitem_id_y ? 1 : 2);
-      return !ThisDimSize || *ThisDimSize != 1;
+      return !HasUniformYZ && (!ThisDimSize || *ThisDimSize != 1);
     }
     default:
       return AMDGPU::isIntrinsicSourceOfDivergence(IID);
@@ -1073,11 +1067,7 @@ bool GCNTTIImpl::isAlwaysUniform(const Value *V) const {
   bool XDimDoesntResetWithinWaves = false;
   if (auto *I = dyn_cast<Instruction>(V)) {
     const Function *F = I->getFunction();
-    std::optional<unsigned> ReqdXDimSize = ST->getReqdWorkGroupSize(*F, 0);
-    XDimDoesntResetWithinWaves =
-        ST->getMaxWorkitemID(*F, 1) == 0 && ST->getMaxWorkitemID(*F, 2) == 0;
-    if (ReqdXDimSize.has_value() && isPowerOf2_32(*ReqdXDimSize))
-      XDimDoesntResetWithinWaves = true;
+    XDimDoesntResetWithinWaves = ST->hasWavefrontsEvenlySplittingXDim(*F);
   }
   using namespace llvm::PatternMatch;
   uint64_t C;

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

@@ -113,11 +113,40 @@ define amdgpu_kernel void @workitem_id_x_not_singlethreaded_dimz() !reqd_work_gr
   ret void
 }
 
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_z_uniform_len_1'
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_z_uniform_len_1(ptr %o) !reqd_work_group_size !4 {
+  %id.z = call i32 @llvm.amdgcn.workitem.id.z()
+  store i32 %id.z, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_div_wavefront_size(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  %id.sg = lshr i32 %id.x, 6
+  store i32 %id.sg, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_y_uniform_in_subgroup'
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_y_uniform_in_subgroup(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.y = call i32 @llvm.amdgcn.workitem.id.y()
+  store i32 %id.y, ptr %o
+  ret void
+}
+
 attributes #0 = { nounwind readnone }
 attributes #1 = { nounwind }
 attributes #2 = { "amdgpu-flat-work-group-size"="1,1" }
+attributes #3 = { "target-cpu"="gfx900" "amdgpu-flat-work-group-size"="256,256" }
 
 !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}
+!4 = !{i32 64, i32 1, i32 1}
+!5 = !{i32 128, i32 2, i32 1}