[AArch64] Improve codegen of volatile load/store of i128

[vhscampos]

Summary

Instead of generating two i64 instructions for each load or store of a
volatile i128 value (two LDRs or STRs), now emit a single LDP or STP.

[vhscampos]

Amending patch to include the transfer of memory operands from the original
SDNode to the new MachineSDNode.

[vhscampos]

1. Create new AArch64ISD nodes specific to load/store of pairs of registers.
2. Custom lower i128 volatile loads/stores to these new AArch64ISD nodes.

[vhscampos]

1. Move new ISD node definitions to TableGen.
2. Instruction selection done in TableGen using patterns.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

3025

Not sure we're actually guaranteed here that StoreNode->getValue() is an i128.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

3025

Actually, we're probably fine here if it's a volatile store. Maybe worth adding an assert, though.

[vhscampos]

Add an assertion.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.h

enum NodeType : unsigned {
   LD4LANEpost,
   ST2LANEpost,
   ST3LANEpost,
   ST4LANEpost,
 
   STG,
   STZG,
   ST2G,
-  STZ2G
+  STZ2G,
 
+  LDP,
+  STP
 };
 
 } // end namespace AArch64ISD
 
 namespace {
 
 // Any instruction that defines a 32-bit result zeros out the high half of the
 // register. Truncate can be lowered to EXTRACT_SUBREG. CopyFromReg may

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::FLT_ROUNDS_, MVT::i32, Custom);
 
   setOperationAction(ISD::ATOMIC_CMP_SWAP, MVT::i128, Custom);
   setOperationAction(ISD::ATOMIC_LOAD_SUB, MVT::i32, Custom);
   setOperationAction(ISD::ATOMIC_LOAD_SUB, MVT::i64, Custom);
   setOperationAction(ISD::ATOMIC_LOAD_AND, MVT::i32, Custom);
   setOperationAction(ISD::ATOMIC_LOAD_AND, MVT::i64, Custom);
 
+  // 128-bit loads and stores can be done without expanding
+  setOperationAction(ISD::LOAD, MVT::i128, Custom);
+  setOperationAction(ISD::STORE, MVT::i128, Custom);
+
   // Lower READCYCLECOUNTER using an mrs from PMCCNTR_EL0.
   // This requires the Performance Monitors extension.
   if (Subtarget->hasPerfMon())
     setOperationAction(ISD::READCYCLECOUNTER, MVT::i64, Legal);
 
   if (getLibcallName(RTLIB::SINCOS_STRET_F32) != nullptr &&
       getLibcallName(RTLIB::SINCOS_STRET_F64) != nullptr) {
     // Issue __sincos_stret if available.
const char *AArch64TargetLowering::getTargetNodeName(unsigned Opcode) const {
   case AArch64ISD::GLD1S_IMM:         return "AArch64ISD::GLD1S_IMM";
   case AArch64ISD::SST1:              return "AArch64ISD::SST1";
   case AArch64ISD::SST1_SCALED:       return "AArch64ISD::SST1_SCALED";
   case AArch64ISD::SST1_SXTW:         return "AArch64ISD::SST1_SXTW";
   case AArch64ISD::SST1_UXTW:         return "AArch64ISD::SST1_UXTW";
   case AArch64ISD::SST1_SXTW_SCALED:  return "AArch64ISD::SST1_SXTW_SCALED";
   case AArch64ISD::SST1_UXTW_SCALED:  return "AArch64ISD::SST1_UXTW_SCALED";
   case AArch64ISD::SST1_IMM:          return "AArch64ISD::SST1_IMM";
+  case AArch64ISD::LDP:               return "AArch64ISD::LDP";
+  case AArch64ISD::STP:               return "AArch64ISD::STP";
   }
   return nullptr;
 }
 
 MachineBasicBlock *
 AArch64TargetLowering::EmitF128CSEL(MachineInstr &MI,
                                     MachineBasicBlock *MBB) const {
   // We materialise the F128CSEL pseudo-instruction as some control flow and a
SDValue ExtractTrunc = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i32,
                                      Trunc, DAG.getConstant(0, DL, MVT::i64));
 
   return DAG.getStore(ST->getChain(), DL, ExtractTrunc,
                       ST->getBasePtr(), ST->getMemOperand());
 }
 
 // Custom lowering for any store, vector or scalar and/or default or with
 // a truncate operations.  Currently only custom lower truncate operation
