[llvm] [AtomicExpand] Add elementwise expansion (PR #196464)

[Antonio Frighetto via llvm-commits]

================
@@ -7759,30 +7762,75 @@ NVPTXTargetLowering::shouldExpandAtomicRMWInIR(const AtomicRMWInst *AI) const {
   return AtomicExpansionKind::CmpXChg;
 }
 
+NVPTXTargetLowering::AtomicExpansionKind
+NVPTXTargetLowering::shouldExpandAtomicRMWInIR(const AtomicRMWInst *AI) const {
+  // TODO: once we support native elementwise vector atoms,
+  // return `AtomicExpansionKind::None` to preserve and lower them.
+  if (AI->isElementwise()) {
+    auto *VecTy = cast<FixedVectorType>(AI->getType());
+    Type *LaneTy = VecTy->getElementType();
+
+    // Collapse <1 x T> elementwise RMWs to scalar RMWs
+    if (VecTy->getNumElements() == 1)
+      return AtomicExpansionKind::Expand;
+
+    // If the scalar lane op is natively supported, return Expand so halving
+    // eventually bottoms out at the scalar base case, where this hook returns
+    // None for each scalar lane and the scalar `atom.*` instruction is
+    // preserved.
+    if (getScalarAtomicRMWExpansion(AI, LaneTy, STI) ==
----------------
antoniofrighetto wrote:

It looks like we wouldn't perform bitcasting to integer anymore now (as we no longer go through shouldCastAtomicRMWIInIR()). Should we take this into account before deciding whether expanding the op into a cmpxchg loop or not (getScalarAtomicRMWExpansion() handles only integer types)?
```llvm
; ./bin/opt -mtriple=nvptx64-nvidia-cuda -passes='require<libcall-lowering-info>,atomic-expand'
define <2 x float> @test(ptr %p, <2 x float> %v) {
  %rv = atomicrmw elementwise xchg ptr %p, <2 x float> %v monotonic
  ret <2 x float> %rv
}
```
```
opt: /work/llvm-project/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp:7698:
TargetLoweringBase::AtomicExpansionKind getScalarAtomicRMWExpansion(const
AtomicRMWInst *, Type *, const NVPTXSubtarget &): Assertion `Ty->isIntegerTy()
&& "Ty should be integer at this point"' failed.
```

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