AArch64: use ldp/stp for 128-bit atomic load/store with v8.4

[t.p.northover]

Summary

Prior to v8.4a[*] the only way to get a 128-bit atomic load or store was via ldxp/stxp (or casp), which is not only inefficient but outright impossible without write access. So v8.4a extended the memory model so that any 16-byte operation aligned to 16 bytes (as all LLVM atomic load/stores must be) is atomic.

This patch implements ISel for these instructions in both SDAG and GISel. In both cases we go for ldp/stp implementations since atomics are much more likely to be GPR-based operations.

[t.p.northover]

Added separate feature to control it (named from ARMARM), since it's optional from v8.2 onwards.

[paquette]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17689

might as well make this an if?

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

1029

I think mi_match could really simplify this function.

if (mi_match(Root, MRI, m_GPtrAdd(m_Reg(NewBase), m_ICst(NewOffset))) && isShiftedInt<7, 3>(NewOffset))) {
  Base = NewBase;
  Offset = NewOffset;
}

... But I'm not sure about the >64 bit widths.

[efriedma]

Missing tests?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

915

Why not custom legalize? Not that it really matters, I guess.

16322

Assert that we have the expected atomic ordering?

16326

Would it make sense to add an AArch64ISD node to use here, instead of doing instruction selection early?

17469

What are you trying to accomplish here?

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

915

It gets converted to a cmpxchg during LegalizeTypes, and I don't think we can hook in there.

16326

Good idea. That should let us use SelectAddrMode... from ISelDAGToDAG instead of reinventing it here.

17469

Ah, I've got a separate patch to use CASP if LSE is present but not this LSE2. They started off together and I think this is a remnant of that bit. I'll remove it.

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

1029

Oh, neat! Didn't know that existed.

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16326

Actually, it looks like someone has beaten me to it with volatile load/store, and I was wrong about not being able to hook into LegalizeTypes. I should be able to reuse most of their work.

[t.p.northover]

• Added tests.
• Switched to existing SDAG LDP/STP route for volatiles.
• Simplified GISel path with matching expressions.
• Think I fixed everything else mentioned.

[efriedma]

LGTM

[t.p.northover]

Thanks. Committed as 13aa102e0769.

llvm/lib/Target/AArch64/AArch64.td

def FeatureCRC : SubtargetFeature<"crc", "HasCRC", "true",
   "Enable ARMv8 CRC-32 checksum instructions">;
 
 def FeatureRAS : SubtargetFeature<"ras", "HasRAS", "true",
   "Enable ARMv8 Reliability, Availability and Serviceability Extensions">;
 
 def FeatureLSE : SubtargetFeature<"lse", "HasLSE", "true",
   "Enable ARMv8.1 Large System Extension (LSE) atomic instructions">;
 
+def FeatureLSE2 : SubtargetFeature<"lse2", "HasLSE2", "true",
+  "Enable ARMv8.4 Large System Extension 2 (LSE2) atomicity rules">;
+
 def FeatureOutlineAtomics : SubtargetFeature<"outline-atomics", "OutlineAtomics", "true",
   "Enable out of line atomics to support LSE instructions">;
 
 def FeatureRDM : SubtargetFeature<"rdm", "HasRDM", "true",
   "Enable ARMv8.1 Rounding Double Multiply Add/Subtract instructions">;
 
 def FeaturePAN : SubtargetFeature<
     "pan", "HasPAN", "true",

 def HasV8_3aOps : SubtargetFeature<"v8.3a", "HasV8_3aOps", "true",
   "Support ARM v8.3a instructions", [HasV8_2aOps, FeatureRCPC, FeaturePAuth,
   FeatureJS, FeatureCCIDX, FeatureComplxNum]>;
 
 def HasV8_4aOps : SubtargetFeature<"v8.4a", "HasV8_4aOps", "true",
   "Support ARM v8.4a instructions", [HasV8_3aOps, FeatureDotProd,
   FeatureNV, FeatureMPAM, FeatureDIT,
   FeatureTRACEV8_4, FeatureAM, FeatureSEL2, FeatureTLB_RMI,
-  FeatureFlagM, FeatureRCPC_IMMO]>;
+  FeatureFlagM, FeatureRCPC_IMMO, FeatureLSE2]>;
 
 def HasV8_5aOps : SubtargetFeature<
   "v8.5a", "HasV8_5aOps", "true", "Support ARM v8.5a instructions",
   [HasV8_4aOps, FeatureAltFPCmp, FeatureFRInt3264, FeatureSpecRestrict,
    FeatureSSBS, FeatureSB, FeaturePredRes, FeatureCacheDeepPersist,
    FeatureBranchTargetId]>;
 
 def HasV8_6aOps : SubtargetFeature<