-// from vector v4i16 to v4i8.
+// from vector v4i16 to v4i8 or volatile stores of i128.
 SDValue AArch64TargetLowering::LowerSTORE(SDValue Op,
                                           SelectionDAG &DAG) const {
   SDLoc Dl(Op);
   StoreSDNode *StoreNode = cast<StoreSDNode>(Op);
   assert (StoreNode && "Can only custom lower store nodes");
 
   SDValue Value = StoreNode->getValue();
 
   EVT VT = Value.getValueType();
   EVT MemVT = StoreNode->getMemoryVT();
 
-  assert (VT.isVector() && "Can only custom lower vector store types");
-
+  if (VT.isVector()) {
     unsigned AS = StoreNode->getAddressSpace();
     unsigned Align = StoreNode->getAlignment();
     if (Align < MemVT.getStoreSize() &&
-      !allowsMisalignedMemoryAccesses(
-          MemVT, AS, Align, StoreNode->getMemOperand()->getFlags(), nullptr)) {
+        !allowsMisalignedMemoryAccesses(MemVT, AS, Align,
+                                        StoreNode->getMemOperand()->getFlags(),
+                                        nullptr)) {
       return scalarizeVectorStore(StoreNode, DAG);
     }
 
     if (StoreNode->isTruncatingStore()) {
       return LowerTruncateVectorStore(Dl, StoreNode, VT, MemVT, DAG);
     }
+  } else if (MemVT == MVT::i128 && StoreNode->isVolatile()) {
+    assert(StoreNode->getValue()->getValueType(0) == MVT::i128);
+    SDValue Lo =
+        DAG.getNode(ISD::EXTRACT_ELEMENT, Dl, MVT::i64, StoreNode->getValue(),
+                    DAG.getConstant(0, Dl, MVT::i64));
+    SDValue Hi =
+        DAG.getNode(ISD::EXTRACT_ELEMENT, Dl, MVT::i64, StoreNode->getValue(),
+                    DAG.getConstant(1, Dl, MVT::i64));
+    SDValue Result = DAG.getMemIntrinsicNode(
+        AArch64ISD::STP, Dl, DAG.getVTList(MVT::Other),
+        {StoreNode->getChain(), Lo, Hi, StoreNode->getBasePtr()},
+        StoreNode->getMemoryVT(), StoreNode->getMemOperand());
+    return Result;
+  }
 
   return SDValue();
 }
 
 SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   LLVM_DEBUG(dbgs() << "Custom lowering: ");
   LLVM_DEBUG(Op.dump());
