[SelectionDAG] Widen scalable-vector loads/stores via VP_LOAD/VP_STORE

[frasercrmck]

Summary

This patch fixes a compiler crash when widening scalable-vector loads
and stores which end up breaking down to element-wise store operations.
It does so by providing a way for targets with support for
vector-predicated loads and stores to use those instead. By widening the
operation but maintaining the original effective operation length via
the EVL, only the intended vector elements are loaded or stored.

This method should in theory be possible and even preferred for
fixed-length vector types, but all fixed-length types can be broken down
into their elements, and regardless I have observed regressions in the
generated code when doing so. I believe this is simply due to
VP_LOAD/VP_STORE not being up to par with LOAD/STORE in terms of
optimization. It does improve performance on smaller self-contained
examples, however, so the potential is there.

While the only target that benefits from this is RISCV, the legalization
is generic and so was placed centrally.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

5051

recusion -> recursion?

5066

I'm distrusting of changeVectorElementType. If WideVT happens to be an MVT, the changed type must also be an MVT since there is no way to get the LLVMContext from the MVT to make an EVT for the i1 vector.

Can we use EVT::getVectorVT here instead?

llvm/test/CodeGen/RISCV/rvv/load-store-sdnode.ll

40 ↗

(On Diff #377606)

How are we able to widen this load? Looks like we inferred an align of 4. I think we used that alignment to infer that it was ok to widen using the logic we have for fixed vectors. But I don't think that can work with vscale.

[frasercrmck]

• fix typo
• use EVT::getVectorVT

[frasercrmck]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

5051

Good spot, thanks.

5066

Quite right. It was my intention to fix that up before posting the patch, but I forgot

llvm/test/CodeGen/RISCV/rvv/load-store-sdnode.ll

40 ↗

(On Diff #377606)

I'm glad you pointed this out. I was pretty surprised by this (I was expecting it to crash like stores do). I then somehow convinced myself that it was okay but that was probably wishful thinking. Is the assumption that the hardware's addressable memory is always a multiple of the alignment and so if <3 x i8> is aligned to 4 there must be an extra byte that it can load without faulting?

If so I agree that vscale doesn't work like this. Do you think this should be patched quickly to crash like stores currently do? Feels pretty blunt but it's better than silently generating wrong code.

[craig.topper]

llvm/test/CodeGen/RISCV/rvv/load-store-sdnode.ll

40 ↗

(On Diff #377606)

Yeah if the alignment is >= than the size of the widened load it assumes it won't fault.

An easy fix might be to squash the alignment to 0 for scalable vectors here. I assume that's why volatile crashes.

// Allow wider loads if they are sufficiently aligned to avoid memory faults   
// and if the original load is simple.                                         
unsigned LdAlign = (!LD->isSimple()) ? 0 : LD->getAlignment();

[frasercrmck]

llvm/test/CodeGen/RISCV/rvv/load-store-sdnode.ll

40 ↗

(On Diff #377606)

I've (finally) created D111885 to fix the wrong code, and to serve as a basis for this work.

[frasercrmck]

• rebase on top of prep patch
• add VP_LOAD support too

[frasercrmck]

• rebase again

[frasercrmck]

• rebase

[frasercrmck]

• rebase
• fix accidental unreachable due to bad rebasing

[frasercrmck]

ping

[craig.topper]

llvm/test/CodeGen/RISCV/rvv/legalize-load-sdnode.ll

7

I thought you fixed vp.load with all ones mask in D113022. What is still missing?

[frasercrmck]

• rebase on top of D113022

[frasercrmck]

llvm/test/CodeGen/RISCV/rvv/legalize-load-sdnode.ll

7

Ach I thought I'd rebased. Done now, cheers.

[craig.topper]

LGTM

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

SDValue DAGTypeLegalizer::WidenVecRes_LOAD(SDNode *N) {
 
   SDValue Result;
   SmallVector<SDValue, 16> LdChain;  // Chain for the series of load
   if (ExtType != ISD::NON_EXTLOAD)
     Result = GenWidenVectorExtLoads(LdChain, LD, ExtType);
   else
     Result = GenWidenVectorLoads(LdChain, LD);
 
-  if (!Result)
-    llvm_unreachable("Unable to widen vector load");
+  if (Result) {
 
     // If we generate a single load, we can use that for the chain.  Otherwise,
     // build a factor node to remember the multiple loads are independent and
     // chain to that.
     SDValue NewChain;
     if (LdChain.size() == 1)
       NewChain = LdChain[0];
     else
       NewChain = DAG.getNode(ISD::TokenFactor, SDLoc(LD), MVT::Other, LdChain);
 
     // Modified the chain - switch anything that used the old chain to use
     // the new one.
     ReplaceValueWith(SDValue(N, 1), NewChain);
 
     return Result;
   }
 
+  // Generate a vector-predicated load if it is custom/legal on the target. To
+  // avoid possible recursion, only do this if the widened mask type is legal.
+  // FIXME: Not all targets may support EVL in VP_LOAD. These will have been
+  // removed from the IR by the ExpandVectorPredication pass but we're
+  // reintroducing them here.
+  EVT LdVT = LD->getMemoryVT();
+  EVT WideVT = TLI.getTypeToTransformTo(*DAG.getContext(), LdVT);
+  EVT WideMaskVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1,
+                                    WideVT.getVectorElementCount());
+  if (ExtType == ISD::NON_EXTLOAD && WideVT.isScalableVector() &&
+      TLI.isOperationLegalOrCustom(ISD::VP_LOAD, WideVT) &&
+      TLI.isTypeLegal(WideMaskVT)) {
+    SDLoc DL(N);
+    SDValue Mask = DAG.getAllOnesConstant(DL, WideMaskVT);
+    MVT EVLVT = TLI.getVPExplicitVectorLengthTy();
+    unsigned NumVTElts = LdVT.getVectorMinNumElements();
+    SDValue EVL =
+        DAG.getVScale(DL, EVLVT, APInt(EVLVT.getScalarSizeInBits(), NumVTElts));
+    const auto *MMO = LD->getMemOperand();
+    return DAG.getLoadVP(WideVT, DL, LD->getChain(), LD->getBasePtr(), Mask,
+                         EVL, MMO->getPointerInfo(), MMO->getAlign(),
+                         MMO->getFlags(), MMO->getAAInfo());
+  }
+
+  llvm_unreachable("Unable to widen vector load");
+}
+
 SDValue DAGTypeLegalizer::WidenVecRes_MLOAD(MaskedLoadSDNode *N) {
 
   EVT WidenVT = TLI.getTypeToTransformTo(*DAG.getContext(),N->getValueType(0));
   SDValue Mask = N->getMask();
   EVT MaskVT = Mask.getValueType();
   SDValue PassThru = GetWidenedVector(N->getPassThru());
   ISD::LoadExtType ExtType = N->getExtensionType();
   SDLoc dl(N);
SDValue DAGTypeLegalizer::WidenVecOp_STORE(SDNode *N) {
 
   if (!ST->getMemoryVT().getScalarType().isByteSized())
     return TLI.scalarizeVectorStore(ST, DAG);
 
   if (ST->isTruncatingStore())
     return TLI.scalarizeVectorStore(ST, DAG);
 
   SmallVector<SDValue, 16> StChain;
   if (GenWidenVectorStores(StChain, ST)) {
     if (StChain.size() == 1)
       return StChain[0];
 
     return DAG.getNode(ISD::TokenFactor, SDLoc(ST), MVT::Other, StChain);
   }
 
+  // Generate a vector-predicated store if it is custom/legal on the target.
+  // To avoid possible recursion, only do this if the widened mask type is
+  // legal.
+  // FIXME: Not all targets may support EVL in VP_STORE. These will have been
+  // removed from the IR by the ExpandVectorPredication pass but we're
+  // reintroducing them here.
+  SDValue StVal = ST->getValue();
+  EVT StVT = StVal.getValueType();
+  EVT WideVT = TLI.getTypeToTransformTo(*DAG.getContext(), StVT);
+  EVT WideMaskVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1,
+                                    WideVT.getVectorElementCount());
+  if (WideVT.isScalableVector() &&
+      TLI.isOperationLegalOrCustom(ISD::VP_STORE, WideVT) &&
+      TLI.isTypeLegal(WideMaskVT)) {
+    // Widen the value.
+    SDLoc DL(N);
+    StVal = GetWidenedVector(StVal);
+    SDValue Mask = DAG.getAllOnesConstant(DL, WideMaskVT);
+    MVT EVLVT = TLI.getVPExplicitVectorLengthTy();
+    unsigned NumVTElts = StVT.getVectorMinNumElements();
+    SDValue EVL =
+        DAG.getVScale(DL, EVLVT, APInt(EVLVT.getScalarSizeInBits(), NumVTElts));
+    const auto *MMO = ST->getMemOperand();
+    return DAG.getStoreVP(ST->getChain(), DL, StVal, ST->getBasePtr(), Mask,
+                          EVL, MMO->getPointerInfo(), MMO->getAlign(),
+                          MMO->getFlags(), MMO->getAAInfo());
+  }
+
   llvm_unreachable("Unable to widen vector store");
 }
 
 SDValue DAGTypeLegalizer::WidenVecOp_MSTORE(SDNode *N, unsigned OpNo) {
   assert((OpNo == 1 || OpNo == 3) &&
          "Can widen only data or mask operand of mstore");
   MaskedStoreSDNode *MST = cast<MaskedStoreSDNode>(N);
   SDValue Mask = MST->getMask();

llvm/test/CodeGen/RISCV/rvv/legalize-load-sdnode.ll

-; RUN: not --crash llc -mtriple=riscv32 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s
-; RUN: not --crash llc -mtriple=riscv64 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=riscv64 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s | FileCheck %s
 
 ; Check that we are able to legalize scalable-vector loads that require widening.
 
-; FIXME: LLVM can't yet widen scalable-vector loads.
+; FIXME: We could be using unmasked loads here.
 
 define <vscale x 3 x i8> @load_nxv3i8(<vscale x 3 x i8>* %ptr) {
+; CHECK-LABEL: load_nxv3i8:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a1, vlenb
+; CHECK-NEXT:    srli a1, a1, 3
+; CHECK-NEXT:    slli a2, a1, 1
+; CHECK-NEXT:    add a1, a2, a1
+; CHECK-NEXT:    vsetvli a2, zero, e8, mf2, ta, mu
+; CHECK-NEXT:    vmset.m v0
+; CHECK-NEXT:    vsetvli zero, a1, e8, mf2, ta, mu
+; CHECK-NEXT:    vle8.v v8, (a0), v0.t
+; CHECK-NEXT:    ret
   %v = load <vscale x 3 x i8>, <vscale x 3 x i8>* %ptr
   ret <vscale x 3 x i8> %v
 }
 
 define <vscale x 5 x half> @load_nxv5f16(<vscale x 5 x half>* %ptr) {
+; CHECK-LABEL: load_nxv5f16:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a1, vlenb
+; CHECK-NEXT:    srli a1, a1, 3
+; CHECK-NEXT:    slli a2, a1, 2
+; CHECK-NEXT:    add a1, a2, a1
+; CHECK-NEXT:    vsetvli a2, zero, e8, m1, ta, mu
+; CHECK-NEXT:    vmset.m v0
+; CHECK-NEXT:    vsetvli zero, a1, e16, m2, ta, mu
+; CHECK-NEXT:    vle16.v v8, (a0), v0.t
+; CHECK-NEXT:    ret
   %v = load <vscale x 5 x half>, <vscale x 5 x half>* %ptr
   ret <vscale x 5 x half> %v
 }

llvm/test/CodeGen/RISCV/rvv/legalize-store-sdnode.ll

-; RUN: not --crash llc -mtriple=riscv32 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s
-; RUN: not --crash llc -mtriple=riscv64 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=riscv64 -mattr=+m,+experimental-v,+experimental-zfh,+f,+d -verify-machineinstrs < %s | FileCheck %s
 
 ; Check that we are able to legalize scalable-vector stores that require widening.
 
-; FIXME: LLVM can't yet widen scalable-vector stores.
-
 define void @store_nxv3i8(<vscale x 3 x i8> %val, <vscale x 3 x i8>* %ptr) {
+; CHECK-LABEL: store_nxv3i8:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a1, vlenb
+; CHECK-NEXT:    srli a1, a1, 3
+; CHECK-NEXT:    slli a2, a1, 1
+; CHECK-NEXT:    add a1, a2, a1
+; CHECK-NEXT:    vsetvli zero, a1, e8, mf2, ta, mu
+; CHECK-NEXT:    vse8.v v8, (a0)
+; CHECK-NEXT:    ret
   store <vscale x 3 x i8> %val, <vscale x 3 x i8>* %ptr
   ret void
 }
 
 define void @store_nxv7f64(<vscale x 7 x double> %val, <vscale x 7 x double>* %ptr) {
+; CHECK-LABEL: store_nxv7f64:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a1, vlenb
+; CHECK-NEXT:    srli a1, a1, 3
+; CHECK-NEXT:    slli a2, a1, 3
+; CHECK-NEXT:    sub a1, a2, a1
+; CHECK-NEXT:    vsetvli zero, a1, e64, m8, ta, mu
+; CHECK-NEXT:    vse64.v v8, (a0)
+; CHECK-NEXT:    ret
   store <vscale x 7 x double> %val, <vscale x 7 x double>* %ptr
   ret void
 }