llvm/lib/Target/AArch64/AArch64ISelLowering.h

public:
 
   Value *emitLoadLinked(IRBuilderBase &Builder, Type *ValueTy, Value *Addr,
                         AtomicOrdering Ord) const override;
   Value *emitStoreConditional(IRBuilderBase &Builder, Value *Val, Value *Addr,
                               AtomicOrdering Ord) const override;
 
   void emitAtomicCmpXchgNoStoreLLBalance(IRBuilderBase &Builder) const override;
 
+  bool isOpSuitableForLDPSTP(const Instruction *I) const;
+  bool shouldInsertFencesForAtomic(const Instruction *I) const override;
+
   TargetLoweringBase::AtomicExpansionKind
   shouldExpandAtomicLoadInIR(LoadInst *LI) const override;
   bool shouldExpandAtomicStoreInIR(StoreInst *SI) const override;
   TargetLoweringBase::AtomicExpansionKind
   shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const override;
 
   TargetLoweringBase::AtomicExpansionKind
   shouldExpandAtomicCmpXchgInIR(AtomicCmpXchgInst *AI) const override;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setTargetDAGCombine(ISD::SIGN_EXTEND_INREG);
   setTargetDAGCombine(ISD::TRUNCATE);
   setTargetDAGCombine(ISD::CONCAT_VECTORS);
   setTargetDAGCombine(ISD::INSERT_SUBVECTOR);
   setTargetDAGCombine(ISD::STORE);
   if (Subtarget->supportsAddressTopByteIgnored())
     setTargetDAGCombine(ISD::LOAD);
 
+  // Aligned 128-bit loads and stores are single-copy atomic according to the
+  // v8.4a spec. Unfortunately i128 is not a legal type so the only opportunity
+  // we have to do anything with them is the very first DAG combine.
+  if (Subtarget->hasLSE2()) {
+    setTargetDAGCombine(ISD::ATOMIC_LOAD);
+    setTargetDAGCombine(ISD::ATOMIC_STORE);
+  }
+
   setTargetDAGCombine(ISD::MUL);
 
   setTargetDAGCombine(ISD::SELECT);
   setTargetDAGCombine(ISD::VSELECT);
 
   setTargetDAGCombine(ISD::INTRINSIC_VOID);
   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
   setTargetDAGCombine(ISD::INSERT_VECTOR_ELT);
if (!T->isSized() ||
     return SDValue();
 
   SDLoc DL(GN);
   SDValue Result = DAG.getGlobalAddress(GV, DL, MVT::i64, Offset);
   return DAG.getNode(ISD::SUB, DL, MVT::i64, Result,
                      DAG.getConstant(MinOffset, DL, MVT::i64));
 }
 
