[AMDGPU][AMDGPULowerKernelAttributes] Use stripAndAccumulateConstantOffset

[jmmartinez]

Summary

This removes the need to handle explicitely the presence or absence of bitcasts when typed pointers are used. It also takes into account the case where we have more than one use of load(call @llvm.amdgc.implicitarg.ptr())

[jmmartinez]

• Handle non-inbounds GEP

[jdoerfert]

Some drive by comments

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

118

This function and uses are unrelated.

123

This does break with the "unique" property below. Two users, one simple one not, will result in a load. Is that expected?

139

A unique load reached via select or phi will not be detected. UniqueLoad && UniqueLoad != L can help with that.

Style: I'd remove braces for single return conditionals.

141

I would use the Attributor::checkForAllUses, also unsure why the loads need to be unique. Depending on the call AAPointerInfo would allow you to handle way more things without any of this manual tracing. Either seems to be a useful rewrite beyond the scope of this. Feel free to ignore for now.

[jmmartinez]

• Rename getUniqueLoadUser into getUniqueSimpleLoadUser.
• Consider the case where the UniqueLoad is 'found' twice.
• Style

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

118

I'm sorry, I do not understand this remark. What do you mean?

123

It's ok if there is another non-simple load.

In fact, this pass could be modified to support multiple loads. My guess is that for simplicity only one is supported for the moment.

I renamed the function instead into getUniqueSimpleLoadUser

141

I'm not familiar with it. I'm going to take a look at it. Thanks!

[jdoerfert]

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

118

The commit is about the load tracking and stripping. This change (outlining getBasePtrIntrinsic and using it) is unrelated to the content/description of the patch. Make it two patches instead.

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

118

Oh! I agree. I created https://reviews.llvm.org/D135266 for that change.

[arsenm]

Needs a test for this multiple load case

[jmmartinez]

• Rebase and move annonymous namespace to remove static keyword

[jmmartinez]

• Extend AMDGPULowerKernelAttributes to handle multiple occurrences of the attributes in the same function
• Fix reminder indices in test
• Add test-cases with multiple users, and more than one in a single function

[jmmartinez]

To split into smaller changes later.

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

 
 Function *getBasePtrIntrinsic(Module &M, bool IsV5OrAbove) {
   auto IntrinsicId = IsV5OrAbove ? Intrinsic::amdgcn_implicitarg_ptr
                                  : Intrinsic::amdgcn_dispatch_ptr;
   StringRef Name = Intrinsic::getName(IntrinsicId);
   return M.getFunction(Name);
 }
 
