[SVE][CodeGen] Legalisation of unpredicated load instructions

[kmclaughlin]

Summary

When splitting a load of a scalable type, the new address is
calculated in SplitVecRes_LOAD using a vscale and an add instruction.

This patch also adds a DAG combiner fold to visitADD for vscale:

• Fold (add (vscale(C0)), (vscale(C1))) to (add (vscale(C0 + C1)))

[kmclaughlin]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

1549

Where LoVT is not scalable we update the PointerInfo with the offset, IncrementSize. For scalable types I've just reused LD->getPointerInfo() for now; I'm not entirely sure how best to handle the PointerInfo for scalable vectors as MachinePointerInfo currently has no knowledge of scaled offsets.

[efriedma]

Please split the addressing mode improvements off into a separate patch from the legalization change.

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

1542

BytesIncrement doesn't need to be declared out here?

1547

getFixedSize?

1549

You can get an "empty" MachinePointerInfo. Something like MachinePointerInfo(LD->getPointerInfo()->getAddrSpace());

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2379

Is there some reason to expect the two integers have the same width?

[kmclaughlin]

• Split the addressing mode changes into a separate patch (D82893)
• Changed MPI to MachinePointerInfo(LD->getPointerInfo()->getAddrSpace()) in SplitVecRes_LOAD
• Use getFixedSize when creating the APInt in SplitVecRes_LOAD

[kmclaughlin]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2379

I thought we could expect that the integers will have the same width here, since we're explicitly checking the opcodes used in the expression. I would have thought if the integers had different widths that the expression would also include a sign/zero extend or a truncate somewhere?

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

1549

Thanks for the suggestion on this!

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2379

Both operands of an add have the same type as the result. Not sure that says anything about the operands of the vscale operations: we don't enforce any rule there. If it's supposed to be the same as the result type, probably SelectionDAG::getNode() should check that.

[kmclaughlin]

• Use sextOrTrunc in getVScale() to ensure that MulImm matches the return type of the vscale
• Add an assertion to getNode() which checks that the operand and return types of vscale match

[david-arm]

LGTM!

[kmclaughlin]

Thanks for reviewing this change, @efriedma & @david-arm!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitADD(SDNode *N) {
 
   // Fold (add (vscale * C0), (vscale * C1)) to (vscale * (C0 + C1)).
   if (N0.getOpcode() == ISD::VSCALE && N1.getOpcode() == ISD::VSCALE) {
     APInt C0 = N0->getConstantOperandAPInt(0);
     APInt C1 = N1->getConstantOperandAPInt(0);
     return DAG.getVScale(DL, VT, C0 + C1);
   }
 
+  // fold a+vscale(c1)+vscale(c2) -> a+vscale(c1+c2)
+  if ((N0.getOpcode() == ISD::ADD) &&
+      (N0.getOperand(1).getOpcode() == ISD::VSCALE) &&
+      (N1.getOpcode() == ISD::VSCALE)) {
+    auto VS0 = N0.getOperand(1)->getConstantOperandAPInt(0);
+    auto VS1 = N1->getConstantOperandAPInt(0);
+    auto VS = DAG.getVScale(DL, VT, VS0 + VS1);
+    return DAG.getNode(ISD::ADD, DL, VT, N0.getOperand(0), VS);
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitADDSAT(SDNode *N) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   EVT VT = N0.getValueType();

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

if (!LoMemVT.isByteSized() || !HiMemVT.isByteSized()) {
     return;
   }
 
   Lo = DAG.getLoad(ISD::UNINDEXED, ExtType, LoVT, dl, Ch, Ptr, Offset,
                    LD->getPointerInfo(), LoMemVT, LD->getOriginalAlign(),
                    MMOFlags, AAInfo);
 
   unsigned IncrementSize = LoMemVT.getSizeInBits()/8;
+
+  SDValue BytesIncrement;
+  MachinePointerInfo MPI;
+  if (LoVT.isScalableVector()) {
+    BytesIncrement = DAG.getVScale(
+        dl, Ptr.getValueType(),
+        APInt(Ptr.getValueType().getSizeInBits().getKnownMinSize(),
+              IncrementSize));
+    MPI = LD->getPointerInfo();
+    Ptr = DAG.getNode(ISD::ADD, dl, Ptr.getValueType(), Ptr, BytesIncrement);
+  } else {
+    BytesIncrement = DAG.getConstant(IncrementSize, dl, Ptr.getValueType());
+    MPI = LD->getPointerInfo().getWithOffset(IncrementSize);
     Ptr = DAG.getObjectPtrOffset(dl, Ptr, IncrementSize);
+  }
+
   Hi = DAG.getLoad(ISD::UNINDEXED, ExtType, HiVT, dl, Ch, Ptr, Offset,
-                   LD->getPointerInfo().getWithOffset(IncrementSize), HiMemVT,
-                   LD->getOriginalAlign(), MMOFlags, AAInfo);
+                   MPI, HiMemVT, LD->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));
 
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

