[AArch64][CGP] Sink zext feeding stxr/stlxr into the same block.

[ab]

Summary

The usual CodeGenPrepare trickery, this time on a target-specific intrinsic.

Without this, the expansion of atomics will usually have the zext be hoisted out of the loop, defeating the various patterns we have to catch precisely this case.

[qcolombet]

Hi Ahmed,

LGTM.

Thanks,
-Quentin

lib/CodeGen/CodeGenPrepare.cpp

if (II) {
     case Intrinsic::masked_store: {
       if (!TTI->isLegalMaskedStore(CI->getArgOperand(0)->getType(), 1)) {
         ScalarizeMaskedStore(CI);
         ModifiedDT = true;
         return true;
       }
       return false;
     }
+    case Intrinsic::aarch64_stlxr:
+    case Intrinsic::aarch64_stxr: {
+      ZExtInst *ExtVal = dyn_cast<ZExtInst>(CI->getArgOperand(0));
+      if (!ExtVal || !ExtVal->hasOneUse() ||
+          ExtVal->getParent() == CI->getParent())
+        return false;
+      // Sink a zext feeding stlxr/stxr before it, so it can be folded into it.
+      ExtVal->moveBefore(CI);
+      return true;
+    }
     }
 
     if (TLI) {
       SmallVector<Value*, 2> PtrOps;
       Type *AccessTy;
       if (TLI->GetAddrModeArguments(II, PtrOps, AccessTy))
         while (!PtrOps.empty())
           if (OptimizeMemoryInst(II, PtrOps.pop_back_val(), AccessTy))

test/CodeGen/AArch64/arm64-atomic.ll

 ; RUN: llc < %s -march=arm64 -verify-machineinstrs -mcpu=cyclone | FileCheck %s
 
 define i32 @val_compare_and_swap(i32* %p, i32 %cmp, i32 %new) {
 ; CHECK-LABEL: val_compare_and_swap:
-; CHECK: ubfx   x[[NEWVAL_REG:[0-9]+]], x2, #0, #32
 ; CHECK: [[LABEL:.?LBB[0-9]+_[0-9]+]]:
 ; CHECK: ldaxr   [[RESULT:w[0-9]+]], [x0]
 ; CHECK: cmp    [[RESULT]], w1
 ; CHECK: b.ne   [[LABEL2:.?LBB[0-9]+_[0-9]+]]
-; CHECK: stxr   [[SCRATCH_REG:w[0-9]+]], w[[NEWVAL_REG]], [x0]
+; CHECK: stxr   [[SCRATCH_REG:w[0-9]+]], w2, [x0]
 ; CHECK: cbnz   [[SCRATCH_REG]], [[LABEL]]
 ; CHECK: [[LABEL2]]:
   %pair = cmpxchg i32* %p, i32 %cmp, i32 %new acquire acquire
   %val = extractvalue { i32, i1 } %pair, 0
   ret i32 %val
 }
 
 define i32 @val_compare_and_swap_rel(i32* %p, i32 %cmp, i32 %new) {
 ; CHECK-LABEL: val_compare_and_swap_rel:
-; CHECK: ubfx   x[[NEWVAL_REG:[0-9]+]], x2, #0, #32
 ; CHECK: [[LABEL:.?LBB[0-9]+_[0-9]+]]:
 ; CHECK: ldaxr  [[RESULT:w[0-9]+]], [x0]
 ; CHECK: cmp    [[RESULT]], w1
 ; CHECK: b.ne   [[LABEL2:.?LBB[0-9]+_[0-9]+]]
-; CHECK: stlxr  [[SCRATCH_REG:w[0-9]+]], w[[NEWVAL_REG]], [x0]
+; CHECK: stlxr  [[SCRATCH_REG:w[0-9]+]], w2, [x0]
 ; CHECK: cbnz   [[SCRATCH_REG]], [[LABEL]]
 ; CHECK: [[LABEL2]]:
   %pair = cmpxchg i32* %p, i32 %cmp, i32 %new acq_rel monotonic
   %val = extractvalue { i32, i1 } %pair, 0
   ret i32 %val
 }
 
 define i64 @val_compare_and_swap_64(i64* %p) {