[WIP][LegalizeTypes][LegalizeDAG] Use misaligned load/store to optimize memory access with non-power2 integer types.

[bcl5980]

Summary

When load integer types with non-power2 aligned size, we will split the load to some no overlap legal type loads.
For example, if we trying to legalize i56, it will be split to i32 + i16 + i8 three loads.
This change trying to use two i32 load with 8bits overlap to reduce the load number.

The motivation comes from ARM64EC (https://reviews.llvm.org/D125418#inline-1267564).
For now, we don't apply it to store because it will involve extra dependency in CPU load store queue.

[bcl5980]

fix test fail.

[efriedma]

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

800

The alignment you're passing in here doesn't seem right; you want to pass in the alignment of the load you're planning to generate, right?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2045

None of this logic looks like it's target-specific; can we just do this in target-independent code?

[bcl5980]

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

800

Ah, yeah. I missed it. Thanks for the catch.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2045

Yes. Of course, we can move this code into target-independent layer.
But there are some regressions on other platform like x86, comes from load/store pair.
For now, I haven't enabled the optimization for store. So, load pattern is different from store if we have a load/store pair.
For example,

define void @i56_or(ptr %a) {
  %aa = load i56, ptr %a, align 1
  %b = or i56 %aa, 384
  store i56 %b, ptr %a, align 1
  ret void
}

Actually, this is a general issue but the gain from load can cover the extra instructions on ARM but x86 can't. So, I limited this to ARM.
Maybe we can enable the optimization for store also to fix the issue. But I still have a little concern on it.

[bcl5980]

Move the code to independent layer.
Send the patch to review but for now x86 still has some regressions.

[efriedma]

Do we have test coverage for load+store patterns involving integers larger than 56 bits?

If the code quality issues are specific to 56 bits, it might make sense to make the initial patch skip handling 56-bit integers, and try to address them in a followup.

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

812

The alignment on this load isn't right.

The "alignment" of a load SDNode isn't the alignment of the load itself; it's the alignment of the base of the MachinePointerInfo. So if you're using getWithOffset() to adjust the MachinePointerInfo, the alignment doesn't change based on the offset. (This distinction helps make alias analysis more accurate.)

[bcl5980]

Update load+store patterns with integers larger than 56 bits. Looks even worse than 56bits.

llvm/include/llvm/CodeGen/TargetLowering.h

#include "llvm/IR/ConstrainedOps.def"
   /// LLT handling variant.
   virtual bool allowsMisalignedMemoryAccesses(
       LLT, unsigned AddrSpace = 0, Align Alignment = Align(1),
       MachineMemOperand::Flags Flags = MachineMemOperand::MONone,
       bool * /*Fast*/ = nullptr) const {
     return false;
   }
 
+  /// Use unaligned memory access for non-power2 types
+  virtual bool allowMisalignedMemForNonPow2Type(
+      EVT, unsigned AddrSpace = 0, Align Alignment = Align(1),
+      MachineMemOperand::Flags Flags = MachineMemOperand::MONone) const {
+    return false;
+  }
+
   /// This function returns true if the memory access is aligned or if the
   /// target allows this specific unaligned memory access. If the access is
   /// allowed, the optional final parameter returns if the access is also fast
   /// (as defined by the target).
   bool allowsMemoryAccessForAlignment(
       LLVMContext &Context, const DataLayout &DL, EVT VT,
       unsigned AddrSpace = 0, Align Alignment = Align(1),
       MachineMemOperand::Flags Flags = MachineMemOperand::MONone,

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

if (SrcWidth != SrcVT.getStoreSizeInBits() &&
 
     if (DL.isLittleEndian()) {
       // EXTLOAD:i24 -> ZEXTLOAD:i16 | (shl EXTLOAD@+2:i8, 16)
       // Load the bottom RoundWidth bits.
       Lo = DAG.getExtLoad(ISD::ZEXTLOAD, dl, Node->getValueType(0), Chain, Ptr,
                           LD->getPointerInfo(), RoundVT, LD->getOriginalAlign(),
                           MMOFlags, AAInfo);
 
+      Align ExtraLoadAlign = LD->getOriginalAlign();
+      unsigned IncSizeBits = RoundWidth;
+      if (TLI.allowMisalignedMemForNonPow2Type(SrcVT, LD->getAddressSpace(),
+                                               ExtraLoadAlign, MMOFlags)) {
+        IncSizeBits = ExtraWidth;
+        ExtraVT = RoundVT;
+        ExtraWidth = RoundWidth;
+        ExtraLoadAlign = Align(1);
+      }
+
       // Load the remaining ExtraWidth bits.
-      IncrementSize = RoundWidth / 8;
+      IncrementSize = IncSizeBits / 8;
       Ptr = DAG.getMemBasePlusOffset(Ptr, TypeSize::Fixed(IncrementSize), dl);
       Hi = DAG.getExtLoad(ExtType, dl, Node->getValueType(0), Chain, Ptr,
                           LD->getPointerInfo().getWithOffset(IncrementSize),
-                          ExtraVT, LD->getOriginalAlign(), MMOFlags, AAInfo);
+                          ExtraVT, ExtraLoadAlign, MMOFlags, AAInfo);
 
       // Build a factor node to remember that this load is independent of
       // the other one.
       Ch = DAG.getNode(ISD::TokenFactor, dl, MVT::Other, Lo.getValue(1),
                        Hi.getValue(1));
 
       // Move the top bits to the right place.
       Hi = DAG.getNode(
           ISD::SHL, dl, Hi.getValueType(), Hi,
-          DAG.getConstant(RoundWidth, dl,
+          DAG.getConstant(IncSizeBits, dl,
                           TLI.getShiftAmountTy(Hi.getValueType(), DL)));
 
       // Join the hi and lo parts.
       Value = DAG.getNode(ISD::OR, dl, Node->getValueType(0), Lo, Hi);
     } else {
       // Big endian - avoid unaligned loads.
       // EXTLOAD:i24 -> (shl EXTLOAD:i16, 8) | ZEXTLOAD@+2:i8
       // Load the top RoundWidth bits.

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

if (N->getMemoryVT().bitsLE(NVT)) {
     // Little-endian - low bits are at low addresses.
     Lo = DAG.getLoad(NVT, dl, Ch, Ptr, N->getPointerInfo(),
                      N->getOriginalAlign(), MMOFlags, AAInfo);
 
     unsigned ExcessBits =
       N->getMemoryVT().getSizeInBits() - NVT.getSizeInBits();
     EVT NEVT = EVT::getIntegerVT(*DAG.getContext(), ExcessBits);
 
+    if (TLI.allowMisalignedMemForNonPow2Type(
+            N->getMemoryVT(), N->getAddressSpace(), N->getAlign(), MMOFlags)) {
+      // use unaligned load to simplify the non-power2 types
+      unsigned IncrementSize = NEVT.getSizeInBits() / 8;
+      Ptr = DAG.getMemBasePlusOffset(Ptr, TypeSize::Fixed(IncrementSize), dl);
+      Hi = DAG.getExtLoad(ExtType, dl, NVT, Ch, Ptr,
+                          N->getPointerInfo().getWithOffset(IncrementSize), NVT,
+                          Align(1), MMOFlags, AAInfo);
+      unsigned Opcode = ExtType == ISD::SEXTLOAD ? ISD::SRA : ISD::SRL;
+      unsigned ShiftCnt = NVT.getSizeInBits() - ExcessBits;
+      Hi = DAG.getNode(Opcode, dl, NVT, Hi, DAG.getConstant(ShiftCnt, dl, NVT));
+    } else {
       // Increment the pointer to the other half.
-    unsigned IncrementSize = NVT.getSizeInBits()/8;
+      unsigned IncrementSize = NVT.getSizeInBits() / 8;
       Ptr = DAG.getMemBasePlusOffset(Ptr, TypeSize::Fixed(IncrementSize), dl);
       Hi = DAG.getExtLoad(ExtType, dl, NVT, Ch, Ptr,
-                        N->getPointerInfo().getWithOffset(IncrementSize), NEVT,
-                        N->getOriginalAlign(), MMOFlags, AAInfo);
-
+                          N->getPointerInfo().getWithOffset(IncrementSize),
+                          NEVT, N->getOriginalAlign(), MMOFlags, AAInfo);
       // Build a factor node to remember that this load is independent of the
       // other one.
       Ch = DAG.getNode(ISD::TokenFactor, dl, MVT::Other, Lo.getValue(1),
                        Hi.getValue(1));
+    }
   } else {
     // Big-endian - high bits are at low addresses.  Favor aligned loads at
     // the cost of some bit-fiddling.
     EVT MemVT = N->getMemoryVT();
     unsigned EBytes = MemVT.getStoreSize();
     unsigned IncrementSize = NVT.getSizeInBits()/8;
     unsigned ExcessBits = (EBytes - IncrementSize)*8;
 

llvm/lib/Target/AArch64/AArch64ISelLowering.h

bool allowsMisalignedMemoryAccesses(
       MachineMemOperand::Flags Flags = MachineMemOperand::MONone,
       bool *Fast = nullptr) const override;
   /// LLT variant.
   bool allowsMisalignedMemoryAccesses(LLT Ty, unsigned AddrSpace,
                                       Align Alignment,
                                       MachineMemOperand::Flags Flags,
                                       bool *Fast = nullptr) const override;
 
+  /// Use unaligned memory access for non-power2 types
+  bool allowMisalignedMemForNonPow2Type(
+      EVT, unsigned AddrSpace, Align Alignment,
+      MachineMemOperand::Flags Flags) const override;
+
   /// Provide custom lowering hooks for some operations.
   SDValue LowerOperation(SDValue Op, SelectionDAG &DAG) const override;
 
   const char *getTargetNodeName(unsigned Opcode) const override;
 
   SDValue PerformDAGCombine(SDNode *N, DAGCombinerInfo &DCI) const override;
 
   /// This method returns a target specific FastISel object, or null if the

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

*Fast = !Subtarget->isMisaligned128StoreSlow() ||
 
             // Disregard v2i64. Memcpy lowering produces those and splitting
             // them regresses performance on micro-benchmarks and olden/bh.
             Ty == LLT::fixed_vector(2, 64);
   }
   return true;
 }
 
+bool AArch64TargetLowering::allowMisalignedMemForNonPow2Type(
+    EVT VT, unsigned AddrSpace, Align Alignment,
+    MachineMemOperand::Flags Flags) const {
+
+  unsigned WidthBits = VT.getSizeInBits().getFixedSize();
+
+  // If pop count is less or equal than 2 we can't get
+  // any benifit from unaligned load/store
+  if (countPopulation(WidthBits) <= 2)
+    return false;
+
+  // If load bits is not byte align or is already 64bit aligned
+  // we don't use unaligned load/store
+  if ((WidthBits & 63) == 0 || (WidthBits & 7) != 0)
+    return false;
+
+  bool fast;
+  if (!allowsMisalignedMemoryAccesses(VT, AddrSpace, Alignment, Flags, &fast))
+    return false;
+
+  return fast;
+}
+
 FastISel *
 AArch64TargetLowering::createFastISel(FunctionLoweringInfo &funcInfo,
                                       const TargetLibraryInfo *libInfo) const {
   return AArch64::createFastISel(funcInfo, libInfo);
 }
 
 const char *AArch64TargetLowering::getTargetNodeName(unsigned Opcode) const {
 #define MAKE_CASE(V)                                                           \

llvm/test/CodeGen/AArch64/arm64-non-pow2-ldst.ll

 ; CHECK-NEXT:    ret
     %r = load i24, i24* %p
     ret i24 %r
 }
 
 define i56 @ldi56(ptr %p) nounwind {
 ; CHECK-LABEL: ldi56:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldrb w8, [x0, #6]
-; CHECK-NEXT:    ldrh w9, [x0, #4]
-; CHECK-NEXT:    ldr w0, [x0]
-; CHECK-NEXT:    bfi w9, w8, #16, #16
-; CHECK-NEXT:    bfi x0, x9, #32, #32
+; CHECK-NEXT:    ldur w8, [x0, #3]
+; CHECK-NEXT:    ldr w9, [x0]
+; CHECK-NEXT:    orr x0, x9, x8, lsl #24
 ; CHECK-NEXT:    ret
     %r = load i56, i56* %p
     ret i56 %r
 }
 
 define i80 @ldi80(ptr %p) nounwind {
 ; CHECK-LABEL: ldi80:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr x8, [x0]
 ; CHECK-NEXT:    ldrh w1, [x0, #8]
 ; CHECK-NEXT:    mov x0, x8
 ; CHECK-NEXT:    ret
     %r = load i80, i80* %p
     ret i80 %r
 }
 
 define i120 @ldi120(ptr %p) nounwind {
 ; CHECK-LABEL: ldi120:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldrb w8, [x0, #14]
-; CHECK-NEXT:    ldrh w9, [x0, #12]
-; CHECK-NEXT:    ldr w1, [x0, #8]
+; CHECK-NEXT:    ldur x8, [x0, #7]
 ; CHECK-NEXT:    ldr x0, [x0]
-; CHECK-NEXT:    bfi w9, w8, #16, #16
-; CHECK-NEXT:    bfi x1, x9, #32, #32
+; CHECK-NEXT:    lsr x1, x8, #8
 ; CHECK-NEXT:    ret
     %r = load i120, i120* %p
     ret i120 %r
 }
 
 define i280 @ldi280(ptr %p) nounwind {
 ; CHECK-LABEL: ldi280:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldp x8, x1, [x0]
-; CHECK-NEXT:    ldrb w9, [x0, #34]
-; CHECK-NEXT:    ldrh w4, [x0, #32]
+; CHECK-NEXT:    ldur x9, [x0, #27]
 ; CHECK-NEXT:    ldp x2, x3, [x0, #16]
 ; CHECK-NEXT:    mov x0, x8
-; CHECK-NEXT:    bfi x4, x9, #16, #8
+; CHECK-NEXT:    lsr x4, x9, #40
 ; CHECK-NEXT:    ret
     %r = load i280, i280* %p
     ret i280 %r
 }
 
 define void @sti24(ptr %p, i24 %a) nounwind {
 ; CHECK-LABEL: sti24:
 ; CHECK:       // %bb.0: