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

[Teja Alaghari via llvm-commits]

TejaX-Alaghari wrote:

Hey @arsenm,

> I'm not sure instcombine is the right place for this. Is there an instcombine run after AMDGPU attributor?

You are absolutely correct! This was the fundamental issue with the original implementation. After analyzing the LLVM optimization pipeline, I discovered that:

- `InstCombine` runs early and repeatedly throughout the O2 pipeline
- AMDGPU Attributor runs at the very end of the pipeline
- No `InstCombine` passes occur after AMDGPU Attributor

This pipeline timing meant the optimization could never practically trigger in real-world code, which is why I moved it to `AMDGPUCodeGenPrepare`. This pass:

- Runs after AMDGPU Attributor in the pipeline
- Has proper access to `reqd_work_group_size` metadata when it's needed
- Is specifically designed for target-specific IR optimizations that depend on metadata

> "Do we actually propagate reqd_work_group_size in attributor?"

Yes, the AMDGPU Attributor does handle `reqd_work_group_size` metadata, though you're right to note it's more complex than typical attributes:

How it works:
1. Clang CodeGen sets the metadata initially from source attributes:

```cpp
// From CodeGenFunction.cpp  
Fn->setMetadata("reqd_work_group_size", llvm::MDNode::get(Context, AttrMDArgs));
```
2. AMDGPU Attributor processes and potentially refines this metadata during optimization
3. AMDGPUCodeGenPrepare accesses it via the subtarget helper:
```cpp
if (auto MaybeX = ST.getReqdWorkGroupSize(*F, 0)) {
  // Optimization can now safely proceed
}
```
The key insight is that the metadata needs to be stable and available when the optimization runs, which is guaranteed in CodeGenPrepare but not in InstCombine.


Please let me know if you have further concerns or suggestions on the current implementation, thanks!

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