[llvm] [AMDGPU]: Rewrite mbcnt_lo/mbcnt_hi to work item ID where applicable (PR #160496)

Teja Alaghari via llvm-commits llvm-commits at lists.llvm.org
Fri Oct 31 00:01:45 PDT 2025


https://github.com/TejaX-Alaghari updated https://github.com/llvm/llvm-project/pull/160496

>From 989aee6fa2f419a95c1fa64619b9671a04815035 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Mon, 29 Sep 2025 13:48:07 +0530
Subject: [PATCH 01/11] [AMDGPU] Move mbcnt optimization from InstCombine to
 AMDGPUCodeGenPrepare

This addresses reviewer concerns about pipeline timing by moving the mbcnt
optimization from InstCombine to AMDGPUCodeGenPrepare. The InstCombine pass
runs before AMDGPU Attributor, which means reqd_work_group_size metadata
may not be available. AMDGPUCodeGenPrepare runs later in the pipeline after
the attributor pass, ensuring proper metadata availability.

Changes:
- Move visitMbcntLo and visitMbcntHi methods to AMDGPUCodeGenPrepare
- Remove complex mbcnt optimization from AMDGPUInstCombineIntrinsic
- Keep simple wave32 mbcnt_hi -> copy optimization in InstCombine
- Move test files from InstCombine/AMDGPU to Transforms/AMDGPU
- Update test RUN lines to use amdgpu-codegenprepare pass

This fixes the pipeline ordering issue where InstCombine runs before
AMDGPU Attributor, preventing the optimization from triggering when
reqd_work_group_size metadata is set by the attributor.
---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    | 147 ++++++++++++++++++
 .../Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll |  25 +++
 .../AMDGPU/mbcnt-to-bitmask-posit.ll          |  26 ++++
 .../AMDGPU/mbcnt-to-workitem-neg.ll           |  20 +++
 .../AMDGPU/mbcnt-to-workitem-posit.ll         |  26 ++++
 .../AMDGPU/mbcnt-to-workitem-wave32-neg.ll    |  66 ++++++++
 .../AMDGPU/mbcnt-to-workitem-wave32.ll        |  25 +++
 7 files changed, 335 insertions(+)
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 8e35ba77d69aa..4f647e34b9f42 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -260,6 +260,8 @@ class AMDGPUCodeGenPrepareImpl
   bool visitIntrinsicInst(IntrinsicInst &I);
   bool visitFMinLike(IntrinsicInst &I);
   bool visitSqrt(IntrinsicInst &I);
+  bool visitMbcntLo(IntrinsicInst &I);
+  bool visitMbcntHi(IntrinsicInst &I);
   bool run();
 };
 
@@ -1910,6 +1912,10 @@ bool AMDGPUCodeGenPrepareImpl::visitIntrinsicInst(IntrinsicInst &I) {
     return visitFMinLike(I);
   case Intrinsic::sqrt:
     return visitSqrt(I);
+  case Intrinsic::amdgcn_mbcnt_lo:
+    return visitMbcntLo(I);
+  case Intrinsic::amdgcn_mbcnt_hi:
+    return visitMbcntHi(I);
   default:
     return false;
   }
@@ -2090,6 +2096,147 @@ INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
 INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)
 
+bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
+  // On wave32 targets, mbcnt.lo(~0, 0) can be replaced with workitem.id.x
+  if (!ST.isWave32())
+    return false;
+
+  // Check for pattern mbcnt.lo(~0, 0)
+  auto *Arg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
+  auto *Arg1C = dyn_cast<ConstantInt>(I.getArgOperand(1));
+  if (!Arg0C || !Arg1C || !Arg0C->isAllOnesValue() || !Arg1C->isZero())
+    return false;
+
+  // Check reqd_work_group_size similar to mbcnt_hi case
+  Function *F = I.getFunction();
+  if (!F)
+    return false;
+
+  unsigned Wave = 0;
+  if (ST.isWaveSizeKnown())
+    Wave = ST.getWavefrontSize();
+
+  if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
+    unsigned XLen = *MaybeX;
+    if (Wave == 0 && XLen == 32)
+      Wave = XLen;
+
+    if (Wave != 0 && XLen == Wave) {
+      IRBuilder<> B(&I);
+      CallInst *NewCall =
+          B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+      NewCall->takeName(&I);
+      ST.makeLIDRangeMetadata(NewCall);
+      I.replaceAllUsesWith(NewCall);
+      I.eraseFromParent();
+      return true;
+    }
+  }
+
+  return false;
+}
+
+bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
+  // exec_hi is all 0, so this is just a copy on wave32.
+  if (ST.isWave32()) {
+    I.replaceAllUsesWith(I.getArgOperand(1));
+    I.eraseFromParent();
+    return true;
+  }
+
+  // Pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0))
+  auto *HiArg1 = dyn_cast<CallInst>(I.getArgOperand(1));
+  if (!HiArg1)
+    return false;
+
+  Function *CalledF = HiArg1->getCalledFunction();
+  if (!CalledF || CalledF->getIntrinsicID() != Intrinsic::amdgcn_mbcnt_lo)
+    return false;
+
+  // hi arg0 must be all-ones
+  auto *HiArg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
+  if (!HiArg0C || !HiArg0C->isAllOnesValue())
+    return false;
+
+  // lo args: arg0 == ~0, arg1 == 0
+  Value *Lo0 = HiArg1->getArgOperand(0);
+  Value *Lo1 = HiArg1->getArgOperand(1);
+  auto *Lo0C = dyn_cast<ConstantInt>(Lo0);
+  auto *Lo1C = dyn_cast<ConstantInt>(Lo1);
+  if (!Lo0C || !Lo1C || !Lo0C->isAllOnesValue() || !Lo1C->isZero())
+    return false;
+
+  // Query reqd_work_group_size via subtarget helper and compare X to wave
+  // size conservatively.
+  Function *F = I.getFunction();
+  if (!F)
+    return false;
+
+  unsigned Wave = 0;
+  if (ST.isWaveSizeKnown())
+    Wave = ST.getWavefrontSize();
+
+  if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
+    unsigned XLen = *MaybeX;
+    if (Wave == 0 && (XLen == 32 || XLen == 64))
+      Wave = XLen; // allow common sizes under test harness
+
+    if (Wave != 0 && XLen == Wave) {
+      IRBuilder<> B(&I);
+      CallInst *NewCall =
+          B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+      NewCall->takeName(&I);
+      // Attach range metadata when available.
+      ST.makeLIDRangeMetadata(NewCall);
+      I.replaceAllUsesWith(NewCall);
+      I.eraseFromParent();
+      return true;
+    }
+    // Optional: if X dimension evenly splits into wavefronts we can
+    // replace lane-id computation with a bitmask when the wave is a
+    // power-of-two. Use the Subtarget helper to conservatively decide
+    // when per-wave tiling is preserved.
+    if (ST.hasWavefrontsEvenlySplittingXDim(*F, /*RequiresUniformYZ=*/true)) {
+      if (Wave != 0 && isPowerOf2_32(Wave)) {
+        // Construct: tid = workitem.id.x(); mask = Wave-1; res = tid & mask
+        IRBuilder<> B(&I);
+        CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+        Tid->takeName(&I);
+        IntegerType *ITy = cast<IntegerType>(Tid->getType());
+        Constant *Mask = ConstantInt::get(ITy, Wave - 1);
+        Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
+        AndInst->takeName(&I);
+        // Attach range metadata for the result if possible.
+        ST.makeLIDRangeMetadata(AndInst);
+        I.replaceAllUsesWith(AndInst);
+        I.eraseFromParent();
+        return true;
+      }
+    }
+  } else {
+    // No reqd_work_group_size metadata: be conservative and only handle the
+    // common test harness cases where reqd_work_group_size metadata exists
+    // and equals 32/64.
+    if (auto *Node = F->getMetadata("reqd_work_group_size")) {
+      if (Node->getNumOperands() == 3) {
+        unsigned XLen =
+            mdconst::extract<ConstantInt>(Node->getOperand(0))->getZExtValue();
+        if (XLen == 32 || XLen == 64) {
+          IRBuilder<> B(&I);
+          CallInst *NewCall =
+              B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+          NewCall->takeName(&I);
+          I.replaceAllUsesWith(NewCall);
+          I.eraseFromParent();
+          return true;
+        }
+      }
+    }
+  }
+
+  return false;
+}
+
 char AMDGPUCodeGenPrepare::ID = 0;
 
 FunctionPass *llvm::createAMDGPUCodeGenPreparePass() {
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll
new file mode 100644
index 0000000000000..470751c3c73f3
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll
@@ -0,0 +1,25 @@
+;; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %sNOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple=amdgcn-unknown-amdhsa -passes=instcombine < %s | FileCheck %s
+
+define i32 @test_mbcnt_non_wave_size() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_non_wave_size(
+; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+!1 = !{i32 48, i32 1, i32 1}
+
+; Declarations
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
+;.
+; CHECK: [[META0]] = !{i32 48, i32 1, i32 1}
+;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
new file mode 100644
index 0000000000000..97697d4d9651f
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+define i32 @test_mbcnt_wave64_to_workitem() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_wave64_to_workitem(
+; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[B:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+!1 = !{i32 64, i32 1, i32 1}
+
+; Declarations
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
+declare i32 @llvm.amdgcn.workitem.id.x()
+;.
+; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
+;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll
new file mode 100644
index 0000000000000..af8d713b798ed
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+define i32 @test_mbcnt_no_reqd_work_group_size() {
+; CHECK-LABEL: define i32 @test_mbcnt_no_reqd_work_group_size() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Declarations
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
+declare i32 @llvm.amdgcn.workitem.id.x()
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
new file mode 100644
index 0000000000000..91aa942df9337
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+define i32 @test_mbcnt_to_workitem() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_to_workitem(
+; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[B:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+!0 = !{i32 64, i32 1, i32 1}
+
+; Declarations
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
+declare i32 @llvm.amdgcn.workitem.id.x()
+;.
+; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
+;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
new file mode 100644
index 0000000000000..403ea7c361250
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
@@ -0,0 +1,66 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+; Test that mbcnt.lo(~0, 0) is NOT optimized on wave32 when work group size doesn't match wave size
+define i32 @test_mbcnt_lo_wave32_non_matching_wgs() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_non_matching_wgs(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; Test that mbcnt.lo(~0, 0) is NOT optimized on wave32 when no reqd_work_group_size is specified
+define i32 @test_mbcnt_lo_wave32_no_wgs() {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_no_wgs(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; Test that mbcnt.lo with non-all-ones first arg is NOT optimized
+define i32 @test_mbcnt_lo_wave32_partial_mask() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_partial_mask(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 15, i32 0)
+; CHECK-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 15, i32 0)
+  ret i32 %a
+}
+
+; Test that mbcnt.lo with non-zero second arg is NOT optimized
+define i32 @test_mbcnt_lo_wave32_non_zero_base() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_non_zero_base(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
+; CHECK-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
+  ret i32 %a
+}
+
+!0 = !{i32 48, i32 1, i32 1}  ; Work group size 48 != wave size 32
+!1 = !{i32 32, i32 1, i32 1}  ; Work group size 32 == wave size 32
+
+; Function Attrs: nounwind readnone speculatable willreturn
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
+
+attributes #0 = { nounwind readnone speculatable willreturn }
+;.
+; CHECK: [[META0]] = !{i32 48, i32 1, i32 1}
+; CHECK: [[META1]] = !{i32 32, i32 1, i32 1}
+;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll
new file mode 100644
index 0000000000000..07a5028ca1ee5
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll
@@ -0,0 +1,25 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+define i32 @test_mbcnt_lo_wave32() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 32) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+!0 = !{i32 32, i32 1, i32 1}
+
+; Function Attrs: nounwind readnone speculatable willreturn
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
+declare i32 @llvm.amdgcn.workitem.id.x() #0
+
+attributes #0 = { nounwind readnone speculatable willreturn }
+;.
+; CHECK: [[META0]] = !{i32 32, i32 1, i32 1}
+;.

>From 61036bdbbec44be73cbe1e60820a8415529bac3b Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Tue, 30 Sep 2025 11:43:42 +0530
Subject: [PATCH 02/11] [AMDGPU] Enhance mbcnt.lo optimization to support
 bitmask cases

Address reviewer feedback by extending visitMbcntLo to handle cases where
X dimension doesn't exactly match wave size but allows even wave distribution.

Key improvements:
- mbcnt.lo(~0, 0) -> workitem.id.x() & (wave_size-1) when X dimension
  allows wavefronts to evenly split using ST.hasWavefrontsEvenlySplittingXDim()
- Handles cases like X=48, X=64, X=128 on wave32 (previously unoptimized)
- Added comprehensive test coverage for bitmask optimization cases
- Updated existing test to reflect new correct behavior

This addresses krzysz00's comment about supporting masking optimizations
for non-exact wave size matches.
---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    | 17 +++++++
 .../Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll  | 44 +++++++++++++++++++
 .../AMDGPU/mbcnt-to-workitem-wave32-neg.ll    |  5 ++-
 3 files changed, 64 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 4f647e34b9f42..f95bf528dd94a 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -2131,6 +2131,23 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
       I.eraseFromParent();
       return true;
     }
+    // Handle bitmask case: when X dimension evenly splits into waves
+    // mbcnt.lo(~0, 0) = workitem.id.x() & (wave_size - 1)
+    if (ST.hasWavefrontsEvenlySplittingXDim(*F, /*RequiresUniformYZ=*/true)) {
+      if (Wave != 0 && isPowerOf2_32(Wave)) {
+        IRBuilder<> B(&I);
+        CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+        ST.makeLIDRangeMetadata(Tid);
+        IntegerType *ITy = cast<IntegerType>(Tid->getType());
+        Constant *Mask = ConstantInt::get(ITy, Wave - 1);
+        Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
+        AndInst->takeName(&I);
+        // Note: Range metadata cannot be applied to 'and' instructions
+        I.replaceAllUsesWith(AndInst);
+        I.eraseFromParent();
+        return true;
+      }
+    }
   }
 
   return false;
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll b/llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll
new file mode 100644
index 0000000000000..f84217705818e
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll
@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+; Test that mbcnt.lo(~0, 0) is optimized to workitem.id.x() & 0x1f on wave32 
+; when work group size is multiple of wave size (64 = 2 * 32)
+define i32 @test_mbcnt_lo_wave32_bitmask() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_bitmask(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call {{.*}} i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; Test with X dimension = 128 (4 * 32 waves)
+define i32 @test_mbcnt_lo_wave32_bitmask_128() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_bitmask_128(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call {{.*}} i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+!0 = !{i32 64, i32 1, i32 1}   ; 64 = 2 * 32 wave size
+!1 = !{i32 128, i32 1, i32 1}  ; 128 = 4 * 32 wave size
+
+; Function Attrs: nounwind readnone speculatable willreturn
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
+declare i32 @llvm.amdgcn.workitem.id.x() #0
+
+attributes #0 = { nounwind readnone speculatable willreturn }
+;.
+; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
+; CHECK: [[META1]] = !{i32 128, i32 1, i32 1}
+;.
\ No newline at end of file
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
index 403ea7c361250..30514562556ea 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
@@ -1,12 +1,13 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
 ; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
 
-; Test that mbcnt.lo(~0, 0) is NOT optimized on wave32 when work group size doesn't match wave size
+; Test that mbcnt.lo(~0, 0) IS optimized on wave32 with bitmask when work group size allows even wave distribution
 define i32 @test_mbcnt_lo_wave32_non_matching_wgs() !reqd_work_group_size !0 {
 ; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_non_matching_wgs(
 ; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call {{.*}} i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
 ; CHECK-NEXT:    ret i32 [[A]]
 ;
 entry:

>From dcbce2a20ab547a0dc7a128dec4fedc2ccb1545f Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Tue, 30 Sep 2025 14:57:34 +0530
Subject: [PATCH 03/11] Fix mbcnt.hi wave32 optimization to respect
 reqd_work_group_size metadata

The wave32 optimization for mbcnt.hi was unconditionally optimizing the
instruction on wave32 targets, which caused CodeGen test failures because
the tests expected to see the mbcnt.hi instruction in the output.

This change makes the wave32 mbcnt.hi optimization conditional on the same
metadata requirements as the mbcnt.lo optimization, ensuring consistency
and preserving existing test behavior for functions without the appropriate
metadata.
---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    | 23 ++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index f95bf528dd94a..6b0b60701ed49 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -2155,10 +2155,27 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
 
 bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   // exec_hi is all 0, so this is just a copy on wave32.
+  // However, only optimize if we have the same conditions as mbcnt.lo.
   if (ST.isWave32()) {
-    I.replaceAllUsesWith(I.getArgOperand(1));
-    I.eraseFromParent();
-    return true;
+    Function *F = I.getFunction();
+    if (!F)
+      return false;
+
+    unsigned Wave = 0;
+    if (ST.isWaveSizeKnown())
+      Wave = ST.getWavefrontSize();
+
+    if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
+      unsigned XLen = *MaybeX;
+      if (Wave == 0 && XLen == 32)
+        Wave = XLen;
+
+      if (Wave != 0 && XLen == Wave) {
+        I.replaceAllUsesWith(I.getArgOperand(1));
+        I.eraseFromParent();
+        return true;
+      }
+    }
   }
 
   // Pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0))

>From 26543cb450241db44c0767a1c1341e4ae37c79be Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Fri, 3 Oct 2025 11:04:32 +0530
Subject: [PATCH 04/11] Address review feedback: Add periods to all comments

As requested by @cdevadas, all comments in the mbcnt optimization
functions now end with a period for consistency with LLVM style.

Note: The use of 'auto' keyword was defended by @krzysz00 as these
appear next to dyn_cast<> which already includes the type information,
following common LLVM patterns.
---
 .../lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 6b0b60701ed49..ae5b49901e2b2 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -2097,17 +2097,17 @@ INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)
 
 bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
-  // On wave32 targets, mbcnt.lo(~0, 0) can be replaced with workitem.id.x
+  // On wave32 targets, mbcnt.lo(~0, 0) can be replaced with workitem.id.x.
   if (!ST.isWave32())
     return false;
 
-  // Check for pattern mbcnt.lo(~0, 0)
+  // Check for pattern mbcnt.lo(~0, 0).
   auto *Arg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
   auto *Arg1C = dyn_cast<ConstantInt>(I.getArgOperand(1));
   if (!Arg0C || !Arg1C || !Arg0C->isAllOnesValue() || !Arg1C->isZero())
     return false;
 
-  // Check reqd_work_group_size similar to mbcnt_hi case
+  // Check reqd_work_group_size similar to mbcnt_hi case.
   Function *F = I.getFunction();
   if (!F)
     return false;
@@ -2131,8 +2131,8 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
       I.eraseFromParent();
       return true;
     }
-    // Handle bitmask case: when X dimension evenly splits into waves
-    // mbcnt.lo(~0, 0) = workitem.id.x() & (wave_size - 1)
+    // Handle bitmask case: when X dimension evenly splits into waves.
+    // mbcnt.lo(~0, 0) = workitem.id.x() & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(*F, /*RequiresUniformYZ=*/true)) {
       if (Wave != 0 && isPowerOf2_32(Wave)) {
         IRBuilder<> B(&I);
@@ -2142,7 +2142,7 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
         Constant *Mask = ConstantInt::get(ITy, Wave - 1);
         Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
         AndInst->takeName(&I);
-        // Note: Range metadata cannot be applied to 'and' instructions
+        // Note: Range metadata cannot be applied to 'and' instructions.
         I.replaceAllUsesWith(AndInst);
         I.eraseFromParent();
         return true;
@@ -2178,7 +2178,7 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
     }
   }
 
-  // Pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0))
+  // Pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0)).
   auto *HiArg1 = dyn_cast<CallInst>(I.getArgOperand(1));
   if (!HiArg1)
     return false;
@@ -2187,12 +2187,12 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   if (!CalledF || CalledF->getIntrinsicID() != Intrinsic::amdgcn_mbcnt_lo)
     return false;
 
-  // hi arg0 must be all-ones
+  // hi arg0 must be all-ones.
   auto *HiArg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
   if (!HiArg0C || !HiArg0C->isAllOnesValue())
     return false;
 
-  // lo args: arg0 == ~0, arg1 == 0
+  // lo args: arg0 == ~0, arg1 == 0.
   Value *Lo0 = HiArg1->getArgOperand(0);
   Value *Lo1 = HiArg1->getArgOperand(1);
   auto *Lo0C = dyn_cast<ConstantInt>(Lo0);

>From 6471dbb16a748fa14c24ad97238cfc402aabf0d0 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Fri, 3 Oct 2025 12:46:35 +0530
Subject: [PATCH 05/11] Address review comments from @arsenm

- Make visitMbcnt* functions const
- Use class member F instead of getFunction() calls
- Use auto* for pointer types and remove unnecessary null checks
- Remove range metadata from AND instruction (only apply to intrinsic calls)
- Simplify wave size logic by checking isWaveSizeKnown() upfront
- Use direct metadata access with F.getMetadata() instead of F->getMetadata()
- Add -mcpu=gfx906 to both test RUN lines to enable wave64 detection
- Update CHECK expectations to account for function attributes added
  when specifying a CPU target

All changes maintain the same optimization functionality while improving
code quality and following LLVM coding conventions.
---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    | 81 ++++++++-----------
 .../AMDGPU/mbcnt-to-bitmask-posit.ll          |  4 +-
 .../AMDGPU/mbcnt-to-workitem-posit.ll         |  4 +-
 3 files changed, 37 insertions(+), 52 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index ae5b49901e2b2..b346952a781a0 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -2107,21 +2107,16 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
   if (!Arg0C || !Arg1C || !Arg0C->isAllOnesValue() || !Arg1C->isZero())
     return false;
 
-  // Check reqd_work_group_size similar to mbcnt_hi case.
-  Function *F = I.getFunction();
-  if (!F)
+  // Abort if wave size is not known.
+  if (!ST.isWaveSizeKnown())
     return false;
 
-  unsigned Wave = 0;
-  if (ST.isWaveSizeKnown())
-    Wave = ST.getWavefrontSize();
+  unsigned Wave = ST.getWavefrontSize();
 
-  if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
+  if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
     unsigned XLen = *MaybeX;
-    if (Wave == 0 && XLen == 32)
-      Wave = XLen;
 
-    if (Wave != 0 && XLen == Wave) {
+    if (XLen == Wave) {
       IRBuilder<> B(&I);
       CallInst *NewCall =
           B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
@@ -2133,8 +2128,8 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
     }
     // Handle bitmask case: when X dimension evenly splits into waves.
     // mbcnt.lo(~0, 0) = workitem.id.x() & (wave_size - 1).
-    if (ST.hasWavefrontsEvenlySplittingXDim(*F, /*RequiresUniformYZ=*/true)) {
-      if (Wave != 0 && isPowerOf2_32(Wave)) {
+    if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
+      if (isPowerOf2_32(Wave)) {
         IRBuilder<> B(&I);
         CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
         ST.makeLIDRangeMetadata(Tid);
@@ -2157,20 +2152,16 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   // exec_hi is all 0, so this is just a copy on wave32.
   // However, only optimize if we have the same conditions as mbcnt.lo.
   if (ST.isWave32()) {
-    Function *F = I.getFunction();
-    if (!F)
+    // Abort if wave size is not known.
+    if (!ST.isWaveSizeKnown())
       return false;
 
-    unsigned Wave = 0;
-    if (ST.isWaveSizeKnown())
-      Wave = ST.getWavefrontSize();
+    unsigned Wave = ST.getWavefrontSize();
 
-    if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
+    if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
       unsigned XLen = *MaybeX;
-      if (Wave == 0 && XLen == 32)
-        Wave = XLen;
 
-      if (Wave != 0 && XLen == Wave) {
+      if (XLen == Wave) {
         I.replaceAllUsesWith(I.getArgOperand(1));
         I.eraseFromParent();
         return true;
@@ -2202,20 +2193,16 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
 
   // Query reqd_work_group_size via subtarget helper and compare X to wave
   // size conservatively.
-  Function *F = I.getFunction();
-  if (!F)
+  // Abort if wave size is not known.
+  if (!ST.isWaveSizeKnown())
     return false;
 
-  unsigned Wave = 0;
-  if (ST.isWaveSizeKnown())
-    Wave = ST.getWavefrontSize();
+  unsigned Wave = ST.getWavefrontSize();
 
-  if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
+  if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
     unsigned XLen = *MaybeX;
-    if (Wave == 0 && (XLen == 32 || XLen == 64))
-      Wave = XLen; // allow common sizes under test harness
 
-    if (Wave != 0 && XLen == Wave) {
+    if (XLen == Wave) {
       IRBuilder<> B(&I);
       CallInst *NewCall =
           B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
@@ -2230,18 +2217,17 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
     // replace lane-id computation with a bitmask when the wave is a
     // power-of-two. Use the Subtarget helper to conservatively decide
     // when per-wave tiling is preserved.
-    if (ST.hasWavefrontsEvenlySplittingXDim(*F, /*RequiresUniformYZ=*/true)) {
-      if (Wave != 0 && isPowerOf2_32(Wave)) {
+    if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
+      if (isPowerOf2_32(Wave)) {
         // Construct: tid = workitem.id.x(); mask = Wave-1; res = tid & mask
         IRBuilder<> B(&I);
         CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-        Tid->takeName(&I);
+        ST.makeLIDRangeMetadata(Tid);
         IntegerType *ITy = cast<IntegerType>(Tid->getType());
         Constant *Mask = ConstantInt::get(ITy, Wave - 1);
         Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
         AndInst->takeName(&I);
-        // Attach range metadata for the result if possible.
-        ST.makeLIDRangeMetadata(AndInst);
+        // Note: Range metadata cannot be applied to 'and' instructions.
         I.replaceAllUsesWith(AndInst);
         I.eraseFromParent();
         return true;
@@ -2251,19 +2237,18 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
     // No reqd_work_group_size metadata: be conservative and only handle the
     // common test harness cases where reqd_work_group_size metadata exists
     // and equals 32/64.
-    if (auto *Node = F->getMetadata("reqd_work_group_size")) {
-      if (Node->getNumOperands() == 3) {
-        unsigned XLen =
-            mdconst::extract<ConstantInt>(Node->getOperand(0))->getZExtValue();
-        if (XLen == 32 || XLen == 64) {
-          IRBuilder<> B(&I);
-          CallInst *NewCall =
-              B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-          NewCall->takeName(&I);
-          I.replaceAllUsesWith(NewCall);
-          I.eraseFromParent();
-          return true;
-        }
+    const MDNode *Node = F.getMetadata("reqd_work_group_size");
+    if (Node && Node->getNumOperands() == 3) {
+      unsigned XLen =
+          mdconst::extract<ConstantInt>(Node->getOperand(0))->getZExtValue();
+      if (XLen == Wave) {
+        IRBuilder<> B(&I);
+        CallInst *NewCall =
+            B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+        NewCall->takeName(&I);
+        I.replaceAllUsesWith(NewCall);
+        I.eraseFromParent();
+        return true;
       }
     }
   }
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
index 97697d4d9651f..ad3993bbc30c1 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
@@ -1,9 +1,9 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck %s
 
 define i32 @test_mbcnt_wave64_to_workitem() !reqd_work_group_size !1 {
 ; CHECK-LABEL: define i32 @test_mbcnt_wave64_to_workitem(
-; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; CHECK-NEXT:    [[B:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
index 91aa942df9337..fe049e1627409 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
@@ -1,9 +1,9 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck %s
 
 define i32 @test_mbcnt_to_workitem() !reqd_work_group_size !0 {
 ; CHECK-LABEL: define i32 @test_mbcnt_to_workitem(
-; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; CHECK-NEXT:    [[B:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()

>From a439df17e0409fd423d741082a5e554731d9f939 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Fri, 10 Oct 2025 14:54:58 +0530
Subject: [PATCH 06/11] Refactored test cases and implementation comments

---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    |  56 ++--
 .../Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll  |  44 ---
 .../Transforms/AMDGPU/mbcnt-negative-cases.ll | 253 ++++++++++++++++++
 .../Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll |  25 --
 .../AMDGPU/mbcnt-to-bitmask-posit.ll          |  26 --
 .../AMDGPU/mbcnt-to-workitem-neg.ll           |  20 --
 .../AMDGPU/mbcnt-to-workitem-posit.ll         |  26 --
 .../AMDGPU/mbcnt-to-workitem-wave32-neg.ll    |  67 -----
 .../AMDGPU/mbcnt-to-workitem-wave32.ll        |  25 --
 .../AMDGPU/mbcnt-wave32-optimizations.ll      | 138 ++++++++++
 .../AMDGPU/mbcnt-wave64-optimizations.ll      |  95 +++++++
 11 files changed, 517 insertions(+), 258 deletions(-)
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
 delete mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
 create mode 100644 llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index b346952a781a0..679941560779e 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -2096,18 +2096,21 @@ INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
 INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)
 
+/// Optimize mbcnt.lo calls on wave32 architectures for lane ID computation.
 bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
-  // On wave32 targets, mbcnt.lo(~0, 0) can be replaced with workitem.id.x.
+  // This optimization only applies to wave32 targets where mbcnt.lo operates on
+  // the full execution mask.
   if (!ST.isWave32())
     return false;
 
-  // Check for pattern mbcnt.lo(~0, 0).
+  // Only optimize the pattern mbcnt.lo(~0, 0) which counts active lanes with
+  // lower IDs.
   auto *Arg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
   auto *Arg1C = dyn_cast<ConstantInt>(I.getArgOperand(1));
   if (!Arg0C || !Arg1C || !Arg0C->isAllOnesValue() || !Arg1C->isZero())
     return false;
 
-  // Abort if wave size is not known.
+  // Abort if wave size is not known at compile time.
   if (!ST.isWaveSizeKnown())
     return false;
 
@@ -2116,6 +2119,8 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
   if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
     unsigned XLen = *MaybeX;
 
+    // When XLen == wave_size, each work group contains exactly one wave, so
+    // mbcnt.lo(~0, 0) directly equals the workitem ID within the group.
     if (XLen == Wave) {
       IRBuilder<> B(&I);
       CallInst *NewCall =
@@ -2126,8 +2131,9 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
       I.eraseFromParent();
       return true;
     }
-    // Handle bitmask case: when X dimension evenly splits into waves.
-    // mbcnt.lo(~0, 0) = workitem.id.x() & (wave_size - 1).
+    // When work group evenly splits into waves and wave size is power-of-2,
+    // we can compute lane ID within wave using bit masking:
+    // lane_id = workitem.id.x & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
       if (isPowerOf2_32(Wave)) {
         IRBuilder<> B(&I);
@@ -2137,7 +2143,6 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
         Constant *Mask = ConstantInt::get(ITy, Wave - 1);
         Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
         AndInst->takeName(&I);
-        // Note: Range metadata cannot be applied to 'and' instructions.
         I.replaceAllUsesWith(AndInst);
         I.eraseFromParent();
         return true;
@@ -2148,11 +2153,12 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
   return false;
 }
 
+/// Optimize mbcnt.hi calls for lane ID computation.
 bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
-  // exec_hi is all 0, so this is just a copy on wave32.
-  // However, only optimize if we have the same conditions as mbcnt.lo.
+  // On wave32, the upper 32 bits of exec are always 0, so mbcnt.hi(mask, val)
+  // always returns val unchanged.
   if (ST.isWave32()) {
-    // Abort if wave size is not known.
+    // Abort if wave size is not known at compile time.
     if (!ST.isWaveSizeKnown())
       return false;
 
@@ -2161,6 +2167,8 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
     if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
       unsigned XLen = *MaybeX;
 
+      // Replace mbcnt.hi(mask, val) with val only when work group size matches
+      // wave size (single wave per work group).
       if (XLen == Wave) {
         I.replaceAllUsesWith(I.getArgOperand(1));
         I.eraseFromParent();
@@ -2169,7 +2177,9 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
     }
   }
 
-  // Pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0)).
+  // Optimize the complete lane ID computation pattern:
+  // mbcnt.hi(~0, mbcnt.lo(~0, 0)) which counts all active lanes with lower IDs
+  // across the full execution mask.
   auto *HiArg1 = dyn_cast<CallInst>(I.getArgOperand(1));
   if (!HiArg1)
     return false;
@@ -2178,12 +2188,12 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   if (!CalledF || CalledF->getIntrinsicID() != Intrinsic::amdgcn_mbcnt_lo)
     return false;
 
-  // hi arg0 must be all-ones.
+  // mbcnt.hi mask must be all-ones (count from upper 32 bits)
   auto *HiArg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
   if (!HiArg0C || !HiArg0C->isAllOnesValue())
     return false;
 
-  // lo args: arg0 == ~0, arg1 == 0.
+  // mbcnt.lo mask must be all-ones (mask=~0, all lanes) and base must be 0.
   Value *Lo0 = HiArg1->getArgOperand(0);
   Value *Lo1 = HiArg1->getArgOperand(1);
   auto *Lo0C = dyn_cast<ConstantInt>(Lo0);
@@ -2191,9 +2201,7 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   if (!Lo0C || !Lo1C || !Lo0C->isAllOnesValue() || !Lo1C->isZero())
     return false;
 
-  // Query reqd_work_group_size via subtarget helper and compare X to wave
-  // size conservatively.
-  // Abort if wave size is not known.
+  // Abort if wave size is not known at compile time.
   if (!ST.isWaveSizeKnown())
     return false;
 
@@ -2202,24 +2210,24 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
     unsigned XLen = *MaybeX;
 
+    // When XLen == wave_size, each work group contains exactly one wave, so
+    // lane_id = workitem.id.x.
     if (XLen == Wave) {
       IRBuilder<> B(&I);
       CallInst *NewCall =
           B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
       NewCall->takeName(&I);
-      // Attach range metadata when available.
       ST.makeLIDRangeMetadata(NewCall);
       I.replaceAllUsesWith(NewCall);
       I.eraseFromParent();
       return true;
     }
-    // Optional: if X dimension evenly splits into wavefronts we can
-    // replace lane-id computation with a bitmask when the wave is a
-    // power-of-two. Use the Subtarget helper to conservatively decide
-    // when per-wave tiling is preserved.
+    // When work group evenly splits into waves and wave size is power-of-2,
+    // we can compute lane ID within wave using bit masking:
+    // lane_id = workitem.id.x & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
       if (isPowerOf2_32(Wave)) {
-        // Construct: tid = workitem.id.x(); mask = Wave-1; res = tid & mask
+        // Construct optimized sequence: workitem.id.x & (wave_size - 1)
         IRBuilder<> B(&I);
         CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
         ST.makeLIDRangeMetadata(Tid);
@@ -2227,16 +2235,14 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
         Constant *Mask = ConstantInt::get(ITy, Wave - 1);
         Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
         AndInst->takeName(&I);
-        // Note: Range metadata cannot be applied to 'and' instructions.
         I.replaceAllUsesWith(AndInst);
         I.eraseFromParent();
         return true;
       }
     }
   } else {
-    // No reqd_work_group_size metadata: be conservative and only handle the
-    // common test harness cases where reqd_work_group_size metadata exists
-    // and equals 32/64.
+    // When ST.getReqdWorkGroupSize() fails, use metadata. And only optimize the
+    // case when work group size = wave size.
     const MDNode *Node = F.getMetadata("reqd_work_group_size");
     if (Node && Node->getNumOperands() == 3) {
       unsigned XLen =
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll b/llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll
deleted file mode 100644
index f84217705818e..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-lo-to-bitmask.ll
+++ /dev/null
@@ -1,44 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
-
-; Test that mbcnt.lo(~0, 0) is optimized to workitem.id.x() & 0x1f on wave32 
-; when work group size is multiple of wave size (64 = 2 * 32)
-define i32 @test_mbcnt_lo_wave32_bitmask() !reqd_work_group_size !0 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_bitmask(
-; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call {{.*}} i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
-; CHECK-NEXT:    ret i32 [[TMP1]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
-; Test with X dimension = 128 (4 * 32 waves)
-define i32 @test_mbcnt_lo_wave32_bitmask_128() !reqd_work_group_size !1 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_bitmask_128(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call {{.*}} i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
-; CHECK-NEXT:    ret i32 [[TMP1]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
-!0 = !{i32 64, i32 1, i32 1}   ; 64 = 2 * 32 wave size
-!1 = !{i32 128, i32 1, i32 1}  ; 128 = 4 * 32 wave size
-
-; Function Attrs: nounwind readnone speculatable willreturn
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
-declare i32 @llvm.amdgcn.workitem.id.x() #0
-
-attributes #0 = { nounwind readnone speculatable willreturn }
-;.
-; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
-; CHECK: [[META1]] = !{i32 128, i32 1, i32 1}
-;.
\ No newline at end of file
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll b/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
new file mode 100644
index 0000000000000..caa2438921a5b
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
@@ -0,0 +1,253 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX9 %s
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX10 %s
+
+; Test negative cases where mbcnt optimizations should NOT be applied
+
+; =============================================================================
+; NO WORK GROUP SIZE METADATA
+; =============================================================================
+
+; Test with no reqd_work_group_size
+define i32 @test_mbcnt_no_work_group_size() {
+; GFX9-LABEL: define i32 @test_mbcnt_no_work_group_size() {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX9-NEXT:    ret i32 [[B]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_no_work_group_size(
+; GFX10-SAME: ) #[[ATTR0:[0-9]+]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX10-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Test mbcnt.lo with no work group size
+define i32 @test_mbcnt_lo_no_work_group_size() {
+; GFX9-LABEL: define i32 @test_mbcnt_lo_no_work_group_size() {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX9-NEXT:    ret i32 [[A]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_lo_no_work_group_size(
+; GFX10-SAME: ) #[[ATTR0]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX10-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; =============================================================================
+; NON-WAVE-MULTIPLE WORK GROUP SIZES
+; =============================================================================
+
+; Test with work group size = not a wave multiple (48)
+define i32 @test_mbcnt_non_wave_multiple() !reqd_work_group_size !0 {
+; GFX9-LABEL: define i32 @test_mbcnt_non_wave_multiple(
+; GFX9-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX9-NEXT:    ret i32 [[B]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_non_wave_multiple(
+; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0:![0-9]+]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
+; GFX10-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
+; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX10-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Test mbcnt.lo with work group size = not a wave multiple (48)
+define i32 @test_mbcnt_lo_non_wave_multiple() !reqd_work_group_size !0 {
+; GFX9-LABEL: define i32 @test_mbcnt_lo_non_wave_multiple(
+; GFX9-SAME: ) !reqd_work_group_size [[META0]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX9-NEXT:    ret i32 [[A]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_lo_non_wave_multiple(
+; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
+; GFX10-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
+; GFX10-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; =============================================================================
+; PARTIAL MASKS AND NON-STANDARD PATTERNS
+; =============================================================================
+
+; Test with partial mask
+define i32 @test_mbcnt_partial_mask() !reqd_work_group_size !1 {
+; GFX9-LABEL: define i32 @test_mbcnt_partial_mask(
+; GFX9-SAME: ) !reqd_work_group_size [[META1:![0-9]+]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
+; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX9-NEXT:    ret i32 [[B]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_partial_mask(
+; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
+; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX10-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Test with non-zero base
+define i32 @test_mbcnt_non_zero_base() !reqd_work_group_size !1 {
+; GFX9-LABEL: define i32 @test_mbcnt_non_zero_base(
+; GFX9-SAME: ) !reqd_work_group_size [[META1]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
+; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX9-NEXT:    ret i32 [[B]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_non_zero_base(
+; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
+; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; GFX10-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; =============================================================================
+; ARCHITECTURE-SPECIFIC NEGATIVE CASES
+; =============================================================================
+
+; Test with work group size = partial wave multiple (48)
+define i32 @test_mbcnt_lo_wave32_partial_wave() !reqd_work_group_size !0 {
+; GFX9-LABEL: define i32 @test_mbcnt_lo_wave32_partial_wave(
+; GFX9-SAME: ) !reqd_work_group_size [[META0]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; GFX9-NEXT:    ret i32 [[A]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_lo_wave32_partial_wave(
+; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
+; GFX10-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
+; GFX10-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; =============================================================================
+; COPY OPTIMIZATION NEGATIVE CASES
+; =============================================================================
+
+; Test with no work group size
+define i32 @test_mbcnt_hi_copy_no_wgs(i32 %val) {
+; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_no_wgs(
+; GFX9-SAME: i32 [[VAL:%.*]]) {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
+; GFX9-NEXT:    ret i32 [[RESULT]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_no_wgs(
+; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
+; GFX10-NEXT:    ret i32 [[RESULT]]
+;
+entry:
+  %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %val)
+  ret i32 %result
+}
+
+; Test with work group size = not a wave multiple (48)
+define i32 @test_mbcnt_hi_copy_non_wave_multiple(i32 %val) !reqd_work_group_size !0 {
+; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
+; GFX9-SAME: i32 [[VAL:%.*]]) !reqd_work_group_size [[META0]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
+; GFX9-NEXT:    ret i32 [[RESULT]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
+; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
+; GFX10-NEXT:    ret i32 [[RESULT]]
+;
+entry:
+  %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %val)
+  ret i32 %result
+}
+
+; Test with zero mask
+define i32 @test_mbcnt_hi_copy_zero_mask(i32 %val) !reqd_work_group_size !1 {
+; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
+; GFX9-SAME: i32 [[VAL:%.*]]) !reqd_work_group_size [[META1]] {
+; GFX9-NEXT:  [[ENTRY:.*:]]
+; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
+; GFX9-NEXT:    ret i32 [[RESULT]]
+;
+; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
+; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META1]] {
+; GFX10-NEXT:  [[ENTRY:.*:]]
+; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
+; GFX10-NEXT:    ret i32 [[RESULT]]
+;
+entry:
+  %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 %val)
+  ret i32 %result
+}
+
+; =============================================================================
+; METADATA
+; =============================================================================
+
+!0 = !{i32 48, i32 1, i32 1}   ; X=48 (not wave multiple)
+!1 = !{i32 64, i32 1, i32 1}   ; X=64 (wave64 or 2*wave32)
+
+; =============================================================================
+; FUNCTION DECLARATIONS
+; =============================================================================
+
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32) #0
+
+attributes #0 = { nounwind readnone speculatable willreturn }
+;.
+; GFX9: [[META0]] = !{i32 48, i32 1, i32 1}
+; GFX9: [[META1]] = !{i32 64, i32 1, i32 1}
+;.
+; GFX10: [[META0]] = !{i32 48, i32 1, i32 1}
+; GFX10: [[META1]] = !{i32 64, i32 1, i32 1}
+;.
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CHECK: {{.*}}
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll
deleted file mode 100644
index 470751c3c73f3..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-neg.ll
+++ /dev/null
@@ -1,25 +0,0 @@
-;; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %sNOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple=amdgcn-unknown-amdhsa -passes=instcombine < %s | FileCheck %s
-
-define i32 @test_mbcnt_non_wave_size() !reqd_work_group_size !1 {
-; CHECK-LABEL: define i32 @test_mbcnt_non_wave_size(
-; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; CHECK-NEXT:    ret i32 [[B]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
-  ret i32 %b
-}
-
-!1 = !{i32 48, i32 1, i32 1}
-
-; Declarations
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
-declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
-;.
-; CHECK: [[META0]] = !{i32 48, i32 1, i32 1}
-;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
deleted file mode 100644
index ad3993bbc30c1..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-bitmask-posit.ll
+++ /dev/null
@@ -1,26 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck %s
-
-define i32 @test_mbcnt_wave64_to_workitem() !reqd_work_group_size !1 {
-; CHECK-LABEL: define i32 @test_mbcnt_wave64_to_workitem(
-; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; CHECK-NEXT:    [[B:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    ret i32 [[B]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
-  ret i32 %b
-}
-
-!1 = !{i32 64, i32 1, i32 1}
-
-; Declarations
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
-declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
-declare i32 @llvm.amdgcn.workitem.id.x()
-;.
-; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
-;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll
deleted file mode 100644
index af8d713b798ed..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-neg.ll
+++ /dev/null
@@ -1,20 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
-
-define i32 @test_mbcnt_no_reqd_work_group_size() {
-; CHECK-LABEL: define i32 @test_mbcnt_no_reqd_work_group_size() {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; CHECK-NEXT:    ret i32 [[B]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
-  ret i32 %b
-}
-
-; Declarations
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
-declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
-declare i32 @llvm.amdgcn.workitem.id.x()
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
deleted file mode 100644
index fe049e1627409..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-posit.ll
+++ /dev/null
@@ -1,26 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck %s
-
-define i32 @test_mbcnt_to_workitem() !reqd_work_group_size !0 {
-; CHECK-LABEL: define i32 @test_mbcnt_to_workitem(
-; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; CHECK-NEXT:    [[B:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    ret i32 [[B]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
-  ret i32 %b
-}
-
-!0 = !{i32 64, i32 1, i32 1}
-
-; Declarations
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)
-declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)
-declare i32 @llvm.amdgcn.workitem.id.x()
-;.
-; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
-;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
deleted file mode 100644
index 30514562556ea..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32-neg.ll
+++ /dev/null
@@ -1,67 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
-
-; Test that mbcnt.lo(~0, 0) IS optimized on wave32 with bitmask when work group size allows even wave distribution
-define i32 @test_mbcnt_lo_wave32_non_matching_wgs() !reqd_work_group_size !0 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_non_matching_wgs(
-; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call {{.*}} i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
-; CHECK-NEXT:    ret i32 [[A]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
-; Test that mbcnt.lo(~0, 0) is NOT optimized on wave32 when no reqd_work_group_size is specified
-define i32 @test_mbcnt_lo_wave32_no_wgs() {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_no_wgs(
-; CHECK-SAME: ) #[[ATTR0]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; CHECK-NEXT:    ret i32 [[A]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
-; Test that mbcnt.lo with non-all-ones first arg is NOT optimized
-define i32 @test_mbcnt_lo_wave32_partial_mask() !reqd_work_group_size !1 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_partial_mask(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 15, i32 0)
-; CHECK-NEXT:    ret i32 [[A]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 15, i32 0)
-  ret i32 %a
-}
-
-; Test that mbcnt.lo with non-zero second arg is NOT optimized
-define i32 @test_mbcnt_lo_wave32_non_zero_base() !reqd_work_group_size !1 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32_non_zero_base(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
-; CHECK-NEXT:    ret i32 [[A]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
-  ret i32 %a
-}
-
-!0 = !{i32 48, i32 1, i32 1}  ; Work group size 48 != wave size 32
-!1 = !{i32 32, i32 1, i32 1}  ; Work group size 32 == wave size 32
-
-; Function Attrs: nounwind readnone speculatable willreturn
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
-
-attributes #0 = { nounwind readnone speculatable willreturn }
-;.
-; CHECK: [[META0]] = !{i32 48, i32 1, i32 1}
-; CHECK: [[META1]] = !{i32 32, i32 1, i32 1}
-;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll b/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll
deleted file mode 100644
index 07a5028ca1ee5..0000000000000
--- a/llvm/test/Transforms/AMDGPU/mbcnt-to-workitem-wave32.ll
+++ /dev/null
@@ -1,25 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
-
-define i32 @test_mbcnt_lo_wave32() !reqd_work_group_size !0 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_wave32(
-; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 32) i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    ret i32 [[TMP0]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
-!0 = !{i32 32, i32 1, i32 1}
-
-; Function Attrs: nounwind readnone speculatable willreturn
-declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
-declare i32 @llvm.amdgcn.workitem.id.x() #0
-
-attributes #0 = { nounwind readnone speculatable willreturn }
-;.
-; CHECK: [[META0]] = !{i32 32, i32 1, i32 1}
-;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll b/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
new file mode 100644
index 0000000000000..975d3f295e9db
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
@@ -0,0 +1,138 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+; Test mbcnt optimizations for wave32 architecture
+; Covers: simple replacement, bitmask optimization, copy optimization
+
+; =============================================================================
+; SIMPLE REPLACEMENT OPTIMIZATIONS (mbcnt.lo -> workitem.id.x)
+; =============================================================================
+
+; Test with work group size = wave size
+define i32 @test_mbcnt_lo_simple_wave32() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_simple_wave32(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 32) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; =============================================================================
+; BITMASK OPTIMIZATIONS (mbcnt.lo -> workitem.id.x & 0x1f)
+; =============================================================================
+
+; Test with work group size = 2 * wave size
+define i32 @test_mbcnt_lo_bitmask_64() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_64(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; Test with work group size = 4 * wave size
+define i32 @test_mbcnt_lo_bitmask_128() !reqd_work_group_size !3 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_128(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META2:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 128) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; Test with work group size = non-power-of-2 but wave-multiple
+define i32 @test_mbcnt_lo_bitmask_96() !reqd_work_group_size !2 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_96(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 96) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
+; =============================================================================
+; COPY OPTIMIZATION (mbcnt.hi(mask, val) -> val)
+; =============================================================================
+
+; Test with mask = wave32 range
+define i32 @test_mbcnt_hi_copy_basic(i32 %val) !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_hi_copy_basic(
+; CHECK-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret i32 [[VAL]]
+;
+entry:
+  %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %val)
+  ret i32 %result
+}
+
+; Test with partial mask
+define i32 @test_mbcnt_hi_copy_partial_mask(i32 %val) !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_hi_copy_partial_mask(
+; CHECK-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret i32 [[VAL]]
+;
+entry:
+  %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 65535, i32 %val)
+  ret i32 %result
+}
+
+; =============================================================================
+; FULL PATTERN OPTIMIZATION (mbcnt.hi(~0, mbcnt.lo(~0, 0)) -> workitem.id.x)
+; =============================================================================
+
+; Test full pattern on wave32
+define i32 @test_mbcnt_full_pattern_wave32() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave32(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 32) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; =============================================================================
+; METADATA
+; =============================================================================
+
+!0 = !{i32 32, i32 1, i32 1}   ; X=32 (1*wave32), Y=1, Z=1
+!1 = !{i32 64, i32 1, i32 1}   ; X=64 (2*wave32), Y=1, Z=1
+!2 = !{i32 96, i32 1, i32 1}   ; X=96 (3*wave32), Y=1, Z=1
+!3 = !{i32 128, i32 1, i32 1}  ; X=128 (4*wave32), Y=1, Z=1
+
+; =============================================================================
+; FUNCTION DECLARATIONS
+; =============================================================================
+
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32) #0
+
+attributes #0 = { nounwind readnone speculatable willreturn }
+;.
+; CHECK: [[META0]] = !{i32 32, i32 1, i32 1}
+; CHECK: [[META1]] = !{i32 64, i32 1, i32 1}
+; CHECK: [[META2]] = !{i32 128, i32 1, i32 1}
+; CHECK: [[META3]] = !{i32 96, i32 1, i32 1}
+;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll b/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
new file mode 100644
index 0000000000000..42ba27a37b094
--- /dev/null
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
@@ -0,0 +1,95 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck %s
+
+; Test full mbcnt pattern optimization for wave64 architecture
+; Covers: full pattern optimization
+
+; =============================================================================
+; FULL PATTERN OPTIMIZATION - mbcnt.hi(~0, mbcnt.lo(~0, 0)) -> workitem.id.x
+; =============================================================================
+
+; Test with work group size = wave size (64)
+define i32 @test_mbcnt_full_pattern_wave64_basic() !reqd_work_group_size !0 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_basic(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Test with work group size = 2 * wave size (128)
+define i32 @test_mbcnt_full_pattern_wave64_128() !reqd_work_group_size !1 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_128(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 128) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[B:%.*]] = and i32 [[TMP0]], 63
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Test with work group size = 4 * wave size (256)
+define i32 @test_mbcnt_full_pattern_wave64_256() !reqd_work_group_size !2 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_256(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META2:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 256) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[B:%.*]] = and i32 [[TMP0]], 63
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; Test with multidimensional work group where X dimension matches pattern
+define i32 @test_mbcnt_full_pattern_wave64_multidim() !reqd_work_group_size !3 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_multidim(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
+; =============================================================================
+; METADATA
+; =============================================================================
+
+!0 = !{i32 64, i32 1, i32 1}   ; X=64 (1*wave), Y=1, Z=1
+!1 = !{i32 128, i32 1, i32 1}  ; X=128 (2*wave), Y=1, Z=1
+!2 = !{i32 256, i32 1, i32 1}  ; X=256 (4*wave), Y=1, Z=1
+!3 = !{i32 64, i32 2, i32 1}   ; X=64 (1*wave), Y=2, Z=1
+
+; =============================================================================
+; FUNCTION DECLARATIONS
+; =============================================================================
+
+declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
+declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32) #0
+
+attributes #0 = { nounwind readnone speculatable willreturn }
+;.
+; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
+; CHECK: [[META1]] = !{i32 128, i32 1, i32 1}
+; CHECK: [[META2]] = !{i32 256, i32 1, i32 1}
+; CHECK: [[META3]] = !{i32 64, i32 2, i32 1}
+;.

>From d554cecf0e773ff8ef2692ade7cd900e88479724 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Wed, 29 Oct 2025 11:08:54 +0530
Subject: [PATCH 07/11] Simplified the implementation per review suggestions -
 1. Remove unnecessary casts 2. Replace manual AND instruction injection with
 ReplaceInstWithValue utility 3. Replace manual pattern matching with match
 utility 4. Rearange the logic to avoid redundancy 5. Add additional logic to
 check if the reqd wave size is a multiple of Wave size for full pattern
 optimization

---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    | 108 +++++++-----------
 1 file changed, 42 insertions(+), 66 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 679941560779e..864ddb6d82fdb 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -35,6 +35,7 @@
 #include "llvm/Support/KnownFPClass.h"
 #include "llvm/Transforms/Utils/IntegerDivision.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include <llvm/Transforms/Utils/BasicBlockUtils.h>
 
 #define DEBUG_TYPE "amdgpu-codegenprepare"
 
@@ -2098,6 +2099,10 @@ INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
 
 /// Optimize mbcnt.lo calls on wave32 architectures for lane ID computation.
 bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
+  // Abort if wave size is not known at compile time.
+  if (!ST.isWaveSizeKnown())
+    return false;
+
   // This optimization only applies to wave32 targets where mbcnt.lo operates on
   // the full execution mask.
   if (!ST.isWave32())
@@ -2110,10 +2115,6 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
   if (!Arg0C || !Arg1C || !Arg0C->isAllOnesValue() || !Arg1C->isZero())
     return false;
 
-  // Abort if wave size is not known at compile time.
-  if (!ST.isWaveSizeKnown())
-    return false;
-
   unsigned Wave = ST.getWavefrontSize();
 
   if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
@@ -2122,13 +2123,11 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
     // When XLen == wave_size, each work group contains exactly one wave, so
     // mbcnt.lo(~0, 0) directly equals the workitem ID within the group.
     if (XLen == Wave) {
-      IRBuilder<> B(&I);
-      CallInst *NewCall =
-          B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-      NewCall->takeName(&I);
+      Function *WorkitemIdFn = Intrinsic::getOrInsertDeclaration(
+          I.getModule(), Intrinsic::amdgcn_workitem_id_x);
+      CallInst *NewCall = CallInst::Create(WorkitemIdFn, I.getName());
+      ReplaceInstWithInst(&I, NewCall);
       ST.makeLIDRangeMetadata(NewCall);
-      I.replaceAllUsesWith(NewCall);
-      I.eraseFromParent();
       return true;
     }
     // When work group evenly splits into waves and wave size is power-of-2,
@@ -2139,12 +2138,10 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
         IRBuilder<> B(&I);
         CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
         ST.makeLIDRangeMetadata(Tid);
-        IntegerType *ITy = cast<IntegerType>(Tid->getType());
-        Constant *Mask = ConstantInt::get(ITy, Wave - 1);
-        Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
-        AndInst->takeName(&I);
-        I.replaceAllUsesWith(AndInst);
-        I.eraseFromParent();
+        Constant *Mask = ConstantInt::get(Tid->getType(), Wave - 1);
+        Value *AndInst = B.CreateAnd(Tid, Mask);
+        BasicBlock::iterator BI(&I);
+        ReplaceInstWithValue(BI, AndInst);
         return true;
       }
     }
@@ -2155,23 +2152,24 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
 
 /// Optimize mbcnt.hi calls for lane ID computation.
 bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
-  // On wave32, the upper 32 bits of exec are always 0, so mbcnt.hi(mask, val)
-  // always returns val unchanged.
-  if (ST.isWave32()) {
-    // Abort if wave size is not known at compile time.
-    if (!ST.isWaveSizeKnown())
-      return false;
+  // Abort if wave size is not known at compile time.
+  if (!ST.isWaveSizeKnown())
+    return false;
 
-    unsigned Wave = ST.getWavefrontSize();
+  // Calculate wave size once at the beginning
+  unsigned Wave = ST.getWavefrontSize();
 
+  // On wave32, the upper 32 bits of execution mask are always 0, so
+  // mbcnt.hi(mask, val) always returns val unchanged.
+  if (ST.isWave32()) {
     if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
       unsigned XLen = *MaybeX;
 
       // Replace mbcnt.hi(mask, val) with val only when work group size matches
       // wave size (single wave per work group).
       if (XLen == Wave) {
-        I.replaceAllUsesWith(I.getArgOperand(1));
-        I.eraseFromParent();
+        BasicBlock::iterator BI(&I);
+        ReplaceInstWithValue(BI, I.getArgOperand(1));
         return true;
       }
     }
@@ -2180,63 +2178,42 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
   // Optimize the complete lane ID computation pattern:
   // mbcnt.hi(~0, mbcnt.lo(~0, 0)) which counts all active lanes with lower IDs
   // across the full execution mask.
-  auto *HiArg1 = dyn_cast<CallInst>(I.getArgOperand(1));
-  if (!HiArg1)
-    return false;
+  using namespace PatternMatch;
 
-  Function *CalledF = HiArg1->getCalledFunction();
-  if (!CalledF || CalledF->getIntrinsicID() != Intrinsic::amdgcn_mbcnt_lo)
+  // Check for pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0))
+  if (!match(I.getArgOperand(0), m_AllOnes()))
     return false;
 
-  // mbcnt.hi mask must be all-ones (count from upper 32 bits)
-  auto *HiArg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
-  if (!HiArg0C || !HiArg0C->isAllOnesValue())
+  if (!match(I.getArgOperand(1),
+             m_Intrinsic<Intrinsic::amdgcn_mbcnt_lo>(m_AllOnes(), m_Zero())))
     return false;
 
-  // mbcnt.lo mask must be all-ones (mask=~0, all lanes) and base must be 0.
-  Value *Lo0 = HiArg1->getArgOperand(0);
-  Value *Lo1 = HiArg1->getArgOperand(1);
-  auto *Lo0C = dyn_cast<ConstantInt>(Lo0);
-  auto *Lo1C = dyn_cast<ConstantInt>(Lo1);
-  if (!Lo0C || !Lo1C || !Lo0C->isAllOnesValue() || !Lo1C->isZero())
-    return false;
-
-  // Abort if wave size is not known at compile time.
-  if (!ST.isWaveSizeKnown())
-    return false;
-
-  unsigned Wave = ST.getWavefrontSize();
-
   if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
     unsigned XLen = *MaybeX;
 
     // When XLen == wave_size, each work group contains exactly one wave, so
     // lane_id = workitem.id.x.
     if (XLen == Wave) {
-      IRBuilder<> B(&I);
-      CallInst *NewCall =
-          B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-      NewCall->takeName(&I);
+      Function *WorkitemIdFn = Intrinsic::getOrInsertDeclaration(
+          I.getModule(), Intrinsic::amdgcn_workitem_id_x);
+      CallInst *NewCall = CallInst::Create(WorkitemIdFn, I.getName());
+      ReplaceInstWithInst(&I, NewCall);
       ST.makeLIDRangeMetadata(NewCall);
-      I.replaceAllUsesWith(NewCall);
-      I.eraseFromParent();
       return true;
     }
     // When work group evenly splits into waves and wave size is power-of-2,
     // we can compute lane ID within wave using bit masking:
     // lane_id = workitem.id.x & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
-      if (isPowerOf2_32(Wave)) {
+      if (XLen % Wave == 0 && isPowerOf2_32(Wave)) {
         // Construct optimized sequence: workitem.id.x & (wave_size - 1)
         IRBuilder<> B(&I);
         CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
         ST.makeLIDRangeMetadata(Tid);
-        IntegerType *ITy = cast<IntegerType>(Tid->getType());
-        Constant *Mask = ConstantInt::get(ITy, Wave - 1);
-        Instruction *AndInst = cast<Instruction>(B.CreateAnd(Tid, Mask));
-        AndInst->takeName(&I);
-        I.replaceAllUsesWith(AndInst);
-        I.eraseFromParent();
+        Constant *Mask = ConstantInt::get(Tid->getType(), Wave - 1);
+        Value *AndInst = B.CreateAnd(Tid, Mask);
+        BasicBlock::iterator BI(&I);
+        ReplaceInstWithValue(BI, AndInst);
         return true;
       }
     }
@@ -2248,12 +2225,11 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
       unsigned XLen =
           mdconst::extract<ConstantInt>(Node->getOperand(0))->getZExtValue();
       if (XLen == Wave) {
-        IRBuilder<> B(&I);
-        CallInst *NewCall =
-            B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-        NewCall->takeName(&I);
-        I.replaceAllUsesWith(NewCall);
-        I.eraseFromParent();
+        Function *WorkitemIdFn = Intrinsic::getOrInsertDeclaration(
+            I.getModule(), Intrinsic::amdgcn_workitem_id_x);
+        CallInst *NewCall = CallInst::Create(WorkitemIdFn, I.getName());
+        ReplaceInstWithInst(&I, NewCall);
+        ST.makeLIDRangeMetadata(NewCall);
         return true;
       }
     }

>From 2868b1c212a661b6225a2ef6fd8b3802ba483cd0 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Wed, 29 Oct 2025 14:12:33 +0530
Subject: [PATCH 08/11] Addressed feedback on - 1. Add const qualifiers to
 mbcnt lo & hi visit functions 2. Reorganize partial wave size cases in the
 test files

---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    |  17 ++-
 .../Transforms/AMDGPU/mbcnt-negative-cases.ll | 117 ++++--------------
 .../AMDGPU/mbcnt-wave32-optimizations.ll      |  58 +++++++--
 .../AMDGPU/mbcnt-wave64-optimizations.ll      |  18 +++
 4 files changed, 95 insertions(+), 115 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 864ddb6d82fdb..dbf07e6aa091f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -261,8 +261,8 @@ class AMDGPUCodeGenPrepareImpl
   bool visitIntrinsicInst(IntrinsicInst &I);
   bool visitFMinLike(IntrinsicInst &I);
   bool visitSqrt(IntrinsicInst &I);
-  bool visitMbcntLo(IntrinsicInst &I);
-  bool visitMbcntHi(IntrinsicInst &I);
+  bool visitMbcntLo(IntrinsicInst &I) const;
+  bool visitMbcntHi(IntrinsicInst &I) const;
   bool run();
 };
 
@@ -2098,7 +2098,7 @@ INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)
 
 /// Optimize mbcnt.lo calls on wave32 architectures for lane ID computation.
-bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
+bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) const {
   // Abort if wave size is not known at compile time.
   if (!ST.isWaveSizeKnown())
     return false;
@@ -2110,9 +2110,8 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
 
   // Only optimize the pattern mbcnt.lo(~0, 0) which counts active lanes with
   // lower IDs.
-  auto *Arg0C = dyn_cast<ConstantInt>(I.getArgOperand(0));
-  auto *Arg1C = dyn_cast<ConstantInt>(I.getArgOperand(1));
-  if (!Arg0C || !Arg1C || !Arg0C->isAllOnesValue() || !Arg1C->isZero())
+  if (!match(&I,
+             m_Intrinsic<Intrinsic::amdgcn_mbcnt_lo>(m_AllOnes(), m_Zero())))
     return false;
 
   unsigned Wave = ST.getWavefrontSize();
@@ -2151,12 +2150,12 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) {
 }
 
 /// Optimize mbcnt.hi calls for lane ID computation.
-bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
+bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) const {
   // Abort if wave size is not known at compile time.
   if (!ST.isWaveSizeKnown())
     return false;
 
-  // Calculate wave size once at the beginning
+  // Calculate wave size
   unsigned Wave = ST.getWavefrontSize();
 
   // On wave32, the upper 32 bits of execution mask are always 0, so
@@ -2205,7 +2204,7 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) {
     // we can compute lane ID within wave using bit masking:
     // lane_id = workitem.id.x & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
-      if (XLen % Wave == 0 && isPowerOf2_32(Wave)) {
+      if (isPowerOf2_32(Wave)) {
         // Construct optimized sequence: workitem.id.x & (wave_size - 1)
         IRBuilder<> B(&I);
         CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll b/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
index caa2438921a5b..758f76358c771 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX9 %s
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX9 %s
 ; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX10 %s
 
 ; Test negative cases where mbcnt optimizations should NOT be applied
@@ -10,7 +10,8 @@
 
 ; Test with no reqd_work_group_size
 define i32 @test_mbcnt_no_work_group_size() {
-; GFX9-LABEL: define i32 @test_mbcnt_no_work_group_size() {
+; GFX9-LABEL: define i32 @test_mbcnt_no_work_group_size(
+; GFX9-SAME: ) #[[ATTR0:[0-9]+]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
@@ -31,7 +32,8 @@ entry:
 
 ; Test mbcnt.lo with no work group size
 define i32 @test_mbcnt_lo_no_work_group_size() {
-; GFX9-LABEL: define i32 @test_mbcnt_lo_no_work_group_size() {
+; GFX9-LABEL: define i32 @test_mbcnt_lo_no_work_group_size(
+; GFX9-SAME: ) #[[ATTR0]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; GFX9-NEXT:    ret i32 [[A]]
@@ -47,68 +49,21 @@ entry:
   ret i32 %a
 }
 
-; =============================================================================
-; NON-WAVE-MULTIPLE WORK GROUP SIZES
-; =============================================================================
-
-; Test with work group size = not a wave multiple (48)
-define i32 @test_mbcnt_non_wave_multiple() !reqd_work_group_size !0 {
-; GFX9-LABEL: define i32 @test_mbcnt_non_wave_multiple(
-; GFX9-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX9-NEXT:    ret i32 [[B]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_non_wave_multiple(
-; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0:![0-9]+]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
-; GFX10-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
-; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX10-NEXT:    ret i32 [[B]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
-  ret i32 %b
-}
-
-; Test mbcnt.lo with work group size = not a wave multiple (48)
-define i32 @test_mbcnt_lo_non_wave_multiple() !reqd_work_group_size !0 {
-; GFX9-LABEL: define i32 @test_mbcnt_lo_non_wave_multiple(
-; GFX9-SAME: ) !reqd_work_group_size [[META0]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX9-NEXT:    ret i32 [[A]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_lo_non_wave_multiple(
-; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
-; GFX10-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
-; GFX10-NEXT:    ret i32 [[A]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
 ; =============================================================================
 ; PARTIAL MASKS AND NON-STANDARD PATTERNS
 ; =============================================================================
 
 ; Test with partial mask
-define i32 @test_mbcnt_partial_mask() !reqd_work_group_size !1 {
+define i32 @test_mbcnt_partial_mask() !reqd_work_group_size !0 {
 ; GFX9-LABEL: define i32 @test_mbcnt_partial_mask(
-; GFX9-SAME: ) !reqd_work_group_size [[META1:![0-9]+]] {
+; GFX9-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
 ; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
 ; GFX9-NEXT:    ret i32 [[B]]
 ;
 ; GFX10-LABEL: define i32 @test_mbcnt_partial_mask(
-; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
+; GFX10-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
 ; GFX10-NEXT:  [[ENTRY:.*:]]
 ; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
 ; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
@@ -121,16 +76,16 @@ entry:
 }
 
 ; Test with non-zero base
-define i32 @test_mbcnt_non_zero_base() !reqd_work_group_size !1 {
+define i32 @test_mbcnt_non_zero_base() !reqd_work_group_size !0 {
 ; GFX9-LABEL: define i32 @test_mbcnt_non_zero_base(
-; GFX9-SAME: ) !reqd_work_group_size [[META1]] {
+; GFX9-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
 ; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
 ; GFX9-NEXT:    ret i32 [[B]]
 ;
 ; GFX10-LABEL: define i32 @test_mbcnt_non_zero_base(
-; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1]] {
+; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
 ; GFX10-NEXT:  [[ENTRY:.*:]]
 ; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
 ; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
@@ -142,30 +97,6 @@ entry:
   ret i32 %b
 }
 
-; =============================================================================
-; ARCHITECTURE-SPECIFIC NEGATIVE CASES
-; =============================================================================
-
-; Test with work group size = partial wave multiple (48)
-define i32 @test_mbcnt_lo_wave32_partial_wave() !reqd_work_group_size !0 {
-; GFX9-LABEL: define i32 @test_mbcnt_lo_wave32_partial_wave(
-; GFX9-SAME: ) !reqd_work_group_size [[META0]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX9-NEXT:    ret i32 [[A]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_lo_wave32_partial_wave(
-; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
-; GFX10-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
-; GFX10-NEXT:    ret i32 [[A]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
 ; =============================================================================
 ; COPY OPTIMIZATION NEGATIVE CASES
 ; =============================================================================
@@ -173,7 +104,7 @@ entry:
 ; Test with no work group size
 define i32 @test_mbcnt_hi_copy_no_wgs(i32 %val) {
 ; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_no_wgs(
-; GFX9-SAME: i32 [[VAL:%.*]]) {
+; GFX9-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
 ; GFX9-NEXT:    ret i32 [[RESULT]]
@@ -190,15 +121,15 @@ entry:
 }
 
 ; Test with work group size = not a wave multiple (48)
-define i32 @test_mbcnt_hi_copy_non_wave_multiple(i32 %val) !reqd_work_group_size !0 {
+define i32 @test_mbcnt_hi_copy_non_wave_multiple(i32 %val) !reqd_work_group_size !1 {
 ; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
-; GFX9-SAME: i32 [[VAL:%.*]]) !reqd_work_group_size [[META0]] {
+; GFX9-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
 ; GFX9-NEXT:    ret i32 [[RESULT]]
 ;
 ; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
-; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
+; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
 ; GFX10-NEXT:  [[ENTRY:.*:]]
 ; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
 ; GFX10-NEXT:    ret i32 [[RESULT]]
@@ -209,15 +140,15 @@ entry:
 }
 
 ; Test with zero mask
-define i32 @test_mbcnt_hi_copy_zero_mask(i32 %val) !reqd_work_group_size !1 {
+define i32 @test_mbcnt_hi_copy_zero_mask(i32 %val) !reqd_work_group_size !0 {
 ; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
-; GFX9-SAME: i32 [[VAL:%.*]]) !reqd_work_group_size [[META1]] {
+; GFX9-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
 ; GFX9-NEXT:  [[ENTRY:.*:]]
 ; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
 ; GFX9-NEXT:    ret i32 [[RESULT]]
 ;
 ; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
-; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META1]] {
+; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
 ; GFX10-NEXT:  [[ENTRY:.*:]]
 ; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
 ; GFX10-NEXT:    ret i32 [[RESULT]]
@@ -231,8 +162,8 @@ entry:
 ; METADATA
 ; =============================================================================
 
-!0 = !{i32 48, i32 1, i32 1}   ; X=48 (not wave multiple)
-!1 = !{i32 64, i32 1, i32 1}   ; X=64 (wave64 or 2*wave32)
+!0 = !{i32 64, i32 1, i32 1}   ; X=64 (wave64 or 2*wave32)
+!1 = !{i32 48, i32 1, i32 1}   ; X=48 (not wave multiple)
 
 ; =============================================================================
 ; FUNCTION DECLARATIONS
@@ -243,11 +174,11 @@ declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32) #0
 
 attributes #0 = { nounwind readnone speculatable willreturn }
 ;.
-; GFX9: [[META0]] = !{i32 48, i32 1, i32 1}
-; GFX9: [[META1]] = !{i32 64, i32 1, i32 1}
+; GFX9: [[META0]] = !{i32 64, i32 1, i32 1}
+; GFX9: [[META1]] = !{i32 48, i32 1, i32 1}
 ;.
-; GFX10: [[META0]] = !{i32 48, i32 1, i32 1}
-; GFX10: [[META1]] = !{i32 64, i32 1, i32 1}
+; GFX10: [[META0]] = !{i32 64, i32 1, i32 1}
+; GFX10: [[META1]] = !{i32 48, i32 1, i32 1}
 ;.
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; CHECK: {{.*}}
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll b/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
index 975d3f295e9db..e2c683ffd5562 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
@@ -8,7 +8,7 @@
 ; SIMPLE REPLACEMENT OPTIMIZATIONS (mbcnt.lo -> workitem.id.x)
 ; =============================================================================
 
-; Test with work group size = wave size
+; Test with work group size = wave size (32)
 define i32 @test_mbcnt_lo_simple_wave32() !reqd_work_group_size !0 {
 ; CHECK-LABEL: define i32 @test_mbcnt_lo_simple_wave32(
 ; CHECK-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
@@ -25,7 +25,7 @@ entry:
 ; BITMASK OPTIMIZATIONS (mbcnt.lo -> workitem.id.x & 0x1f)
 ; =============================================================================
 
-; Test with work group size = 2 * wave size
+; Test with work group size = 2 * wave size (64)
 define i32 @test_mbcnt_lo_bitmask_64() !reqd_work_group_size !1 {
 ; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_64(
 ; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
@@ -39,12 +39,12 @@ entry:
   ret i32 %a
 }
 
-; Test with work group size = 4 * wave size
-define i32 @test_mbcnt_lo_bitmask_128() !reqd_work_group_size !3 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_128(
+; Test with work group size = 3 * wave size (96)
+define i32 @test_mbcnt_lo_bitmask_96() !reqd_work_group_size !2 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_96(
 ; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META2:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 128) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 96) i32 @llvm.amdgcn.workitem.id.x()
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
@@ -53,12 +53,12 @@ entry:
   ret i32 %a
 }
 
-; Test with work group size = non-power-of-2 but wave-multiple
-define i32 @test_mbcnt_lo_bitmask_96() !reqd_work_group_size !2 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_96(
+; Test with work group size = 4 * wave size (128)
+define i32 @test_mbcnt_lo_bitmask_128() !reqd_work_group_size !3 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_128(
 ; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 96) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 128) i32 @llvm.amdgcn.workitem.id.x()
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
@@ -67,6 +67,20 @@ entry:
   ret i32 %a
 }
 
+; Test with work group size = 0.75 * wave size (48)
+define i32 @test_mbcnt_lo_bitmask_48() !reqd_work_group_size !4 {
+; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_48(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META4:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[A]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  ret i32 %a
+}
+
 ; =============================================================================
 ; COPY OPTIMIZATION (mbcnt.hi(mask, val) -> val)
 ; =============================================================================
@@ -99,7 +113,7 @@ entry:
 ; FULL PATTERN OPTIMIZATION (mbcnt.hi(~0, mbcnt.lo(~0, 0)) -> workitem.id.x)
 ; =============================================================================
 
-; Test full pattern on wave32
+; Test with work group size = wave size (32)
 define i32 @test_mbcnt_full_pattern_wave32() !reqd_work_group_size !0 {
 ; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave32(
 ; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
@@ -113,6 +127,22 @@ entry:
   ret i32 %b
 }
 
+; Test with work group size = 0.75 * wave size (48)
+define i32 @test_mbcnt_full_pattern_wave32_partial() !reqd_work_group_size !4 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave32_partial(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META4]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[B:%.*]] = and i32 [[TMP0]], 31
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
 ; =============================================================================
 ; METADATA
 ; =============================================================================
@@ -121,6 +151,7 @@ entry:
 !1 = !{i32 64, i32 1, i32 1}   ; X=64 (2*wave32), Y=1, Z=1
 !2 = !{i32 96, i32 1, i32 1}   ; X=96 (3*wave32), Y=1, Z=1
 !3 = !{i32 128, i32 1, i32 1}  ; X=128 (4*wave32), Y=1, Z=1
+!4 = !{i32 48, i32 1, i32 1}   ; X=48 (1.5*wave32), Y=1, Z=1
 
 ; =============================================================================
 ; FUNCTION DECLARATIONS
@@ -133,6 +164,7 @@ attributes #0 = { nounwind readnone speculatable willreturn }
 ;.
 ; CHECK: [[META0]] = !{i32 32, i32 1, i32 1}
 ; CHECK: [[META1]] = !{i32 64, i32 1, i32 1}
-; CHECK: [[META2]] = !{i32 128, i32 1, i32 1}
-; CHECK: [[META3]] = !{i32 96, i32 1, i32 1}
+; CHECK: [[META2]] = !{i32 96, i32 1, i32 1}
+; CHECK: [[META3]] = !{i32 128, i32 1, i32 1}
+; CHECK: [[META4]] = !{i32 48, i32 1, i32 1}
 ;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll b/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
index 42ba27a37b094..d70871b6eda3b 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
@@ -70,6 +70,22 @@ entry:
   ret i32 %b
 }
 
+; Test with work group size = 0.75 * wave size (48)
+define i32 @test_mbcnt_full_pattern_wave64_partial() !reqd_work_group_size !4 {
+; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_partial(
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META4:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[B:%.*]] = and i32 [[TMP0]], 63
+; CHECK-NEXT:    ret i32 [[B]]
+;
+entry:
+  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
+  ret i32 %b
+}
+
 ; =============================================================================
 ; METADATA
 ; =============================================================================
@@ -78,6 +94,7 @@ entry:
 !1 = !{i32 128, i32 1, i32 1}  ; X=128 (2*wave), Y=1, Z=1
 !2 = !{i32 256, i32 1, i32 1}  ; X=256 (4*wave), Y=1, Z=1
 !3 = !{i32 64, i32 2, i32 1}   ; X=64 (1*wave), Y=2, Z=1
+!4 = !{i32 48, i32 1, i32 1}   ; X=48 (0.75*wave), Y=1, Z=1
 
 ; =============================================================================
 ; FUNCTION DECLARATIONS
@@ -92,4 +109,5 @@ attributes #0 = { nounwind readnone speculatable willreturn }
 ; CHECK: [[META1]] = !{i32 128, i32 1, i32 1}
 ; CHECK: [[META2]] = !{i32 256, i32 1, i32 1}
 ; CHECK: [[META3]] = !{i32 64, i32 2, i32 1}
+; CHECK: [[META4]] = !{i32 48, i32 1, i32 1}
 ;.

>From 9b62e55c1232547cebb0a63397f1584f27a28802 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Wed, 29 Oct 2025 16:14:59 +0530
Subject: [PATCH 09/11] Remove redundant isPower2_32 check and combine 2 match
 functions in visitMbcntHi

---
 .../Target/AMDGPU/AMDGPUCodeGenPrepare.cpp    | 57 ++++++++-----------
 1 file changed, 25 insertions(+), 32 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index dbf07e6aa091f..51ef0681442c7 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -2129,20 +2129,18 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntLo(IntrinsicInst &I) const {
       ST.makeLIDRangeMetadata(NewCall);
       return true;
     }
-    // When work group evenly splits into waves and wave size is power-of-2,
-    // we can compute lane ID within wave using bit masking:
-    // lane_id = workitem.id.x & (wave_size - 1).
+    // When work group evenly splits into waves, we can compute lane ID within
+    // wave using bit masking: lane_id = workitem.id.x & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
-      if (isPowerOf2_32(Wave)) {
-        IRBuilder<> B(&I);
-        CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-        ST.makeLIDRangeMetadata(Tid);
-        Constant *Mask = ConstantInt::get(Tid->getType(), Wave - 1);
-        Value *AndInst = B.CreateAnd(Tid, Mask);
-        BasicBlock::iterator BI(&I);
-        ReplaceInstWithValue(BI, AndInst);
-        return true;
-      }
+      // Construct optimized sequence: workitem.id.x & (wave_size - 1)
+      IRBuilder<> B(&I);
+      CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+      ST.makeLIDRangeMetadata(Tid);
+      Constant *Mask = ConstantInt::get(Tid->getType(), Wave - 1);
+      Value *AndInst = B.CreateAnd(Tid, Mask);
+      BasicBlock::iterator BI(&I);
+      ReplaceInstWithValue(BI, AndInst);
+      return true;
     }
   }
 
@@ -2180,11 +2178,9 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) const {
   using namespace PatternMatch;
 
   // Check for pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0))
-  if (!match(I.getArgOperand(0), m_AllOnes()))
-    return false;
-
-  if (!match(I.getArgOperand(1),
-             m_Intrinsic<Intrinsic::amdgcn_mbcnt_lo>(m_AllOnes(), m_Zero())))
+  if (!match(&I, m_Intrinsic<Intrinsic::amdgcn_mbcnt_hi>(
+                     m_AllOnes(), m_Intrinsic<Intrinsic::amdgcn_mbcnt_lo>(
+                                      m_AllOnes(), m_Zero()))))
     return false;
 
   if (auto MaybeX = ST.getReqdWorkGroupSize(F, 0)) {
@@ -2200,21 +2196,18 @@ bool AMDGPUCodeGenPrepareImpl::visitMbcntHi(IntrinsicInst &I) const {
       ST.makeLIDRangeMetadata(NewCall);
       return true;
     }
-    // When work group evenly splits into waves and wave size is power-of-2,
-    // we can compute lane ID within wave using bit masking:
-    // lane_id = workitem.id.x & (wave_size - 1).
+    // When work group evenly splits into waves, we can compute lane ID within
+    // wave using bit masking: lane_id = workitem.id.x & (wave_size - 1).
     if (ST.hasWavefrontsEvenlySplittingXDim(F, /*RequiresUniformYZ=*/true)) {
-      if (isPowerOf2_32(Wave)) {
-        // Construct optimized sequence: workitem.id.x & (wave_size - 1)
-        IRBuilder<> B(&I);
-        CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
-        ST.makeLIDRangeMetadata(Tid);
-        Constant *Mask = ConstantInt::get(Tid->getType(), Wave - 1);
-        Value *AndInst = B.CreateAnd(Tid, Mask);
-        BasicBlock::iterator BI(&I);
-        ReplaceInstWithValue(BI, AndInst);
-        return true;
-      }
+      // Construct optimized sequence: workitem.id.x & (wave_size - 1)
+      IRBuilder<> B(&I);
+      CallInst *Tid = B.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
+      ST.makeLIDRangeMetadata(Tid);
+      Constant *Mask = ConstantInt::get(Tid->getType(), Wave - 1);
+      Value *AndInst = B.CreateAnd(Tid, Mask);
+      BasicBlock::iterator BI(&I);
+      ReplaceInstWithValue(BI, AndInst);
+      return true;
     }
   } else {
     // When ST.getReqdWorkGroupSize() fails, use metadata. And only optimize the

>From 7a58080a156dea9ed820fee416ccc22ae2dcb921 Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Wed, 29 Oct 2025 16:31:29 +0530
Subject: [PATCH 10/11] Removed 4*Wave case from the tests

---
 .../AMDGPU/mbcnt-wave32-optimizations.ll      | 28 ++++-----------
 .../AMDGPU/mbcnt-wave64-optimizations.ll      | 34 +++++--------------
 2 files changed, 14 insertions(+), 48 deletions(-)

diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll b/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
index e2c683ffd5562..1ec25da09d6a1 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-wave32-optimizations.ll
@@ -53,24 +53,10 @@ entry:
   ret i32 %a
 }
 
-; Test with work group size = 4 * wave size (128)
-define i32 @test_mbcnt_lo_bitmask_128() !reqd_work_group_size !3 {
-; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_128(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 128) i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 31
-; CHECK-NEXT:    ret i32 [[TMP1]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  ret i32 %a
-}
-
 ; Test with work group size = 0.75 * wave size (48)
-define i32 @test_mbcnt_lo_bitmask_48() !reqd_work_group_size !4 {
+define i32 @test_mbcnt_lo_bitmask_48() !reqd_work_group_size !3 {
 ; CHECK-LABEL: define i32 @test_mbcnt_lo_bitmask_48(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META4:![0-9]+]] {
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
 ; CHECK-NEXT:    [[A:%.*]] = and i32 [[TMP0]], 31
@@ -128,9 +114,9 @@ entry:
 }
 
 ; Test with work group size = 0.75 * wave size (48)
-define i32 @test_mbcnt_full_pattern_wave32_partial() !reqd_work_group_size !4 {
+define i32 @test_mbcnt_full_pattern_wave32_partial() !reqd_work_group_size !3 {
 ; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave32_partial(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META4]] {
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
@@ -150,8 +136,7 @@ entry:
 !0 = !{i32 32, i32 1, i32 1}   ; X=32 (1*wave32), Y=1, Z=1
 !1 = !{i32 64, i32 1, i32 1}   ; X=64 (2*wave32), Y=1, Z=1
 !2 = !{i32 96, i32 1, i32 1}   ; X=96 (3*wave32), Y=1, Z=1
-!3 = !{i32 128, i32 1, i32 1}  ; X=128 (4*wave32), Y=1, Z=1
-!4 = !{i32 48, i32 1, i32 1}   ; X=48 (1.5*wave32), Y=1, Z=1
+!3 = !{i32 48, i32 1, i32 1}   ; X=48 (1.5*wave32), Y=1, Z=1
 
 ; =============================================================================
 ; FUNCTION DECLARATIONS
@@ -165,6 +150,5 @@ attributes #0 = { nounwind readnone speculatable willreturn }
 ; CHECK: [[META0]] = !{i32 32, i32 1, i32 1}
 ; CHECK: [[META1]] = !{i32 64, i32 1, i32 1}
 ; CHECK: [[META2]] = !{i32 96, i32 1, i32 1}
-; CHECK: [[META3]] = !{i32 128, i32 1, i32 1}
-; CHECK: [[META4]] = !{i32 48, i32 1, i32 1}
+; CHECK: [[META3]] = !{i32 48, i32 1, i32 1}
 ;.
diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll b/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
index d70871b6eda3b..595ca0026cf3a 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-wave64-optimizations.ll
@@ -39,26 +39,10 @@ entry:
   ret i32 %b
 }
 
-; Test with work group size = 4 * wave size (256)
-define i32 @test_mbcnt_full_pattern_wave64_256() !reqd_work_group_size !2 {
-; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_256(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META2:![0-9]+]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 256) i32 @llvm.amdgcn.workitem.id.x()
-; CHECK-NEXT:    [[B:%.*]] = and i32 [[TMP0]], 63
-; CHECK-NEXT:    ret i32 [[B]]
-;
-entry:
-  %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-  %b = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %a)
-  ret i32 %b
-}
-
 ; Test with multidimensional work group where X dimension matches pattern
-define i32 @test_mbcnt_full_pattern_wave64_multidim() !reqd_work_group_size !3 {
+define i32 @test_mbcnt_full_pattern_wave64_multidim() !reqd_work_group_size !2 {
 ; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_multidim(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META2:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
@@ -71,9 +55,9 @@ entry:
 }
 
 ; Test with work group size = 0.75 * wave size (48)
-define i32 @test_mbcnt_full_pattern_wave64_partial() !reqd_work_group_size !4 {
+define i32 @test_mbcnt_full_pattern_wave64_partial() !reqd_work_group_size !3 {
 ; CHECK-LABEL: define i32 @test_mbcnt_full_pattern_wave64_partial(
-; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META4:![0-9]+]] {
+; CHECK-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META3:![0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
 ; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 0, 48) i32 @llvm.amdgcn.workitem.id.x()
@@ -92,9 +76,8 @@ entry:
 
 !0 = !{i32 64, i32 1, i32 1}   ; X=64 (1*wave), Y=1, Z=1
 !1 = !{i32 128, i32 1, i32 1}  ; X=128 (2*wave), Y=1, Z=1
-!2 = !{i32 256, i32 1, i32 1}  ; X=256 (4*wave), Y=1, Z=1
-!3 = !{i32 64, i32 2, i32 1}   ; X=64 (1*wave), Y=2, Z=1
-!4 = !{i32 48, i32 1, i32 1}   ; X=48 (0.75*wave), Y=1, Z=1
+!2 = !{i32 64, i32 2, i32 1}   ; X=64 (1*wave), Y=2, Z=1
+!3 = !{i32 48, i32 1, i32 1}   ; X=48 (0.75*wave), Y=1, Z=1
 
 ; =============================================================================
 ; FUNCTION DECLARATIONS
@@ -107,7 +90,6 @@ attributes #0 = { nounwind readnone speculatable willreturn }
 ;.
 ; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
 ; CHECK: [[META1]] = !{i32 128, i32 1, i32 1}
-; CHECK: [[META2]] = !{i32 256, i32 1, i32 1}
-; CHECK: [[META3]] = !{i32 64, i32 2, i32 1}
-; CHECK: [[META4]] = !{i32 48, i32 1, i32 1}
+; CHECK: [[META2]] = !{i32 64, i32 2, i32 1}
+; CHECK: [[META3]] = !{i32 48, i32 1, i32 1}
 ;.

>From 9b58595a89c58462b3afb89cfa04e47efaa70b8b Mon Sep 17 00:00:00 2001
From: TejaX-Alaghari <Teja.Alaghari at amd.com>
Date: Fri, 31 Oct 2025 12:15:14 +0530
Subject: [PATCH 11/11] Simplified negative cases test file as the arch
 specific cases are moved to a different file now

---
 .../Transforms/AMDGPU/mbcnt-negative-cases.ll | 131 ++++++------------
 1 file changed, 39 insertions(+), 92 deletions(-)

diff --git a/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll b/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
index 758f76358c771..0b9000138e5b3 100644
--- a/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
+++ b/llvm/test/Transforms/AMDGPU/mbcnt-negative-cases.ll
@@ -1,6 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx906 -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX9 %s
-; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -mcpu=gfx1030 -passes=amdgpu-codegenprepare < %s | FileCheck --check-prefixes=CHECK,GFX10 %s
+; RUN: opt -S -mtriple amdgcn-unknown-amdhsa -passes=amdgpu-codegenprepare < %s | FileCheck %s
 
 ; Test negative cases where mbcnt optimizations should NOT be applied
 
@@ -10,19 +9,11 @@
 
 ; Test with no reqd_work_group_size
 define i32 @test_mbcnt_no_work_group_size() {
-; GFX9-LABEL: define i32 @test_mbcnt_no_work_group_size(
-; GFX9-SAME: ) #[[ATTR0:[0-9]+]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX9-NEXT:    ret i32 [[B]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_no_work_group_size(
-; GFX10-SAME: ) #[[ATTR0:[0-9]+]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX10-NEXT:    ret i32 [[B]]
+; CHECK-LABEL: define i32 @test_mbcnt_no_work_group_size() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; CHECK-NEXT:    ret i32 [[B]]
 ;
 entry:
   %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
@@ -32,17 +23,10 @@ entry:
 
 ; Test mbcnt.lo with no work group size
 define i32 @test_mbcnt_lo_no_work_group_size() {
-; GFX9-LABEL: define i32 @test_mbcnt_lo_no_work_group_size(
-; GFX9-SAME: ) #[[ATTR0]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX9-NEXT:    ret i32 [[A]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_lo_no_work_group_size(
-; GFX10-SAME: ) #[[ATTR0]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
-; GFX10-NEXT:    ret i32 [[A]]
+; CHECK-LABEL: define i32 @test_mbcnt_lo_no_work_group_size() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
+; CHECK-NEXT:    ret i32 [[A]]
 ;
 entry:
   %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0)
@@ -55,19 +39,12 @@ entry:
 
 ; Test with partial mask
 define i32 @test_mbcnt_partial_mask() !reqd_work_group_size !0 {
-; GFX9-LABEL: define i32 @test_mbcnt_partial_mask(
-; GFX9-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
-; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX9-NEXT:    ret i32 [[B]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_partial_mask(
-; GFX10-SAME: ) #[[ATTR0:[0-9]+]] !reqd_work_group_size [[META0:![0-9]+]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
-; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX10-NEXT:    ret i32 [[B]]
+; CHECK-LABEL: define i32 @test_mbcnt_partial_mask(
+; CHECK-SAME: ) !reqd_work_group_size [[META0:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
+; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; CHECK-NEXT:    ret i32 [[B]]
 ;
 entry:
   %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 65535, i32 0)
@@ -77,19 +54,12 @@ entry:
 
 ; Test with non-zero base
 define i32 @test_mbcnt_non_zero_base() !reqd_work_group_size !0 {
-; GFX9-LABEL: define i32 @test_mbcnt_non_zero_base(
-; GFX9-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
-; GFX9-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX9-NEXT:    ret i32 [[B]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_non_zero_base(
-; GFX10-SAME: ) #[[ATTR0]] !reqd_work_group_size [[META0]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
-; GFX10-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
-; GFX10-NEXT:    ret i32 [[B]]
+; CHECK-LABEL: define i32 @test_mbcnt_non_zero_base(
+; CHECK-SAME: ) !reqd_work_group_size [[META0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
+; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[A]])
+; CHECK-NEXT:    ret i32 [[B]]
 ;
 entry:
   %a = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 5)
@@ -103,17 +73,11 @@ entry:
 
 ; Test with no work group size
 define i32 @test_mbcnt_hi_copy_no_wgs(i32 %val) {
-; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_no_wgs(
-; GFX9-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
-; GFX9-NEXT:    ret i32 [[RESULT]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_no_wgs(
-; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
-; GFX10-NEXT:    ret i32 [[RESULT]]
+; CHECK-LABEL: define i32 @test_mbcnt_hi_copy_no_wgs(
+; CHECK-SAME: i32 [[VAL:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
 ;
 entry:
   %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %val)
@@ -122,17 +86,11 @@ entry:
 
 ; Test with work group size = not a wave multiple (48)
 define i32 @test_mbcnt_hi_copy_non_wave_multiple(i32 %val) !reqd_work_group_size !1 {
-; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
-; GFX9-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
-; GFX9-NEXT:    ret i32 [[RESULT]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
-; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META1:![0-9]+]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
-; GFX10-NEXT:    ret i32 [[RESULT]]
+; CHECK-LABEL: define i32 @test_mbcnt_hi_copy_non_wave_multiple(
+; CHECK-SAME: i32 [[VAL:%.*]]) !reqd_work_group_size [[META1:![0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 [[VAL]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
 ;
 entry:
   %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %val)
@@ -141,17 +99,11 @@ entry:
 
 ; Test with zero mask
 define i32 @test_mbcnt_hi_copy_zero_mask(i32 %val) !reqd_work_group_size !0 {
-; GFX9-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
-; GFX9-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
-; GFX9-NEXT:  [[ENTRY:.*:]]
-; GFX9-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
-; GFX9-NEXT:    ret i32 [[RESULT]]
-;
-; GFX10-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
-; GFX10-SAME: i32 [[VAL:%.*]]) #[[ATTR0]] !reqd_work_group_size [[META0]] {
-; GFX10-NEXT:  [[ENTRY:.*:]]
-; GFX10-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
-; GFX10-NEXT:    ret i32 [[RESULT]]
+; CHECK-LABEL: define i32 @test_mbcnt_hi_copy_zero_mask(
+; CHECK-SAME: i32 [[VAL:%.*]]) !reqd_work_group_size [[META0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 [[VAL]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
 ;
 entry:
   %result = call i32 @llvm.amdgcn.mbcnt.hi(i32 0, i32 %val)
@@ -174,11 +126,6 @@ declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32) #0
 
 attributes #0 = { nounwind readnone speculatable willreturn }
 ;.
-; GFX9: [[META0]] = !{i32 64, i32 1, i32 1}
-; GFX9: [[META1]] = !{i32 48, i32 1, i32 1}
-;.
-; GFX10: [[META0]] = !{i32 64, i32 1, i32 1}
-; GFX10: [[META1]] = !{i32 48, i32 1, i32 1}
+; CHECK: [[META0]] = !{i32 64, i32 1, i32 1}
+; CHECK: [[META1]] = !{i32 48, i32 1, i32 1}
 ;.
-;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
-; CHECK: {{.*}}



More information about the llvm-commits mailing list