+static void matchLDPSTPAddrMode(SDValue Addr, SDValue &Base, SDValue &Offset,
+                                SelectionDAG &DAG) {
+  SDLoc DL(Addr);
+  Base = Addr;
+  Offset = DAG.getTargetConstant(0, DL, MVT::i32);
+
+  if (Addr.getOpcode() != ISD::ADD || !isa<ConstantSDNode>(Addr->getOperand(1)))
+    return;
+
+  int64_t Val = cast<ConstantSDNode>(Base->getOperand(1))->getSExtValue();
+  if (Val % 8 != 0 || !isInt<7>(Val / 8))
+    return;
+
+  Base = Base->getOperand(0);
+  Offset = DAG.getTargetConstant(Val / 8, DL, MVT::i32);
+}
+
+static SDValue performAtomic128Load(SDNode *N, SelectionDAG &DAG,
+                                    TargetLowering::DAGCombinerInfo &DCI) {
+  MemSDNode *MN = cast<MemSDNode>(N);
+  assert(MN->getAlignment() >= 16 && "ldp only works for aligned addresses");
+
+  SDLoc DL(N);
+  SDValue Base, Offset;
+  matchLDPSTPAddrMode(MN->getBasePtr(), Base, Offset, DAG);
+
+  SDNode *NewLoad =
+      DAG.getMachineNode(AArch64::LDPXi, DL, MVT::i64, MVT::i64, MVT::Other,
+                         Base, Offset, N->getOperand(0));
+  DAG.setNodeMemRefs(cast<MachineSDNode>(NewLoad), MN->getMemOperand());
+
+  SDValue NewVal = DAG.getNode(ISD::BUILD_PAIR, DL, MVT::i128,
+                               SDValue(NewLoad, 0), SDValue(NewLoad, 1));
+
+  DCI.CombineTo(N, {NewVal, SDValue(NewLoad, 2)});
+  return SDValue(N, 0);
+}
+
+static SDValue performAtomic128Store(SDNode *N, SelectionDAG &DAG,
+                                     TargetLowering::DAGCombinerInfo &DCI) {
+  MemSDNode *MN = cast<MemSDNode>(N);
+  assert(MN->getAlignment() >= 16 && "ldp only works for aligned addresses");
+
+  SDLoc DL(N);
+  SDValue Lo = DAG.getNode(ISD::EXTRACT_ELEMENT, DL, MVT::i64, N->getOperand(2),
+                           DAG.getIntPtrConstant(0, DL));
+  SDValue Hi = DAG.getNode(ISD::EXTRACT_ELEMENT, DL, MVT::i64, N->getOperand(2),
+                           DAG.getIntPtrConstant(1, DL));
+
+  SDValue Base, Offset;
+  matchLDPSTPAddrMode(MN->getBasePtr(), Base, Offset, DAG);
+
+  SDValue Ops[] = {Lo, Hi, Base, Offset, N->getOperand(0)};
+  SDNode *NewStore = DAG.getMachineNode(AArch64::STPXi, DL, MVT::Other, Ops);
+  DAG.setNodeMemRefs(cast<MachineSDNode>(NewStore), MN->getMemOperand());
+
+  return SDValue(NewStore, 0);
+}
 // Turns the vector of indices into a vector of byte offstes by scaling Offset
 // by (BitWidth / 8).
 static SDValue getScaledOffsetForBitWidth(SelectionDAG &DAG, SDValue Offset,
                                           SDLoc DL, unsigned BitWidth) {
   assert(Offset.getValueType().isScalableVector() &&
          "This method is only for scalable vectors of offsets");
 
   SDValue Shift = DAG.getConstant(Log2_32(BitWidth / 8), DL, MVT::i64);
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case ISD::INSERT_SUBVECTOR:
     return performInsertSubvectorCombine(N, DCI, DAG);
   case ISD::SELECT:
     return performSelectCombine(N, DCI);
   case ISD::VSELECT:
     return performVSelectCombine(N, DCI.DAG);
   case ISD::SETCC:
     return performSETCCCombine(N, DAG);
+  case ISD::ATOMIC_LOAD:
+    if (N->getValueType(0) == MVT::i128)
+      return performAtomic128Load(N, DAG, DCI);
+    break;
+  case ISD::ATOMIC_STORE:
+    if (N->getOperand(2)->getValueType(0) == MVT::i128)
+      return performAtomic128Store(N, DAG, DCI);
+    break;
   case ISD::LOAD:
     if (performTBISimplification(N->getOperand(1), DCI, DAG))
       return SDValue(N, 0);
     break;
   case ISD::STORE:
     return performSTORECombine(N, DCI, DAG, Subtarget);
   case ISD::VECTOR_SPLICE:
     return performSVESpliceCombine(N, DAG);
SDValue Ops[] = {
         createGPRPairNode(DAG, N->getOperand(3)), // Store value
         N->getOperand(1), // Ptr
         N->getOperand(0), // Chain in
     };
 
     unsigned Opcode;
     switch (MemOp->getMergedOrdering()) {
     case AtomicOrdering::Monotonic:
+    case AtomicOrdering::Unordered:
       Opcode = AArch64::CASPX;
       break;
     case AtomicOrdering::Acquire:
       Opcode = AArch64::CASPAX;
       break;
     case AtomicOrdering::Release:
       Opcode = AArch64::CASPLX;
       break;
AArch64TargetLowering::getPreferredVectorAction(MVT VT) const {
   // v4i16, v2i32 instead of to promote.
   if (VT == MVT::v1i8 || VT == MVT::v1i16 || VT == MVT::v1i32 ||
       VT == MVT::v1f32)
     return TypeWidenVector;
 
   return TargetLoweringBase::getPreferredVectorAction(VT);
 }
 
+// In v8.4a, ldp and stp instructions are guaranteed to be single-copy atomic
+// provided the address is 16-byte aligned.
+bool AArch64TargetLowering::isOpSuitableForLDPSTP(const Instruction *I) const {
+  if (!Subtarget->hasLSE2())
+    return false;
+
+  if (auto LI = dyn_cast<LoadInst>(I))
+    return LI->getType()->getPrimitiveSizeInBits() == 128 &&
+           LI->getAlignment() >= 16;
+  else if (auto SI = dyn_cast<StoreInst>(I))
+    return SI->getValueOperand()->getType()->getPrimitiveSizeInBits() == 128 &&
+           SI->getAlignment() >= 16;
+
+  return false;
+}
+
+bool AArch64TargetLowering::shouldInsertFencesForAtomic(
+    const Instruction *I) const {
+  return isOpSuitableForLDPSTP(I);
+}
+
 // Loads and stores less than 128-bits are already atomic; ones above that
 // are doomed anyway, so defer to the default libcall and blame the OS when
 // things go wrong.
 bool AArch64TargetLowering::shouldExpandAtomicStoreInIR(StoreInst *SI) const {
   unsigned Size = SI->getValueOperand()->getType()->getPrimitiveSizeInBits();
-  return Size == 128;
+  if (Size != 128)
+    return false;
+
+  return !isOpSuitableForLDPSTP(SI);
 }
 
 // Loads and stores less than 128-bits are already atomic; ones above that
 // are doomed anyway, so defer to the default libcall and blame the OS when
 // things go wrong.
 TargetLowering::AtomicExpansionKind
 AArch64TargetLowering::shouldExpandAtomicLoadInIR(LoadInst *LI) const {
   unsigned Size = LI->getType()->getPrimitiveSizeInBits();
-  return Size == 128 ? AtomicExpansionKind::LLSC : AtomicExpansionKind::None;
+
+  if (Size != 128 || isOpSuitableForLDPSTP(LI))
+    return AtomicExpansionKind::None;
+
+  return AtomicExpansionKind::LLSC;
 }
 
 // For the real atomic operations, we have ldxr/stxr up to 128 bits,
 TargetLowering::AtomicExpansionKind
 AArch64TargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
   if (AI->isFloatingPointOperation())
     return AtomicExpansionKind::CmpXChg;
 

llvm/lib/Target/AArch64/AArch64Subtarget.h

protected:
   bool HasCONTEXTIDREL2 = false;
 
   bool HasFPARMv8 = false;
   bool HasNEON = false;
   bool HasCrypto = false;
   bool HasDotProd = false;
   bool HasCRC = false;
   bool HasLSE = false;
+  bool HasLSE2 = false;
   bool HasRAS = false;
   bool HasRDM = false;
   bool HasPerfMon = false;
   bool HasFullFP16 = false;
   bool HasFP16FML = false;
   bool HasSPE = false;
 
   // ARMv8.1 extensions
public:
   }
   bool hasCustomCallingConv() const { return CustomCallSavedXRegs.any(); }
   bool hasFPARMv8() const { return HasFPARMv8; }
   bool hasNEON() const { return HasNEON; }
   bool hasCrypto() const { return HasCrypto; }
   bool hasDotProd() const { return HasDotProd; }
   bool hasCRC() const { return HasCRC; }
   bool hasLSE() const { return HasLSE; }
+  bool hasLSE2() const { return HasLSE2; }
   bool hasRAS() const { return HasRAS; }
   bool hasRDM() const { return HasRDM; }
   bool hasSM4() const { return HasSM4; }
   bool hasSHA3() const { return HasSHA3; }
   bool hasSHA2() const { return HasSHA2; }
   bool hasAES() const { return HasAES; }
   bool hasCONTEXTIDREL2() const { return HasCONTEXTIDREL2; }
   bool balanceFPOps() const { return BalanceFPOps; }

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

case AArch64::FPRRegBankID:
     case 8:
       return isStore ? AArch64::STRBui : AArch64::LDRBui;
     case 16:
       return isStore ? AArch64::STRHui : AArch64::LDRHui;
     case 32:
       return isStore ? AArch64::STRSui : AArch64::LDRSui;
     case 64:
       return isStore ? AArch64::STRDui : AArch64::LDRDui;
+    case 128:
+      return isStore ? AArch64::STRQui : AArch64::LDRQui;
     }
     break;
   }
   return GenericOpc;
 }
 
 #ifndef NDEBUG
 /// Helper function that verifies that we have a valid copy at the end of

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