void AArch64TargetLowering::ReplaceNodeResults(
   case ISD::FP_TO_UINT:
   case ISD::FP_TO_SINT:
     assert(N->getValueType(0) == MVT::i128 && "unexpected illegal conversion");
     // Let normal code take care of it by not adding anything to Results.
     return;
   case ISD::ATOMIC_CMP_SWAP:
     ReplaceCMP_SWAP_128Results(N, Results, DAG, Subtarget);
     return;
+  case ISD::LOAD: {
+    assert(SDValue(N, 0).getValueType() == MVT::i128 &&
+           "unexpected load's value type");
+    LoadSDNode *LoadNode = cast<LoadSDNode>(N);
+    if (!LoadNode->isVolatile() || LoadNode->getMemoryVT() != MVT::i128) {
+      // Non-volatile loads are optimized later in AArch64's load/store
+      // optimizer.
+      return;
+    }
+
+    SDValue Result = DAG.getMemIntrinsicNode(
+        AArch64ISD::LDP, SDLoc(N),
+        DAG.getVTList({MVT::i64, MVT::i64, MVT::Other}),
+        {LoadNode->getChain(), LoadNode->getBasePtr()}, LoadNode->getMemoryVT(),
+        LoadNode->getMemOperand());
+
+    SDValue Pair = DAG.getNode(ISD::BUILD_PAIR, SDLoc(N), MVT::i128,
+                               Result.getValue(0), Result.getValue(1));
+    Results.append({Pair, Result.getValue(2) /* Chain */});
+    return;
+  }
   }
 }
 
 bool AArch64TargetLowering::useLoadStackGuardNode() const {
   if (Subtarget->isTargetAndroid() || Subtarget->isTargetFuchsia())
     return TargetLowering::useLoadStackGuardNode();
   return true;
 }

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 def SDT_AArch64TCRET : SDTypeProfile<0, 2, [SDTCisPtrTy<0>]>;
 def SDT_AArch64PREFETCH : SDTypeProfile<0, 2, [SDTCisVT<0, i32>, SDTCisPtrTy<1>]>;
 
 def SDT_AArch64ITOF  : SDTypeProfile<1, 1, [SDTCisFP<0>, SDTCisSameAs<0,1>]>;
 
 def SDT_AArch64TLSDescCall : SDTypeProfile<0, -2, [SDTCisPtrTy<0>,
                                                  SDTCisPtrTy<1>]>;
 