multiclass sve_prefetch<SDPatternOperator prefetch, ValueType PredTy, Instruction RegImmInst, Instruction RegRegInst, int scale, ComplexPattern AddrCP> {
   defm : unpred_store<         store, nxv4f16, ST1H_S_IMM, PTRUE_S>;
   defm : unpred_store<         store, nxv2f16, ST1H_D_IMM, PTRUE_D>;
   defm : unpred_store<         store, nxv4f32,   ST1W_IMM, PTRUE_S>;
   defm : unpred_store<         store, nxv4f32, ST1W_D_IMM, PTRUE_D>;
   defm : unpred_store<         store, nxv2f64,   ST1D_IMM, PTRUE_D>;
 
   multiclass unpred_load<PatFrag Load, ValueType Ty, Instruction RegImmInst,
                          Instruction PTrue> {
-    def : Pat<(Ty (Load  (am_sve_fi GPR64sp:$base, simm4s1:$offset))),
+    let AddedComplexity = 1 in {
+      def _imm: Pat<(Ty (Load  (am_sve_indexed_s4 GPR64sp:$base, simm4s1:$offset))),
+                    (RegImmInst (PTrue 31), GPR64sp:$base, simm4s1:$offset)>;
+    }
+
+    let AddedComplexity = 2 in {
+      def _fi : Pat<(Ty (Load  (am_sve_fi GPR64sp:$base, simm4s1:$offset))),
                     (RegImmInst (PTrue 31), GPR64sp:$base, simm4s1:$offset)>;
+    }
 
     def : Pat<(Ty (Load GPR64:$base)),
               (RegImmInst (PTrue 31), GPR64:$base, (i64 0))>;
   }
 
   defm : unpred_load<        load, nxv16i8,    LD1B_IMM, PTRUE_B>;
   defm : unpred_load< zextloadvi8, nxv8i16,  LD1B_H_IMM, PTRUE_H>;
   defm : unpred_load< zextloadvi8, nxv4i32,  LD1B_S_IMM, PTRUE_S>;

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+; LOAD
+
+define <vscale x 4 x i16> @load_promote_4i8(<vscale x 4 x i16>* %a) {
+; CHECK-LABEL: load_promote_4i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ld1h { z0.s }, p0/z, [x0]
+; CHECK-NEXT:    ret
+  %load = load <vscale x 4 x i16>, <vscale x 4 x i16>* %a
+  ret <vscale x 4 x i16> %load
+}
+
+define <vscale x 16 x i16> @load_split_i16(<vscale x 16 x i16>* %a) {
+; CHECK-LABEL: load_split_i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    ld1h { z0.h }, p0/z, [x0]
+; CHECK-NEXT:    ld1h { z1.h }, p0/z, [x0, #1, mul vl]
+; CHECK-NEXT:    ret
+  %load = load <vscale x 16 x i16>, <vscale x 16 x i16>* %a
+  ret <vscale x 16 x i16> %load
+}
+
+define <vscale x 32 x i16> @load_split_32i16(<vscale x 32 x i16>* %a) {
+; CHECK-LABEL: load_split_32i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    ld1h { z0.h }, p0/z, [x0]
+; CHECK-NEXT:    ld1h { z1.h }, p0/z, [x0, #1, mul vl]
+; CHECK-NEXT:    ld1h { z2.h }, p0/z, [x0, #2, mul vl]
+; CHECK-NEXT:    ld1h { z3.h }, p0/z, [x0, #3, mul vl]
+; CHECK-NEXT:    ret
+  %load = load <vscale x 32 x i16>, <vscale x 32 x i16>* %a
+  ret <vscale x 32 x i16> %load
+}
+
+define <vscale x 16 x i64> @load_split_16i64(<vscale x 16 x i64>* %a) {
+; CHECK-LABEL: load_split_16i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    ld1d { z0.d }, p0/z, [x0]
+; CHECK-NEXT:    ld1d { z1.d }, p0/z, [x0, #1, mul vl]
+; CHECK-NEXT:    ld1d { z2.d }, p0/z, [x0, #2, mul vl]
+; CHECK-NEXT:    ld1d { z3.d }, p0/z, [x0, #3, mul vl]
+; CHECK-NEXT:    ld1d { z4.d }, p0/z, [x0, #4, mul vl]
+; CHECK-NEXT:    ld1d { z5.d }, p0/z, [x0, #5, mul vl]
+; CHECK-NEXT:    ld1d { z6.d }, p0/z, [x0, #6, mul vl]
+; CHECK-NEXT:    ld1d { z7.d }, p0/z, [x0, #7, mul vl]
+; CHECK-NEXT:    ret
+  %load = load <vscale x 16 x i64>, <vscale x 16 x i64>* %a
+  ret <vscale x 16 x i64> %load
+}