-} // end anonymous namespace
+LoadInst *getUniqueSimpleLoadUser(Value *V) {
+  if (LoadInst *L = dyn_cast<LoadInst>(V)) {
+    return L->isSimple() ? L : nullptr;
+  }
+
+  if (!V->getType()->isPointerTy()) {
+    return nullptr;
+  }
+
+  LoadInst *UniqueLoad = nullptr;
+  for (User *U : V->users()) {
+    if (LoadInst *L = getUniqueSimpleLoadUser(U)) {
+      if (UniqueLoad && UniqueLoad != L)
+        return nullptr;
+      UniqueLoad = L;
+    }
+  }
+  return UniqueLoad;
+}
+
+auto m_MatchAnyOf(ArrayRef<Value *> V) {
+  struct anyofval_ty {
+    const ArrayRef<Value *> Vals;
+    anyofval_ty(const ArrayRef<Value *> V) : Vals(V) {}
+    bool match(Value *V) { return is_contained(Vals, V); }
+  };
+  return anyofval_ty(V);
+}
 
 static bool processUse(CallInst *CI, bool IsV5OrAbove) {
-  Function *F = CI->getParent()->getParent();
+  Function *F = CI->getFunction();
 
   auto MD = F->getMetadata("reqd_work_group_size");
   const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;
 
   const bool HasUniformWorkGroupSize =
     F->getFnAttribute("uniform-work-group-size").getValueAsBool();
 
   if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
     return false;
 
-  Value *BlockCounts[3] = {nullptr, nullptr, nullptr};
-  Value *GroupSizes[3]  = {nullptr, nullptr, nullptr};
-  Value *Remainders[3]  = {nullptr, nullptr, nullptr};
-  Value *GridSizes[3]   = {nullptr, nullptr, nullptr};
+  SmallVector<Value *, 1> BlockCounts[3], GroupSizes[3], Remainders,
+      GridSizes[3];
 
   const DataLayout &DL = F->getParent()->getDataLayout();
 
   // We expect to see several GEP users, casted to the appropriate type and
   // loaded.
   for (User *U : CI->users()) {
-    if (!U->hasOneUse())
-      continue;
-
-    int64_t Offset = 0;
-    auto *Load = dyn_cast<LoadInst>(U); // Load from ImplicitArgPtr/DispatchPtr?
-    auto *BCI = dyn_cast<BitCastInst>(U);
-    if (!Load && !BCI) {
-      if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
-        continue;
-      Load = dyn_cast<LoadInst>(*U->user_begin()); // Load from GEP?
-      BCI = dyn_cast<BitCastInst>(*U->user_begin());
-    }
-
-    if (BCI) {
-      if (!BCI->hasOneUse())
-        continue;
-      Load = dyn_cast<LoadInst>(*BCI->user_begin()); // Load from BCI?
-    }
+    LoadInst *Load = getUniqueSimpleLoadUser(U);
+    if (!Load)
 
-    if (!Load || !Load->isSimple())
+    APInt Offset(64, 0U);
+    if (Load != U &&
+        U->stripAndAccumulateConstantOffsets(DL, Offset, true) != CI)
       continue;
 
     unsigned LoadSize = DL.getTypeStoreSize(Load->getType());
 
     // TODO: Handle merged loads.
+    auto const OffsetValue = Offset.getSExtValue();
     if (IsV5OrAbove) { // Base is ImplicitArgPtr.
-      switch (Offset) {
+      switch (OffsetValue) {
       case HIDDEN_BLOCK_COUNT_X:
         if (LoadSize == 4)
-          BlockCounts[0] = Load;
+          BlockCounts[0].push_back(Load);
         break;
       case HIDDEN_BLOCK_COUNT_Y:
         if (LoadSize == 4)
-          BlockCounts[1] = Load;
+          BlockCounts[1].push_back(Load);
         break;
       case HIDDEN_BLOCK_COUNT_Z:
         if (LoadSize == 4)
-          BlockCounts[2] = Load;
+          BlockCounts[2].push_back(Load);
         break;
       case HIDDEN_GROUP_SIZE_X:
         if (LoadSize == 2)
-          GroupSizes[0] = Load;
+          GroupSizes[0].push_back(Load);
         break;
       case HIDDEN_GROUP_SIZE_Y:
         if (LoadSize == 2)
-          GroupSizes[1] = Load;
+          GroupSizes[1].push_back(Load);
         break;
       case HIDDEN_GROUP_SIZE_Z:
         if (LoadSize == 2)
-          GroupSizes[2] = Load;
+          GroupSizes[2].push_back(Load);
         break;
       case HIDDEN_REMAINDER_X:
-        if (LoadSize == 2)
-          Remainders[0] = Load;
-        break;
       case HIDDEN_REMAINDER_Y:
-        if (LoadSize == 2)
-          Remainders[1] = Load;
-        break;
       case HIDDEN_REMAINDER_Z:
         if (LoadSize == 2)
-          Remainders[2] = Load;
+          Remainders.push_back(Load);
         break;
       default:
         break;
       }
     } else { // Base is DispatchPtr.
-      switch (Offset) {
+      switch (OffsetValue) {
       case WORKGROUP_SIZE_X:
         if (LoadSize == 2)
-          GroupSizes[0] = Load;
+          GroupSizes[0].push_back(Load);
         break;
       case WORKGROUP_SIZE_Y:
         if (LoadSize == 2)
-          GroupSizes[1] = Load;
+          GroupSizes[1].push_back(Load);
         break;
       case WORKGROUP_SIZE_Z:
         if (LoadSize == 2)
-          GroupSizes[2] = Load;
+          GroupSizes[2].push_back(Load);
         break;
       case GRID_SIZE_X:
         if (LoadSize == 4)
-          GridSizes[0] = Load;
+          GridSizes[0].push_back(Load);
         break;
       case GRID_SIZE_Y:
         if (LoadSize == 4)
-          GridSizes[1] = Load;
+          GridSizes[1].push_back(Load);
         break;
       case GRID_SIZE_Z:
         if (LoadSize == 4)
-          GridSizes[2] = Load;
+          GridSizes[2].push_back(Load);
         break;
       default:
         break;
       }
     }
   }
 
   bool MadeChange = false;
   if (IsV5OrAbove && HasUniformWorkGroupSize) {
     // Under v5  __ockl_get_local_size returns the value computed by the expression:
     //
     //   workgroup_id < hidden_block_count ? hidden_group_size : hidden_remainder
     //
     // For functions with the attribute uniform-work-group-size=true. we can evaluate
     // workgroup_id < hidden_block_count as true, and thus hidden_group_size is returned
     // for __ockl_get_local_size.
     for (int I = 0; I < 3; ++I) {
-      Value *BlockCount = BlockCounts[I];
-      if (!BlockCount)
-        continue;
-
       using namespace llvm::PatternMatch;
       auto GroupIDIntrin =
           I == 0 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>()
                  : (I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>()
                            : m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());
 
+      for (Value *BlockCount : BlockCounts[I]) {
         for (User *ICmp : BlockCount->users()) {
           ICmpInst::Predicate Pred;
-        if (match(ICmp, m_ICmp(Pred, GroupIDIntrin, m_Specific(BlockCount)))) {
+          if (match(ICmp,
+                    m_ICmp(Pred, GroupIDIntrin, m_Specific(BlockCount)))) {
             if (Pred != ICmpInst::ICMP_ULT)
               continue;
-          ICmp->replaceAllUsesWith(llvm::ConstantInt::getTrue(ICmp->getType()));
+            ICmp->replaceAllUsesWith(
+                llvm::ConstantInt::getTrue(ICmp->getType()));
             MadeChange = true;
           }
         }
       }
+    }
 
     // All remainders should be 0 with uniform work group size.
     for (Value *Remainder : Remainders) {
-      if (!Remainder)
-        continue;
       Remainder->replaceAllUsesWith(Constant::getNullValue(Remainder->getType()));
       MadeChange = true;
     }
   } else if (HasUniformWorkGroupSize) { // Pre-V5.
     // Pattern match the code used to handle partial workgroup dispatches in the
     // library implementation of get_local_size, so the entire function can be
     // constant folded with a known group size.
     //
     // uint r = grid_size - group_id * group_size;
     // get_local_size = (r < group_size) ? r : group_size;
     //
     // If we have uniform-work-group-size (which is the default in OpenCL 1.2),
     // the grid_size is required to be a multiple of group_size). In this case:
     //
     // grid_size - (group_id * group_size) < group_size
     // ->
     // grid_size < group_size + (group_id * group_size)
     //
     // (grid_size / group_size) < 1 + group_id
     //
     // grid_size / group_size is at least 1, so we can conclude the select
     // condition is false (except for group_id == 0, where the select result is
     // the same).
     for (int I = 0; I < 3; ++I) {
-      Value *GroupSize = GroupSizes[I];
-      Value *GridSize = GridSizes[I];
-      if (!GroupSize || !GridSize)
-        continue;
-
       using namespace llvm::PatternMatch;
       auto GroupIDIntrin =
           I == 0 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>()
                  : (I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>()
                            : m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());
+      auto MatchGroupSizesZExt = m_ZExt(m_MatchAnyOf(GroupSizes[I]));
 
-      for (User *U : GroupSize->users()) {
-        auto *ZextGroupSize = dyn_cast<ZExtInst>(U);
-        if (!ZextGroupSize)
+      for (Value *GridSize : GridSizes[I]) {
+        if (GroupSizes[I].empty())
           continue;
 
-        for (User *UMin : ZextGroupSize->users()) {
-          if (match(UMin,
-                    m_UMin(m_Sub(m_Specific(GridSize),
-                                 m_Mul(GroupIDIntrin, m_Specific(ZextGroupSize))),
-                           m_Specific(ZextGroupSize)))) {
+        for (User *Sub : GridSize->users()) {
+          Instruction *ZExtGroupSize;
+          auto MatchAndBindGroupSizesZExt =
+              m_CombineAnd(MatchGroupSizesZExt, m_Instruction(ZExtGroupSize));
+          if (!match(Sub,
+                     m_Sub(m_Specific(GridSize),
+                           m_Mul(GroupIDIntrin, MatchAndBindGroupSizesZExt))))
+            continue;
+          for (User *UMin : Sub->users()) {
+            if (match(UMin, m_UMin(m_Specific(Sub), MatchGroupSizesZExt))) {
               if (HasReqdWorkGroupSize) {
-              ConstantInt *KnownSize
-                = mdconst::extract<ConstantInt>(MD->getOperand(I));
+                ConstantInt *KnownSize =
+                    mdconst::extract<ConstantInt>(MD->getOperand(I));
                 UMin->replaceAllUsesWith(ConstantExpr::getIntegerCast(
                     KnownSize, UMin->getType(), false));
               } else {
-              UMin->replaceAllUsesWith(ZextGroupSize);
+                UMin->replaceAllUsesWith(ZExtGroupSize);
               }
 
               MadeChange = true;
             }
           }
         }
       }
     }
+  }
 
   // If reqd_work_group_size is set, we can replace work group size with it.
   if (!HasReqdWorkGroupSize)
     return MadeChange;
 
   for (int I = 0; I < 3; I++) {
-    Value *GroupSize = GroupSizes[I];
-    if (!GroupSize)
-      continue;
-
+    for (Value *GroupSize : GroupSizes[I]) {
       ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(I));
       GroupSize->replaceAllUsesWith(
           ConstantExpr::getIntegerCast(KnownSize, GroupSize->getType(), false));
       MadeChange = true;
     }
+  }
 
   return MadeChange;
 }
 
+} // end anonymous namespace
 
 // TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
 // TargetPassConfig for subtarget.
 bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
   bool MadeChange = false;
   bool IsV5OrAbove = AMDGPU::getAmdhsaCodeObjectVersion() >= 5;
   Function *BasePtr = getBasePtrIntrinsic(M, IsV5OrAbove);
 

llvm/test/CodeGen/AMDGPU/implicit-arg-v5-opt-opaque-ptr.ll

;
   %cmp.id.count = icmp ult i32 %group.id, %block.count.z
   %local.size.offset = select i1 %cmp.id.count, i64 16, i64 22
   %gep.local.size = getelementptr inbounds i8, ptr addrspace(4) %implicitarg.ptr, i64 %local.size.offset
   %local.size = load i16, ptr addrspace(4) %gep.local.size, align 2
   store i16 %local.size, i16 addrspace(1)* %out
   ret void
 }
 
+; Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
+define amdgpu_kernel void @get_local_size_x_2_block_count_users_opaque_pointer(i16 addrspace(1)* %out) #0 {
+; GCN-LABEL: @get_local_size_x_2_block_count_users_opaque_pointer(
+; GCN-NEXT:    [[IMPLICITARG_PTR:%.*]] = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; GCN-NEXT:    [[GEP_LOCAL_SIZE:%.*]] = getelementptr inbounds i8, ptr addrspace(4) [[IMPLICITARG_PTR]], i64 12
+; GCN-NEXT:    [[LOCAL_SIZE:%.*]] = load i16, ptr addrspace(4) [[GEP_LOCAL_SIZE]], align 4
+; GCN-NEXT:    store volatile i16 [[LOCAL_SIZE]], ptr addrspace(1) [[OUT:%.*]], align 2
+; GCN-NEXT:    store volatile i16 [[LOCAL_SIZE]], ptr addrspace(1) [[OUT:%.*]], align 2
+; GCN-NEXT:    ret void
+;
+  %group.id = tail call i32 @llvm.amdgcn.workgroup.id.x()
+  %implicitarg.ptr = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %block.count.x = load i32, ptr addrspace(4) %implicitarg.ptr, align 4
+  %cmp.id.count = icmp ult i32 %group.id, %block.count.x
+  %local.size.offset = select i1 %cmp.id.count, i64 12, i64 18
+  %gep.local.size = getelementptr inbounds i8, ptr addrspace(4) %implicitarg.ptr, i64 %local.size.offset
+  %local.size = load i16, ptr addrspace(4) %gep.local.size, align 2
+  store volatile i16 %local.size, i16 addrspace(1)* %out
+  %cmp.id.count.1 = icmp ult i32 %group.id, %block.count.x
+  %local.size.offset.1 = select i1 %cmp.id.count.1, i64 12, i64 18
+  %gep.local.size.1 = getelementptr inbounds i8, ptr addrspace(4) %implicitarg.ptr, i64 %local.size.offset.1
+  %local.size.1 = load i16, ptr addrspace(4) %gep.local.size.1, align 2
+  store volatile i16 %local.size.1, i16 addrspace(1)* %out
+  ret void
+}
+
 declare ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr() #1
 declare i32 @llvm.amdgcn.workgroup.id.x() #1
 declare i32 @llvm.amdgcn.workgroup.id.y() #1
 declare i32 @llvm.amdgcn.workgroup.id.z() #1
 
 !llvm.module.flags = !{!1}
 
 attributes #0 = { nounwind "uniform-work-group-size"="true" }
 attributes #1 = { nounwind readnone speculatable }
 !0 = !{i32 8, i32 16, i32 2}
 !1 = !{i32 1, !"amdgpu_code_object_version", i32 500}

llvm/test/CodeGen/AMDGPU/implicit-arg-v5-opt.ll

 
 ; Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
 define amdgpu_kernel void @get_remainder_y(i16 addrspace(1)* %out) #0 {
 ; GCN-LABEL: @get_remainder_y(
 ; GCN-NEXT:    store i16 0, i16 addrspace(1)* [[OUT:%.*]], align 2
 ; GCN-NEXT:    ret void
 ;
   %implicitarg.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
-  %gep.y = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 18
+  %gep.y = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 20
   %bc.y = bitcast i8 addrspace(4)* %gep.y to i16 addrspace(4)*
   %remainder.y = load i16, i16 addrspace(4)* %bc.y, align 2
   store i16 %remainder.y, i16 addrspace(1)* %out
   ret void
 }
 
 ; Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
 define amdgpu_kernel void @get_remainder_z(i16 addrspace(1)* %out) #0 {
 ; GCN-LABEL: @get_remainder_z(
 ; GCN-NEXT:    store i16 0, i16 addrspace(1)* [[OUT:%.*]], align 2
 ; GCN-NEXT:    ret void
 ;
   %implicitarg.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
-  %gep.z = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 18
+  %gep.z = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 22
   %bc.z = bitcast i8 addrspace(4)* %gep.z to i16 addrspace(4)*
   %remainder.z = load i16, i16 addrspace(4)* %bc.z, align 2
   store i16 %remainder.z, i16 addrspace(1)* %out
   ret void
 }
 
 ; Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
+define amdgpu_kernel void @get_remainder_xyz(i16 addrspace(1)* %out) #0 {
+; GCN-LABEL: @get_remainder_xyz(
+; GCN-NEXT:    store volatile i16 0, i16 addrspace(1)* [[OUT:%.*]], align 2
+; GCN-NEXT:    store volatile i16 0, i16 addrspace(1)* [[OUT:%.*]], align 2
+; GCN-NEXT:    store volatile i16 0, i16 addrspace(1)* [[OUT:%.*]], align 2
+; GCN-NEXT:    ret void
+;
+  %implicitarg.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
+  %gep.x = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 18
+  %bc.x = bitcast i8 addrspace(4)* %gep.x to i16 addrspace(4)*
+  %remainder.x = load i16, i16 addrspace(4)* %bc.x, align 2
+  store volatile i16 %remainder.x, i16 addrspace(1)* %out
+  %gep.y = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 20
+  %bc.y = bitcast i8 addrspace(4)* %gep.y to i16 addrspace(4)*
+  %remainder.y = load i16, i16 addrspace(4)* %bc.y, align 2
+  store volatile i16 %remainder.y, i16 addrspace(1)* %out
+  %gep.z = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 22
+  %bc.z = bitcast i8 addrspace(4)* %gep.z to i16 addrspace(4)*
+  %remainder.z = load i16, i16 addrspace(4)* %bc.z, align 2
+  store volatile i16 %remainder.z, i16 addrspace(1)* %out
+  ret void
+}
+
+; Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
 define amdgpu_kernel void @get_work_group_size_x(i16 addrspace(1)* %out) #0 {
 ; GCN-LABEL: @get_work_group_size_x(
 ; GCN-NEXT:    [[IMPLICITARG_PTR:%.*]] = tail call i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
 ; GCN-NEXT:    [[GEP_X:%.*]] = getelementptr inbounds i8, i8 addrspace(4)* [[IMPLICITARG_PTR]], i64 12
 ; GCN-NEXT:    [[BC_X:%.*]] = bitcast i8 addrspace(4)* [[GEP_X]] to i16 addrspace(4)*
 ; GCN-NEXT:    [[GROUP_SIZE_X:%.*]] = load i16, i16 addrspace(4)* [[BC_X]], align 4
 ; GCN-NEXT:    store i16 [[GROUP_SIZE_X]], i16 addrspace(1)* [[OUT:%.*]], align 2
 ; GCN-NEXT:    ret void
;
   %implicitarg.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
   %gep.z = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 16
   %bc.z = bitcast i8 addrspace(4)* %gep.z to i16 addrspace(4)*
   %group.size.z = load i16, i16 addrspace(4)* %bc.z, align 2
   store i16 %group.size.z, i16 addrspace(1)* %out
   ret void
 }
 
+define amdgpu_kernel void @get_all_work_group_sizes_reqd(i16 addrspace(1)* %out) #0 !reqd_work_group_size !0 {
+; GCN-LABEL: @get_all_work_group_sizes_reqd(
+; GCN-NEXT:    store volatile i16 8, i16 addrspace(1)* [[OUT:%.*]], align 2
+; GCN-NEXT:    store volatile i16 16, i16 addrspace(1)* [[OUT:%.*]], align 2
+; GCN-NEXT:    store volatile i16 2, i16 addrspace(1)* [[OUT:%.*]], align 2
+; GCN-NEXT:    ret void
+;
+  %implicitarg.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
+  %gep.x = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 12
+  %gep.y = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 14
+  %gep.z = getelementptr inbounds i8, i8 addrspace(4)* %implicitarg.ptr, i64 16
+  %bc.x = bitcast i8 addrspace(4)* %gep.x to i16 addrspace(4)*
+  %bc.y.0 = bitcast i8 addrspace(4)* %gep.y to i16 addrspace(4)*
+  %bc.y = bitcast i16 addrspace(4)* %bc.y.0 to i16 addrspace(4)*
+  %bc.z = bitcast i8 addrspace(4)* %gep.z to i16 addrspace(4)*
+  %bc.z.select = select i1 true, i16 addrspace(4)* %bc.z, i16 addrspace(4)* %bc.z
+  %group.size.x = load i16, i16 addrspace(4)* %bc.x, align 2
+  %group.size.y = load i16, i16 addrspace(4)* %bc.y, align 2
+  %group.size.z = load i16, i16 addrspace(4)* %bc.z.select, align 2
+  store volatile i16 %group.size.x, i16 addrspace(1)* %out
+  store volatile i16 %group.size.y, i16 addrspace(1)* %out
+  store volatile i16 %group.size.z, i16 addrspace(1)* %out
+  ret void
+}
 
 declare i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr() #1
 declare i32 @llvm.amdgcn.workgroup.id.x() #1
 declare i32 @llvm.amdgcn.workgroup.id.y() #1
 declare i32 @llvm.amdgcn.workgroup.id.z() #1
 
 !llvm.module.flags = !{!1}
 
 attributes #0 = { nounwind "uniform-work-group-size"="true" }
 attributes #1 = { nounwind readnone speculatable }
 !0 = !{i32 8, i32 16, i32 2}
 !1 = !{i32 1, !"amdgpu_code_object_version", i32 500}

llvm/test/CodeGen/AMDGPU/reqd-work-group-size.ll

define amdgpu_kernel void @all_local_size(i64 addrspace(1)* nocapture readnone %out) #0 !reqd_work_group_size !0 {
   %umin2 = call i32 @llvm.umin.i32(i32 %tmp31.i4, i32 %tmp29.i2)
   %tmp34.i7 = zext i32 %umin2 to i64
   store volatile i64 %tmp34.i, i64 addrspace(1)* %out, align 4
   store volatile i64 %tmp34.i14, i64 addrspace(1)* %out, align 4
   store volatile i64 %tmp34.i7, i64 addrspace(1)* %out, align 4
   ret void
 }
 
+; CHECK-LABEL: @all_local_size_twice(
+; CHECK-NEXT: store volatile i64 8, i64 addrspace(1)* %out, align 4
+; CHECK-NEXT: store volatile i64 16, i64 addrspace(1)* %out, align 4
+; CHECK-NEXT: store volatile i64 2, i64 addrspace(1)* %out, align 4
+; CHECK-NEXT: store volatile i64 8, i64 addrspace(1)* %out, align 4
+; CHECK-NEXT: store volatile i64 16, i64 addrspace(1)* %out, align 4
+; CHECK-NEXT: store volatile i64 2, i64 addrspace(1)* %out, align 4
+define amdgpu_kernel void @all_local_size_twice(i64 addrspace(1)* nocapture readnone %out) #0 !reqd_work_group_size !0 {
+  %tmp.i = tail call i8 addrspace(4)* @llvm.amdgcn.dispatch.ptr() #0
+  %tmp2.i = tail call i32 @llvm.amdgcn.workgroup.id.x() #0
+  %tmp3.i = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 12
+  %tmp4.i = bitcast i8 addrspace(4)* %tmp3.i to i32 addrspace(4)*
+  %tmp5.i = load i32, i32 addrspace(4)* %tmp4.i, align 4
+  %tmp6.i = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 4
+  %tmp7.i = bitcast i8 addrspace(4)* %tmp6.i to i16 addrspace(4)*
+  %tmp8.i = load i16, i16 addrspace(4)* %tmp7.i, align 4
+  %tmp29.i = zext i16 %tmp8.i to i32
+  %tmp30.i = mul i32 %tmp2.i, %tmp29.i
+  %tmp31.i = sub i32 %tmp5.i, %tmp30.i
+  %umin0 = call i32 @llvm.umin.i32(i32 %tmp31.i, i32 %tmp29.i)
+  %tmp34.i = zext i32 %umin0 to i64
+  %tmp10.i = tail call i32 @llvm.amdgcn.workgroup.id.y() #0
+  %tmp11.i = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 16
+  %tmp12.i = bitcast i8 addrspace(4)* %tmp11.i to i32 addrspace(4)*
+  %tmp13.i = load i32, i32 addrspace(4)* %tmp12.i, align 8
+  %tmp14.i = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 6
+  %tmp15.i = bitcast i8 addrspace(4)* %tmp14.i to i16 addrspace(4)*
+  %tmp16.i = load i16, i16 addrspace(4)* %tmp15.i, align 2
+  %tmp29.i9 = zext i16 %tmp16.i to i32
+  %tmp30.i10 = mul i32 %tmp10.i, %tmp29.i9
+  %tmp31.i11 = sub i32 %tmp13.i, %tmp30.i10
+  %umin1 = call i32 @llvm.umin.i32(i32 %tmp31.i11, i32 %tmp29.i9)
+  %tmp34.i14 = zext i32 %umin1 to i64
+  %tmp18.i = tail call i32 @llvm.amdgcn.workgroup.id.z() #0
+  %tmp19.i = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 20
+  %tmp20.i = bitcast i8 addrspace(4)* %tmp19.i to i32 addrspace(4)*
+  %tmp21.i = load i32, i32 addrspace(4)* %tmp20.i, align 4
+  %tmp22.i = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 8
+  %tmp23.i = bitcast i8 addrspace(4)* %tmp22.i to i16 addrspace(4)*
+  %tmp24.i = load i16, i16 addrspace(4)* %tmp23.i, align 8
+  %tmp29.i2 = zext i16 %tmp24.i to i32
+  %tmp30.i3 = mul i32 %tmp18.i, %tmp29.i2
+  %tmp31.i4 = sub i32 %tmp21.i, %tmp30.i3
+  %umin2 = call i32 @llvm.umin.i32(i32 %tmp31.i4, i32 %tmp29.i2)
+  %tmp34.i7 = zext i32 %umin2 to i64
+  %tmp2.i.1 = tail call i32 @llvm.amdgcn.workgroup.id.x() #0
+  %tmp3.i.1 = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 12
+  %tmp4.i.1 = bitcast i8 addrspace(4)* %tmp3.i.1 to i32 addrspace(4)*
+  %tmp5.i.1 = load i32, i32 addrspace(4)* %tmp4.i.1, align 4
+  %tmp6.i.1 = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 4
+  %tmp7.i.1 = bitcast i8 addrspace(4)* %tmp6.i.1 to i16 addrspace(4)*
+  %tmp8.i.1 = load i16, i16 addrspace(4)* %tmp7.i.1, align 4
+  %tmp29.i.1 = zext i16 %tmp8.i.1 to i32
+  %tmp30.i.1 = mul i32 %tmp2.i.1, %tmp29.i.1
+  %tmp31.i.1 = sub i32 %tmp5.i.1, %tmp30.i.1
+  %umin0.1 = call i32 @llvm.umin.i32(i32 %tmp31.i.1, i32 %tmp29.i.1)
+  %tmp34.i.1 = zext i32 %umin0.1 to i64
+  %tmp10.i.1 = tail call i32 @llvm.amdgcn.workgroup.id.y() #0
+  %tmp11.i.1 = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 16
+  %tmp12.i.1 = bitcast i8 addrspace(4)* %tmp11.i.1 to i32 addrspace(4)*
+  %tmp13.i.1 = load i32, i32 addrspace(4)* %tmp12.i.1, align 8
+  %tmp14.i.1 = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 6
+  %tmp15.i.1 = bitcast i8 addrspace(4)* %tmp14.i.1 to i16 addrspace(4)*
+  %tmp16.i.1 = load i16, i16 addrspace(4)* %tmp15.i.1, align 2
+  %tmp29.i9.1 = zext i16 %tmp16.i.1 to i32
+  %tmp30.i10.1 = mul i32 %tmp10.i.1, %tmp29.i9.1
+  %tmp31.i11.1 = sub i32 %tmp13.i.1, %tmp30.i10.1
+  %umin1.1 = call i32 @llvm.umin.i32(i32 %tmp31.i11.1, i32 %tmp29.i9) ; use tmp29.i9 instead of tmp29.i9.1
+  %tmp34.i14.1 = zext i32 %umin1.1 to i64
+  %tmp18.i.1 = tail call i32 @llvm.amdgcn.workgroup.id.z() #0
+  %tmp19.i.1 = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 20
+  %tmp20.i.1 = bitcast i8 addrspace(4)* %tmp19.i.1 to i32 addrspace(4)*
+  %tmp21.i.1 = load i32, i32 addrspace(4)* %tmp20.i.1, align 4
+  %tmp22.i.1 = getelementptr inbounds i8, i8 addrspace(4)* %tmp.i, i64 8
+  %tmp23.i.1 = bitcast i8 addrspace(4)* %tmp22.i.1 to i16 addrspace(4)*
+  %tmp24.i.1 = load i16, i16 addrspace(4)* %tmp23.i.1, align 8
+  %tmp29.i2.1 = zext i16 %tmp24.i.1 to i32
+  %tmp30.i3.1 = mul i32 %tmp18.i.1, %tmp29.i2.1
+  %tmp31.i4.1 = sub i32 %tmp21.i.1, %tmp30.i3.1
+  %umin2.1 = call i32 @llvm.umin.i32(i32 %tmp31.i4.1, i32 %tmp29.i2.1)
+  %tmp34.i7.1 = zext i32 %umin2.1 to i64
+  store volatile i64 %tmp34.i, i64 addrspace(1)* %out, align 4
+  store volatile i64 %tmp34.i14, i64 addrspace(1)* %out, align 4
+  store volatile i64 %tmp34.i7, i64 addrspace(1)* %out, align 4
+  store volatile i64 %tmp34.i.1, i64 addrspace(1)* %out, align 4
+  store volatile i64 %tmp34.i14.1, i64 addrspace(1)* %out, align 4
+  store volatile i64 %tmp34.i7.1, i64 addrspace(1)* %out, align 4
+  ret void
+}
+
 ; TODO: Should be able to handle this, but not much reason to.
 ; CHECK-LABEL: @partial_load_group_size_x(
 ; CHECK-NEXT: %dispatch.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.dispatch.ptr()
 ; CHECK-NEXT: %gep.group.size.x = getelementptr inbounds i8, i8 addrspace(4)* %dispatch.ptr, i64 4
 ; CHECK-NEXT: %group.size.x.lo = load i8, i8 addrspace(4)* %gep.group.size.x, align 4
 ; CHECK-NEXT: store i8 %group.size.x.lo, i8 addrspace(1)* %out, align 1
 define amdgpu_kernel void @partial_load_group_size_x(i8 addrspace(1)* %out) #0 !reqd_work_group_size !0 {
   %dispatch.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.dispatch.ptr()