+def SDT_AArch64ldp : SDTypeProfile<2, 1, [SDTCisVT<0, i64>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
+def SDT_AArch64stp : SDTypeProfile<0, 3, [SDTCisVT<0, i64>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
+
 // Generates the general dynamic sequences, i.e.
 //  adrp  x0, :tlsdesc:var
 //  ldr   x1, [x0, #:tlsdesc_lo12:var]
 //  add   x0, x0, #:tlsdesc_lo12:var
 //  .tlsdesccall var
 //  blr   x1
 
 // (the TPIDR_EL0 offset is put directly in X0, hence no "result" here)

 def SDT_AArch64unpk : SDTypeProfile<1, 1, [
     SDTCisInt<0>, SDTCisInt<1>, SDTCisOpSmallerThanOp<1, 0>
 ]>;
 def AArch64sunpkhi : SDNode<"AArch64ISD::SUNPKHI", SDT_AArch64unpk>;
 def AArch64sunpklo : SDNode<"AArch64ISD::SUNPKLO", SDT_AArch64unpk>;
 def AArch64uunpkhi : SDNode<"AArch64ISD::UUNPKHI", SDT_AArch64unpk>;
 def AArch64uunpklo : SDNode<"AArch64ISD::UUNPKLO", SDT_AArch64unpk>;
 
+def AArch64ldp : SDNode<"AArch64ISD::LDP", SDT_AArch64ldp, [SDNPHasChain, SDNPMayLoad, SDNPMemOperand]>;
+def AArch64stp : SDNode<"AArch64ISD::STP", SDT_AArch64stp, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;
+
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 
 // AArch64 Instruction Predicate Definitions.
 // We could compute these on a per-module basis but doing so requires accessing
 // the Function object through the <Target>Subtarget and objections were raised
 // to that (see post-commit review comments for r301750).

 
 // Pair (no allocate)
 defm LDNPW : LoadPairNoAlloc<0b00, 0, GPR32z, simm7s4, "ldnp">;
 defm LDNPX : LoadPairNoAlloc<0b10, 0, GPR64z, simm7s8, "ldnp">;
 defm LDNPS : LoadPairNoAlloc<0b00, 1, FPR32Op, simm7s4, "ldnp">;
 defm LDNPD : LoadPairNoAlloc<0b01, 1, FPR64Op, simm7s8, "ldnp">;
 defm LDNPQ : LoadPairNoAlloc<0b10, 1, FPR128Op, simm7s16, "ldnp">;
 
+def : Pat<(AArch64ldp (am_indexed7s64 GPR64sp:$Rn, simm7s8:$offset)),
+          (LDPXi GPR64sp:$Rn, simm7s8:$offset)>;
+
 //---
 // (register offset)
 //---
 
 // Integer
 defm LDRBB : Load8RO<0b00,  0, 0b01, GPR32, "ldrb", i32, zextloadi8>;
 defm LDRHH : Load16RO<0b01, 0, 0b01, GPR32, "ldrh", i32, zextloadi16>;
 defm LDRW  : Load32RO<0b10, 0, 0b01, GPR32, "ldr", i32, load>;

 
 // Pair (no allocate)
 defm STNPW : StorePairNoAlloc<0b00, 0, GPR32z, simm7s4, "stnp">;
 defm STNPX : StorePairNoAlloc<0b10, 0, GPR64z, simm7s8, "stnp">;
 defm STNPS : StorePairNoAlloc<0b00, 1, FPR32Op, simm7s4, "stnp">;
 defm STNPD : StorePairNoAlloc<0b01, 1, FPR64Op, simm7s8, "stnp">;
 defm STNPQ : StorePairNoAlloc<0b10, 1, FPR128Op, simm7s16, "stnp">;
 
+def : Pat<(AArch64stp GPR64z:$Rt, GPR64z:$Rt2, (am_indexed7s64 GPR64sp:$Rn, simm7s8:$offset)),
+          (STPXi GPR64z:$Rt, GPR64z:$Rt2, GPR64sp:$Rn, simm7s8:$offset)>;
+
 //---
 // (Register offset)
 
 // Integer
 defm STRBB : Store8RO< 0b00, 0, 0b00, GPR32, "strb", i32, truncstorei8>;
 defm STRHH : Store16RO<0b01, 0, 0b00, GPR32, "strh", i32, truncstorei16>;
 defm STRW  : Store32RO<0b10, 0, 0b00, GPR32, "str",  i32, store>;
 defm STRX  : Store64RO<0b11, 0, 0b00, GPR64, "str",  i64, store>;

llvm/test/CodeGen/AArch64/cmpxchg-O0.ll

 
 ; Original implementation assumed the desired & new arguments had already been
 ; type-legalized into some kind of BUILD_PAIR operation and crashed when this
 ; was false.
 @var128 = global i128 0
 define {i128, i1} @test_cmpxchg_128_unsplit(i128* %addr) {
 ; CHECK-LABEL: test_cmpxchg_128_unsplit:
 ; CHECK:     add x[[VAR128:[0-9]+]], {{x[0-9]+}}, :lo12:var128
-; CHECK:     ldr [[DESIRED_HI:x[0-9]+]], [x[[VAR128]], #8]
-; CHECK:     ldr [[DESIRED_LO:x[0-9]+]], [x[[VAR128]]]
-; CHECK:     ldr [[NEW_HI:x[0-9]+]], [x[[VAR128]], #8]
-; CHECK:     ldr [[NEW_LO:x[0-9]+]], [x[[VAR128]]]
+; CHECK:     ldp [[DESIRED_LO:x[0-9]+]], [[DESIRED_HI:x[0-9]+]], [x[[VAR128]]]
+; CHECK:     ldp [[NEW_LO:x[0-9]+]], [[NEW_HI:x[0-9]+]], [x[[VAR128]]]
 ; CHECK: [[RETRY:.LBB[0-9]+_[0-9]+]]:
 ; CHECK:     ldaxp [[OLD_LO:x[0-9]+]], [[OLD_HI:x[0-9]+]], [x0]
 ; CHECK:     cmp [[OLD_LO]], [[DESIRED_LO]]
 ; CHECK:     cset [[CMP_TMP:w[0-9]+]], ne
 ; CHECK:     cmp [[OLD_HI]], [[DESIRED_HI]]
 ; CHECK:     cinc [[CMP:w[0-9]+]], [[CMP_TMP]], ne
 ; CHECK:     cbnz [[CMP]], [[DONE:.LBB[0-9]+_[0-9]+]]
 ; CHECK:     stlxp [[STATUS:w[0-9]+]], [[NEW_LO]], [[NEW_HI]], [x0]
 ; CHECK:     cbnz [[STATUS]], [[RETRY]]
 ; CHECK: [[DONE]]:
 
   %desired = load volatile i128, i128* @var128
   %new = load volatile i128, i128* @var128
   %val = cmpxchg i128* %addr, i128 %desired, i128 %new seq_cst seq_cst
   ret { i128, i1 } %val
 }

llvm/test/CodeGen/AArch64/i128_volatile_load_store.ll

+; RUN: llc -mtriple=aarch64 %s -o - | FileCheck %s
+
+@x = common dso_local global i128 0
+@y = common dso_local global i128 0
+
+define void @test1() {
+; CHECK-LABEL: test1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    ldp x8, x9, [x8]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    stp x8, x9, [x10]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* @x
+  store volatile i128 %tmp, i128* @y
+  ret void
+}
+
+define void @test2() {
+; CHECK-LABEL: test2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    ldp x8, x9, [x8, #504]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    stp x8, x9, [x10, #504]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @x to i8*), i64 504) to i128*)
+  store volatile i128 %tmp, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @y to i8*), i64 504) to i128*)
+  ret void
+}
+
+define void @test3() {
+; CHECK-LABEL: test3:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    add x8, x8, #512 // =512
+; CHECK-NEXT:    ldp x8, x9, [x8]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    add x10, x10, #512 // =512
+; CHECK-NEXT:    stp x8, x9, [x10]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @x to i8*), i64 512) to i128*)
+  store volatile i128 %tmp, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @y to i8*), i64 512) to i128*)
+  ret void
+}
+
+define void @test4() {
+; CHECK-LABEL: test4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    ldp x8, x9, [x8, #-512]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    stp x8, x9, [x10, #-512]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @x to i8*), i64 -512) to i128*)
+  store volatile i128 %tmp, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @y to i8*), i64 -512) to i128*)
+  ret void
+}
+
+define void @test5() {
+; CHECK-LABEL: test5:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    sub x8, x8, #520 // =520
+; CHECK-NEXT:    ldp x8, x9, [x8]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    sub x10, x10, #520 // =520
+; CHECK-NEXT:    stp x8, x9, [x10]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @x to i8*), i64 -520) to i128*)
+  store volatile i128 %tmp, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @y to i8*), i64 -520) to i128*)
+  ret void
+}
+
+define void @test6() {
+; CHECK-LABEL: test6:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    sub x8, x8, #520 // =520
+; CHECK-NEXT:    ldp x8, x9, [x8]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    sub x10, x10, #520 // =520
+; CHECK-NEXT:    stp x8, x9, [x10]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @x to i8*), i64 -520) to i128*)
+  store volatile i128 %tmp, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @y to i8*), i64 -520) to i128*)
+  ret void
+}
+
+define void @test7() {
+; CHECK-LABEL: test7:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    adrp x8, x
+; CHECK-NEXT:    add x8, x8, :lo12:x
+; CHECK-NEXT:    add x8, x8, #503 // =503
+; CHECK-NEXT:    ldp x8, x9, [x8]
+; CHECK-NEXT:    adrp x10, y
+; CHECK-NEXT:    add x10, x10, :lo12:y
+; CHECK-NEXT:    add x10, x10, #503 // =503
+; CHECK-NEXT:    stp x8, x9, [x10]
+; CHECK-NEXT:    ret
+  %tmp = load volatile i128, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @x to i8*), i64 503) to i128*)
+  store volatile i128 %tmp, i128* bitcast (i8* getelementptr (i8, i8* bitcast (i128* @y to i8*), i64 503) to i128*)
+  ret void
+}