[WIP][AArch64][CodeGen] Support select address mode load/store

[paulwalker-arm]

I don't believe we want this functionality. We originally added it but then it got removed by D88994 as unsound because there's no instruction available to load/store predicates that are smaller than <vscale x 16 x i1>. At the C/C++ level we don't expose these "smaller" predicate types and so there shouldn't really be a route needed to load/store them. @Allen Do you have a real world use case where loading/storing them is required?

[Allen]

In D127596#3577439, @paulwalker-arm wrote:

I don't believe we want this functionality. We originally added it but then it got removed by D88994 as unsound because there's no instruction available to load/store predicates that are smaller than <vscale x 16 x i1>. At the C/C++ level we don't expose these "smaller" predicate types and so there shouldn't really be a route needed to load/store them. @Allen Do you have a real world use case where loading/storing them is required?

Thanks @paulwalker-arm very much, I don't have a C case. But only my colleague found that the case compilation failed and tried to rectify the fault.
You have a good memory to find this patch D88994, which helps me understand the problem, thanks again!

[paulwalker-arm]

Just trying to cleanup my review list since we're agreed either more or different work is required.

[Allen]

Yes, I think this is not worthless

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

if (Found) {
       }
     }
 
     // Otherwise, base only, materialize address in register.
     Base = N;
     OffImm = CurDAG->getTargetConstant(0, SDLoc(N), MVT::i64);
     return true;
   }
+  bool SelectAddrModePred(SDValue N, SDValue &Base, SDValue &OffImm);
 
   template<int Width>
   bool SelectAddrModeWRO(SDValue N, SDValue &Base, SDValue &Offset,
                          SDValue &SignExtend, SDValue &DoShift) {
     return SelectAddrModeWRO(N, Width / 8, Base, Offset, SignExtend, DoShift);
   }
 
   template<int Width>
bool AArch64DAGToDAGISel::SelectAddrModeIndexed(SDValue N, unsigned Size,
   // the memory is accessed.
   //    add x0, Xbase, #offset
   //    ldr x0, [x0]
   Base = N;
   OffImm = CurDAG->getTargetConstant(0, dl, MVT::i64);
   return true;
 }
 
+bool AArch64DAGToDAGISel::SelectAddrModePred(SDValue N, SDValue &Base,
+                                             SDValue &OffImm) {
+  SDLoc dl(N);
+
+  // If this is not a frame index, load directly from this address
+  if (N->getOpcode() != ISD::FrameIndex) {
+    Base = N;
+    OffImm = CurDAG->getTargetConstant(0, dl, MVT::i64);
+    return true;
+  }
+
+  // Otherwise, match it for the frame address
+  const DataLayout &DL = CurDAG->getDataLayout();
+  const TargetLowering *TLI = getTargetLowering();
+  int FI = cast<FrameIndexSDNode>(N)->getIndex();
+  Base = CurDAG->getTargetFrameIndex(FI, TLI->getPointerTy(DL));
+  OffImm = CurDAG->getTargetConstant(0, dl, MVT::i64);
+  return true;
+}
+
 /// SelectAddrModeUnscaled - Select a "register plus unscaled signed 9-bit
 /// immediate" address.  This should only match when there is an offset that
 /// is not valid for a scaled immediate addressing mode.  The "Size" argument
 /// is the size in bytes of the memory reference, which is needed here to know
 /// what is valid for a scaled immediate.
 bool AArch64DAGToDAGISel::SelectAddrModeUnscaled(SDValue N, unsigned Size,
                                                  SDValue &Base,
                                                  SDValue &OffImm) {

llvm/lib/Target/AArch64/AArch64InstrFormats.td

 //---
 
 def am_unscaled8 :  ComplexPattern<iPTR, 2, "SelectAddrModeUnscaled8", []>;
 def am_unscaled16 : ComplexPattern<iPTR, 2, "SelectAddrModeUnscaled16", []>;
 def am_unscaled32 : ComplexPattern<iPTR, 2, "SelectAddrModeUnscaled32", []>;
 def am_unscaled64 : ComplexPattern<iPTR, 2, "SelectAddrModeUnscaled64", []>;
 def am_unscaled128 :ComplexPattern<iPTR, 2, "SelectAddrModeUnscaled128", []>;
 
+def am_sve_pred : ComplexPattern<iPTR, 2, "SelectAddrModePred", []>;
+
 def gi_am_unscaled8 :
     GIComplexOperandMatcher<s64, "selectAddrModeUnscaled8">,
     GIComplexPatternEquiv<am_unscaled8>;
 def gi_am_unscaled16 :
     GIComplexOperandMatcher<s64, "selectAddrModeUnscaled16">,
     GIComplexPatternEquiv<am_unscaled16>;
 def gi_am_unscaled32 :
     GIComplexOperandMatcher<s64, "selectAddrModeUnscaled32">,

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def : Pat<(nxv16i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_B 31), PPR:$Ps1, PPR:$Ps2)>;
   def : Pat<(nxv8i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_H 31), PPR:$Ps1, PPR:$Ps2)>;
   def : Pat<(nxv4i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_S 31), PPR:$Ps1, PPR:$Ps2)>;
   def : Pat<(nxv2i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_D 31), PPR:$Ps1, PPR:$Ps2)>;
 