auto IsPtrVecPred = [=](const LegalityQuery &Query) {
     const LLT &ValTy = Query.Types[0];
     if (!ValTy.isVector())
       return false;
     const LLT EltTy = ValTy.getElementType();
     return EltTy.isPointer() && EltTy.getAddressSpace() == 0;
   };
 
   getActionDefinitionsBuilder(G_LOAD)
+      .customIf([=](const LegalityQuery &Query) {
+        return Query.Types[0] == s128 &&
+               Query.MMODescrs[0].Ordering != AtomicOrdering::NotAtomic;
+      })
       .legalForTypesWithMemDesc({{s8, p0, s8, 8},
                                  {s16, p0, s16, 8},
                                  {s32, p0, s32, 8},
                                  {s64, p0, s64, 8},
                                  {p0, p0, s64, 8},
                                  {s128, p0, s128, 8},
                                  {v8s8, p0, s64, 8},
                                  {v16s8, p0, s128, 8},
                                  {v4s16, p0, s64, 8},
                                  {v8s16, p0, s128, 8},
                                  {v2s32, p0, s64, 8},
                                  {v4s32, p0, s128, 8},
                                  {v2s64, p0, s128, 8}})
       // These extends are also legal
       .legalForTypesWithMemDesc({{s32, p0, s8, 8}, {s32, p0, s16, 8}})
       .widenScalarToNextPow2(0, /* MinSize = */8)
       .lowerIfMemSizeNotPow2()
       .clampScalar(0, s8, s64)
       .narrowScalarIf([=](const LegalityQuery &Query) {
         // Clamp extending load results to 32-bits.
         return Query.Types[0].isScalar() &&
           Query.Types[0] != Query.MMODescrs[0].MemoryTy &&
           Query.Types[0].getSizeInBits() > 32;
         },
         changeTo(0, s32))
       // Lower any any-extending loads left into G_ANYEXT and G_LOAD
       .lowerIf([=](const LegalityQuery &Query) {
         return Query.Types[0] != Query.MMODescrs[0].MemoryTy;
       })
       .clampMaxNumElements(0, s8, 16)
       .clampMaxNumElements(0, s16, 8)
       .clampMaxNumElements(0, s32, 4)
       .clampMaxNumElements(0, s64, 2)
       .clampMaxNumElements(0, p0, 2)
       .customIf(IsPtrVecPred)
       .scalarizeIf(typeIs(0, v2s16), 0);
 
   getActionDefinitionsBuilder(G_STORE)
+      .customIf([=](const LegalityQuery &Query) {
+        return Query.Types[0] == s128 &&
+               Query.MMODescrs[0].Ordering != AtomicOrdering::NotAtomic;
+      })
       .legalForTypesWithMemDesc({{s8, p0, s8, 8},
                                  {s16, p0, s8, 8}, // truncstorei8 from s16
                                  {s32, p0, s8, 8}, // truncstorei8 from s32
                                  {s64, p0, s8, 8}, // truncstorei8 from s64
                                  {s16, p0, s16, 8},
                                  {s32, p0, s16, 8}, // truncstorei16 from s32
                                  {s64, p0, s16, 8}, // truncstorei16 from s64
                                  {s32, p0, s8, 8},
                                  {s32, p0, s16, 8},
if (Amount > 31)
     return true; // This will have to remain a register variant.
   auto ExtCst = MIRBuilder.buildConstant(LLT::scalar(64), Amount);
   Observer.changingInstr(MI);
   MI.getOperand(2).setReg(ExtCst.getReg(0));
   Observer.changedInstr(MI);
   return true;
 }
 
+static void matchLDPSTPAddrMode(Register Root, Register &Base, int &Offset,
+                                MachineRegisterInfo &MRI) {
+  Base = Root;
+  Offset = 0;
+
+  MachineInstr *RootDef = MRI.getVRegDef(Root);
+  if (!RootDef || RootDef->getOpcode() != TargetOpcode::G_PTR_ADD)
+    return;
+
+  MachineOperand &OffImm = RootDef->getOperand(2);
+  if (!OffImm.isReg())
+    return;
+  MachineInstr *RHS = MRI.getVRegDef(OffImm.getReg());
+  if (!RHS || RHS->getOpcode() != TargetOpcode::G_CONSTANT)
+    return;
+  int64_t RHSC;
+  MachineOperand &RHSOp1 = RHS->getOperand(1);
+  if (!RHSOp1.isCImm() || RHSOp1.getCImm()->getBitWidth() > 64)
+    return;
+
+  RHSC = RHSOp1.getCImm()->getSExtValue();
+  if (isShiftedInt<7, 3>(RHSC)) {
+    Base = RootDef->getOperand(1).getReg();
+    Offset = RHSC;
+  }
+}
+
 // FIXME: This should be removed and replaced with the generic bitcast legalize
 // action.
 bool AArch64LegalizerInfo::legalizeLoadStore(
     MachineInstr &MI, MachineRegisterInfo &MRI, MachineIRBuilder &MIRBuilder,
     GISelChangeObserver &Observer) const {
   assert(MI.getOpcode() == TargetOpcode::G_STORE ||
          MI.getOpcode() == TargetOpcode::G_LOAD);
   // Here we just try to handle vector loads/stores where our value type might
   // have pointer elements, which the SelectionDAG importer can't handle. To
   // allow the existing patterns for s64 to fire for p0, we just try to bitcast
   // the value to use s64 types.
 
   // Custom legalization requires the instruction, if not deleted, must be fully
   // legalized. In order to allow further legalization of the inst, we create
   // a new instruction and erase the existing one.
 
   Register ValReg = MI.getOperand(0).getReg();
   const LLT ValTy = MRI.getType(ValReg);
 
+  if (ValTy == LLT::scalar(128)) {
+    assert((*MI.memoperands_begin())->getSuccessOrdering() ==
+               AtomicOrdering::Monotonic ||
+           (*MI.memoperands_begin())->getSuccessOrdering() ==
+               AtomicOrdering::Unordered);
+    assert(ST->hasLSE2() && "ldp/stp not single copy atomic without +lse2");
+    LLT s64 = LLT::scalar(64);
+    MachineInstrBuilder NewI;
+    if (MI.getOpcode() == TargetOpcode::G_LOAD) {
+      NewI = MIRBuilder.buildInstr(AArch64::LDPXi, {s64, s64}, {});
+      MIRBuilder.buildMerge(ValReg, {NewI->getOperand(0), NewI->getOperand(1)});
+    } else {
+      auto Split = MIRBuilder.buildUnmerge(s64, MI.getOperand(0));
+      NewI = MIRBuilder.buildInstr(
+          AArch64::STPXi, {}, {Split->getOperand(0), Split->getOperand(1)});
+    }
+    Register Base;
+    int Offset;
+    matchLDPSTPAddrMode(MI.getOperand(1).getReg(), Base, Offset, MRI);
+    NewI.addUse(Base);
+    NewI.addImm(Offset / 8);
+
+    NewI.cloneMemRefs(MI);
+    constrainSelectedInstRegOperands(*NewI, *ST->getInstrInfo(),
+                                     *MRI.getTargetRegisterInfo(),
+                                     *ST->getRegBankInfo());
+    MI.eraseFromParent();
+    return true;
+  }
+
   if (!ValTy.isVector() || !ValTy.getElementType().isPointer() ||
       ValTy.getElementType().getAddressSpace() != 0) {
     LLVM_DEBUG(dbgs() << "Tried to do custom legalization on wrong load/store");
     return false;
   }
 
   unsigned PtrSize = ValTy.getElementType().getSizeInBits();
   const LLT NewTy = LLT::vector(ValTy.getElementCount(), PtrSize);

llvm/test/CodeGen/AArch64/GlobalISel/arm64-atomic-128.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=arm64-linux-gnu -verify-machineinstrs -global-isel -global-isel-abort=1 | FileCheck %s --check-prefix=CHECK-LLSC-O1
-; RUN: llc < %s -mtriple=arm64-linux-gnu -verify-machineinstrs -mcpu=apple-a13 -global-isel -global-isel-abort=1 | FileCheck %s --check-prefix=CHECK-CAS-O1
+; RUN: llc < %s -mtriple=arm64-linux-gnu -verify-machineinstrs -mattr=+lse -global-isel -global-isel-abort=1 | FileCheck %s --check-prefix=CHECK-CAS-O1
 ; RUN: llc < %s -mtriple=arm64-linux-gnu -verify-machineinstrs -O0 -global-isel -global-isel-abort=1 | FileCheck %s --check-prefix=CHECK-LLSC-O0
-; RUN: llc < %s -mtriple=arm64-linux-gnu -verify-machineinstrs -O0 -mcpu=apple-a13 -global-isel -global-isel-abort=1 | FileCheck %s --check-prefix=CHECK-CAS-O0
+; RUN: llc < %s -mtriple=arm64-linux-gnu -verify-machineinstrs -O0 -mattr=+lse -global-isel -global-isel-abort=1 | FileCheck %s --check-prefix=CHECK-CAS-O0
 @var = global i128 0
 
 define void @val_compare_and_swap(i128* %p, i128 %oldval, i128 %newval) {
 ; CHECK-LLSC-O1-LABEL: val_compare_and_swap:
 ; CHECK-LLSC-O1:       // %bb.0:
 ; CHECK-LLSC-O1-NEXT:  .LBB0_1: // =>This Inner Loop Header: Depth=1
 ; CHECK-LLSC-O1-NEXT:    ldaxp x8, x9, [x0]
 ; CHECK-LLSC-O1-NEXT:    cmp x8, x2

 ; CHECK-CAS-O0-NEXT:    .cfi_def_cfa_offset 48
 ; CHECK-CAS-O0-NEXT:    str x2, [sp, #32] // 8-byte Folded Spill
 ; CHECK-CAS-O0-NEXT:    str x3, [sp, #40] // 8-byte Folded Spill
 ; CHECK-CAS-O0-NEXT:    b .LBB4_1
 ; CHECK-CAS-O0-NEXT:  .LBB4_1: // %atomicrmw.start
 ; CHECK-CAS-O0-NEXT:    // =>This Inner Loop Header: Depth=1
 ; CHECK-CAS-O0-NEXT:    ldr x11, [sp, #32] // 8-byte Folded Reload
 ; CHECK-CAS-O0-NEXT:    ldxp x9, x10, [x11]
-; CHECK-CAS-O0-NEXT:    mov x8, #0
+; CHECK-CAS-O0-NEXT:    mov x8, xzr
 ; CHECK-CAS-O0-NEXT:    orr x9, x9, x8
 ; CHECK-CAS-O0-NEXT:    orr x10, x8, x10
 ; CHECK-CAS-O0-NEXT:    // implicit-def: $q0
 ; CHECK-CAS-O0-NEXT:    mov v0.d[0], x9
 ; CHECK-CAS-O0-NEXT:    str q0, [sp] // 16-byte Folded Spill
 ; CHECK-CAS-O0-NEXT:    mov v0.d[1], x10
 ; CHECK-CAS-O0-NEXT:    str q0, [sp, #16] // 16-byte Folded Spill
 ; CHECK-CAS-O0-NEXT:    stxp w8, x9, x10, [x11]

llvm/test/CodeGen/AArch64/atomic-ops-lse.ll

 ; RUN: llc -mtriple=aarch64-none-linux-gnu -disable-post-ra -verify-machineinstrs -mattr=+lse < %s | FileCheck %s
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -disable-post-ra -verify-machineinstrs -mattr=+lse -mattr=+outline-atomics < %s | FileCheck %s
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -disable-post-ra -verify-machineinstrs -mattr=+outline-atomics < %s | FileCheck %s --check-prefix=OUTLINE-ATOMICS
 ; RUN: llc -mtriple=aarch64_be-none-linux-gnu -disable-post-ra -verify-machineinstrs -mattr=+lse < %s | FileCheck %s
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -disable-post-ra -verify-machineinstrs -mattr=+lse < %s | FileCheck %s --check-prefix=CHECK-REG
-; RUN: llc -mtriple=aarch64-none-linux-gnu -disable-post-ra -verify-machineinstrs -mcpu=saphira < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-none-linux-gnu -disable-post-ra -verify-machineinstrs -mcpu=saphira -mattr=-lse2 < %s | FileCheck %s
 
 ; Point of CHECK-REG is to make sure UNPREDICTABLE instructions aren't created
 ; (i.e. reusing a register for status & data in store exclusive).
 ; CHECK-REG-NOT: stlxrb w[[NEW:[0-9]+]], w[[NEW]], [x{{[0-9]+}}]
 ; CHECK-REG-NOT: stlxrb w[[NEW:[0-9]+]], x[[NEW]], [x{{[0-9]+}}]
 
 @var8 = dso_local global i8 0
 @var16 = dso_local global i16 0