[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

1547

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

1540

BytesIncrement doesn't need to be declared out here?

1545

getFixedSize?

1547

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

1547

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/include/llvm/CodeGen/SelectionDAG.h

#endif
   SDValue getUNDEF(EVT VT) {
     return getNode(ISD::UNDEF, SDLoc(), VT);
   }
 
   /// Return a node that represents the runtime scaling 'MulImm * RuntimeVL'.
   SDValue getVScale(const SDLoc &DL, EVT VT, APInt MulImm) {
     assert(MulImm.getMinSignedBits() <= VT.getSizeInBits() &&
            "Immediate does not fit VT");
-    return getNode(ISD::VSCALE, DL, VT, getConstant(MulImm, DL, VT));
+    return getNode(ISD::VSCALE, DL, VT,
+                   getConstant(MulImm.sextOrTrunc(VT.getSizeInBits()), DL, VT));
   }
 
   /// Return a GLOBAL_OFFSET_TABLE node. This does not have a useful SDLoc.
   SDValue getGLOBAL_OFFSET_TABLE(EVT VT) {
     return getNode(ISD::GLOBAL_OFFSET_TABLE, SDLoc(), VT);
   }
 
   /// Gets or creates the specified node.

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()) {
     ReplaceValueWith(SDValue(LD, 1), NewChain);
     return;
   }
 
   Lo = DAG.getLoad(ISD::UNINDEXED, ExtType, LoVT, dl, Ch, Ptr, Offset,
                    LD->getPointerInfo(), LoMemVT, LD->getOriginalAlign(),
                    MMOFlags, AAInfo);
 
-  unsigned IncrementSize = LoMemVT.getSizeInBits()/8;
+  unsigned IncrementSize = LoMemVT.getSizeInBits().getKnownMinSize() / 8;
+
+  MachinePointerInfo MPI;
+  if (LoVT.isScalableVector()) {
+    SDValue BytesIncrement = DAG.getVScale(
+        dl, Ptr.getValueType(),
+        APInt(Ptr.getValueSizeInBits().getFixedSize(), IncrementSize));
+    MPI = MachinePointerInfo(LD->getPointerInfo().getAddrSpace());
+    Ptr = DAG.getNode(ISD::ADD, dl, Ptr.getValueType(), Ptr, BytesIncrement);
+  } else {
+    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);
+  }
+
+  Hi = DAG.getLoad(ISD::UNINDEXED, ExtType, HiVT, dl, Ch, Ptr, Offset, 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/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::FNEG:
 
     if (OpOpcode == ISD::FNEG)  // --X -> X
       return Operand.getOperand(0);
     break;
   case ISD::FABS:
     if (OpOpcode == ISD::FNEG)  // abs(-X) -> abs(X)
       return getNode(ISD::FABS, DL, VT, Operand.getOperand(0));
     break;
+  case ISD::VSCALE:
+    assert(VT == Operand.getValueType() && "Unexpected VT!");
+    break;
   }
 
   SDNode *N;
   SDVTList VTs = getVTList(VT);
   SDValue Ops[] = {Operand};
   if (VT != MVT::Glue) { // Don't CSE flag producing nodes
     FoldingSetNodeID ID;
     AddNodeIDNode(ID, Opcode, VTs, Ops);

llvm/test/CodeGen/AArch64/sve-split-load.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
+}