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

[Chris Jackson via llvm-commits]

================
@@ -1312,9 +1317,122 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
     break;
   }
   case Intrinsic::amdgcn_mbcnt_hi: {
-    // exec_hi is all 0, so this is just a copy.
-    if (ST->isWave32())
+    // exec_hi is all 0, so this is just a copy on wave32.
+    if (ST && ST->isWave32())
       return IC.replaceInstUsesWith(II, II.getArgOperand(1));
+
+    // Pattern: mbcnt.hi(~0, mbcnt.lo(~0, 0))
+    if (auto *HiArg1 = dyn_cast<CallInst>(II.getArgOperand(1))) {
+      Function *CalledF = HiArg1->getCalledFunction();
+      bool IsMbcntLo = false;
+      if (CalledF) {
+        // Fast-path: if this is a declared intrinsic, check the intrinsic ID.
+        if (CalledF->getIntrinsicID() == Intrinsic::amdgcn_mbcnt_lo) {
+          IsMbcntLo = true;
+        } else {
+          // Fallback: accept a declared function with the canonical name, but
+          // verify its signature to be safe: i32(i32,i32). Use the name
+          // comparison only when there's no intrinsic ID match.
+          if (CalledF->getName() == "llvm.amdgcn.mbcnt.lo") {
+            if (FunctionType *FT = CalledF->getFunctionType()) {
+              if (FT->getNumParams() == 2 &&
+                  FT->getReturnType()->isIntegerTy(32) &&
+                  FT->getParamType(0)->isIntegerTy(32) &&
+                  FT->getParamType(1)->isIntegerTy(32))
+                IsMbcntLo = true;
+            }
+          }
+        }
+      }
+
+      if (!IsMbcntLo)
+        break;
+
+      // hi arg0 must be all-ones
+      if (auto *HiArg0C = dyn_cast<ConstantInt>(II.getArgOperand(0))) {
----------------
chrisjbris wrote:

```suggestion
      auto *HiArg0C = dyn_cast<ConstantInt>(II.getArgOperand(0));
      if(!HiArg0C || (HiArg0C && !HiArg0C->isAllOnesValue()))
          break;
```

https://github.com/llvm/llvm-project/pull/160496