+  // Add select addressing mode load/store
+  foreach type = ["nxv16i1", "nxv8i1", "nxv4i1", "nxv2i1"] in {
+    def : Pat< (!cast<ValueType>(type)
+                         (load (am_sve_pred GPR64sp:$base, simm9:$offset))),
+                (LDR_PXI GPR64sp:$base, simm9:$offset)>;
+    def : Pat<(store (!cast<ValueType>(type) PPR:$val),
+                         (am_sve_pred GPR64sp:$base, simm9:$offset)),
+               (STR_PXI PPR:$val, GPR64sp:$base, simm9:$offset)>;
+  }
+
   // Add more complex addressing modes here as required
   multiclass pred_load<ValueType Ty, ValueType PredTy, SDPatternOperator Load,
                        Instruction RegRegInst, Instruction RegImmInst, ComplexPattern AddrCP> {
     let AddedComplexity = 1 in {
       def _reg_reg_z : Pat<(Ty (Load (AddrCP GPR64:$base, GPR64:$offset), (PredTy PPR:$gp), (SVEDup0Undef))),
                            (RegRegInst PPR:$gp, GPR64:$base, GPR64:$offset)>;
     }
     let AddedComplexity = 2 in {

llvm/test/CodeGen/AArch64/aarch64-sve-ldst.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -asm-verbose=1 < %s | FileCheck %s
+
+; Support a new select address mode load/store
+
+;;; LOAD
+
+define <vscale x 16 x i1> @sve_load16(<vscale x 16 x i1>* %base) {
+; CHECK-LABEL: sve_load16:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ldr p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %src = load <vscale x 16 x i1>, <vscale x 16 x i1>* %base, align 1
+  ret <vscale x 16 x i1> %src
+}
+
+define <vscale x 8 x i1> @sve_load8(<vscale x 8 x i1>* %base) {
+; CHECK-LABEL: sve_load8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ldr p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %src = load <vscale x 8 x i1>, <vscale x 8 x i1>* %base, align 1
+  ret <vscale x 8 x i1> %src
+}
+
+define <vscale x 4 x i1> @sve_load4(<vscale x 4 x i1>* %base) {
+; CHECK-LABEL: sve_load4:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ldr p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %src = load <vscale x 4 x i1>, <vscale x 4 x i1>* %base, align 1
+  ret <vscale x 4 x i1> %src
+}
+
+define <vscale x 2 x i1> @sve_load2(<vscale x 2 x i1>* %base) {
+; CHECK-LABEL: sve_load2:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ldr p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %src = load <vscale x 2 x i1>, <vscale x 2 x i1>* %base, align 1
+  ret <vscale x 2 x i1> %src
+}
+
+;;; STORE
+
+define void @sve_store16(<vscale x 16 x i1>* %base, <vscale x 16 x i1> %src) {
+; CHECK-LABEL: sve_store16:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 16 x i1> %src, <vscale x 16 x i1>* %base, align 1
+  ret void
+}
+
+define void @sve_store8(<vscale x 8 x i1>* %base, <vscale x 8 x i1> %src) {
+; CHECK-LABEL: sve_store8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 8 x i1> %src, <vscale x 8 x i1>* %base, align 1
+  ret void
+}
+
+define void @sve_store4(<vscale x 4 x i1>* %base, <vscale x 4 x i1> %src) {
+; CHECK-LABEL: sve_store4:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 4 x i1> %src, <vscale x 4 x i1>* %base, align 1
+  ret void
+}
+
+define void @sve_store2(<vscale x 2 x i1>* %base, <vscale x 2 x i1> %src) {
+; CHECK-LABEL: sve_store2:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 2 x i1> %src, <vscale x 2 x i1>* %base, align 1
+  ret void
+}
+