[llvm] [CodeGen] Fixed vector halving bug for masked load

[aartbik]

Summary

Given a VL=14 that is enveloped by a proper VL=16, splitting the
masked load using the enveloping halving VL=8/8 should yields
should eventually yield V=8/5. This fixes various assert failures
in getHalfNumVectorElementsVT() and IncrementMemoryAddress().

Note, I suspect similar fixes will be needed for other masked
operations, but for now I send out a fix for masked load only.

Bugzilla issue 45563
https://bugs.llvm.org/show_bug.cgi?id=45563

[craig.topper]

I'm concerned that there might be a way to get HiMemVT to contain 0 elements. I think v17i32 would trigger it. We'd widen to v32i32, then split to v16i32 which will have a memvt of v1i32 and v16i32. Then we'll split the v16i32 pieces. And now we'll try to split that v1i32 memory vt.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9468

EVTs should be passed by value.

[aartbik]

In D78608#1996059, @craig.topper wrote:

I'm concerned that there might be a way to get HiMemVT to contain 0 elements. I think v17i32 would trigger it. We'd widen to v32i32, then split to v16i32 which will have a memvt of v1i32 and v16i32. Then we'll split the v16i32 pieces. And now we'll try to split that v1i32 memory vt.

But do you agree this solution goes in the right direction?
Do you have any suggestions for the v17i32 case?
Happy to test that tomorrow and think about it some more.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9468

I mimicked the prototypes of related methods in this file? Are those wrong too?
Happy to change, just checking what best practices are here?

[craig.topper]

In D78608#1996085, @aartbik wrote:

In D78608#1996059, @craig.topper wrote:

I'm concerned that there might be a way to get HiMemVT to contain 0 elements. I think v17i32 would trigger it. We'd widen to v32i32, then split to v16i32 which will have a memvt of v1i32 and v16i32. Then we'll split the v16i32 pieces. And now we'll try to split that v1i32 memory vt.

But do you agree this solution goes in the right direction?
Do you have any suggestions for the v17i32 case?
Happy to test that tomorrow and think about it some more.

I wonder if we should fix this in WidenVecRes_MLOAD by doing something closer to WidenVecRes_LOAD. Finding a legal type and slicing up the vector as we widen it.

[aartbik]

do not generate zero storage size masked loads

[aartbik]

PTAL

You had a very sharp eye that v17f32 would cause a split with an empty hi masked load eventually.
It is somewhat harmless since the load is not used and thus removed from the DAG, but this revision
is even cleaner by not generating the hi masked load at all!

[aartbik]

removed left behind debug code

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9495

Stray 'k' got added to this line.

[aartbik]

fixed some stray diffs

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

1596

Could we return PassThruHi directly?

llvm/test/CodeGen/X86/pr45563-2.ll

2

Please use update_llc_test_checks on tests this size

[aartbik]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

1596

I experimented with different values, this seemed the least code change, since it folds directly in the DAG.getNode below, and is removed right after that. Any thing else would require more changes to deal with the 1 vs 2 result values.

llvm/test/CodeGen/X86/pr45563-2.ll

2

Ah, I saw some auto generated messages, but was not aware of this tool. But, forgive my noobness, what will this auto-generate? I am really just counting the number of masked moves that result, I don't want to pin down the full assembly?

[craig.topper]

llvm/test/CodeGen/X86/pr45563-2.ll

2

It will generate checks for the full assembly, but the X86 maintainers are used to that. We'll just rerun the script and look at the diff if it fails for some change in the future.

[aartbik]

used update_llc_test_checks (cool!)

[aartbik]

llvm/test/CodeGen/X86/pr45563-2.ll

2

Nice utility! Having the full assembly makes it a bit sensitive to changes, but I can also see that the X86 maintainers want to be aware of even the slightest of changes and with this utility that becomes a lot simpler. Cool!

[craig.topper]

LGTM

llvm/include/llvm/CodeGen/SelectionDAG.h

if (auto A = InferPtrAlign(Ptr))
       return A->value();
     return 0;
   }
 
   /// Compute the VTs needed for the low/hi parts of a type
   /// which is split (or expanded) into two not necessarily identical pieces.
   std::pair<EVT, EVT> GetSplitDestVTs(const EVT &VT) const;
 
+  /// Compute the VTs needed for the low/hi parts of a type, dependent on an
+  /// enveloping VT that has been split into two identical pieces. Sets the
+  /// HisIsEmpty flag when hi type has zero storage size.
+  std::pair<EVT, EVT> GetDependentSplitDestVTs(const EVT &VT, const EVT &EnvVT,
+                                               bool *HiIsEmpty) const;
+
   /// Split the vector with EXTRACT_SUBVECTOR using the provides
   /// VTs and return the low/high part.
   std::pair<SDValue, SDValue> SplitVector(const SDValue &N, const SDLoc &DL,
                                           const EVT &LoVT, const EVT &HiVT);
 
   /// Split the vector with EXTRACT_SUBVECTOR and return the low/high part.
   std::pair<SDValue, SDValue> SplitVector(const SDValue &N, const SDLoc &DL) {
     EVT LoVT, HiVT;

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

if (Mask.getOpcode() == ISD::SETCC) {
     if (getTypeAction(Mask.getValueType()) == TargetLowering::TypeSplitVector)
       GetSplitVector(Mask, MaskLo, MaskHi);
     else
       std::tie(MaskLo, MaskHi) = DAG.SplitVector(Mask, dl);
   }
 
   EVT MemoryVT = MLD->getMemoryVT();
   EVT LoMemVT, HiMemVT;
-  std::tie(LoMemVT, HiMemVT) = DAG.GetSplitDestVTs(MemoryVT);
+  bool HiIsEmpty = false;
+  std::tie(LoMemVT, HiMemVT) =
+      DAG.GetDependentSplitDestVTs(MemoryVT, LoVT, &HiIsEmpty);
 
   SDValue PassThruLo, PassThruHi;
   if (getTypeAction(PassThru.getValueType()) == TargetLowering::TypeSplitVector)
     GetSplitVector(PassThru, PassThruLo, PassThruHi);
   else
     std::tie(PassThruLo, PassThruHi) = DAG.SplitVector(PassThru, dl);
 
   MachineMemOperand *MMO = DAG.getMachineFunction().getMachineMemOperand(
       MLD->getPointerInfo(), MachineMemOperand::MOLoad, LoMemVT.getStoreSize(),
       Alignment, MLD->getAAInfo(), MLD->getRanges());
 
   Lo = DAG.getMaskedLoad(LoVT, dl, Ch, Ptr, Offset, MaskLo, PassThruLo, LoMemVT,
                          MMO, MLD->getAddressingMode(), ExtType,
                          MLD->isExpandingLoad());
 
+  if (HiIsEmpty) {
+    // The hi masked load has zero storage size. We therefore simply set
+    // it to the low masked load and rely on subsequent optimization to
+    // remove the the unused element in the chain.
+    Hi = Lo;
+  } else {
+    // Generate hi masked load.
     Ptr = TLI.IncrementMemoryAddress(Ptr, MaskLo, dl, LoMemVT, DAG,
                                      MLD->isExpandingLoad());
     unsigned HiOffset = LoMemVT.getStoreSize();
 
     MMO = DAG.getMachineFunction().getMachineMemOperand(
-      MLD->getPointerInfo().getWithOffset(HiOffset), MachineMemOperand::MOLoad,
-      HiMemVT.getStoreSize(), Alignment, MLD->getAAInfo(), MLD->getRanges());
+        MLD->getPointerInfo().getWithOffset(HiOffset),
+        MachineMemOperand::MOLoad, HiMemVT.getStoreSize(), Alignment,
+        MLD->getAAInfo(), MLD->getRanges());
 
-  Hi = DAG.getMaskedLoad(HiVT, dl, Ch, Ptr, Offset, MaskHi, PassThruHi, HiMemVT,
-                         MMO, MLD->getAddressingMode(), ExtType,
+    Hi = DAG.getMaskedLoad(HiVT, dl, Ch, Ptr, Offset, MaskHi, PassThruHi,
+                           HiMemVT, MMO, MLD->getAddressingMode(), ExtType,
                            MLD->isExpandingLoad());
+  }
 
   // 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.
   ReplaceValueWith(SDValue(MLD, 1), Ch);
-
 }
 
 void DAGTypeLegalizer::SplitVecRes_MGATHER(MaskedGatherSDNode *MGT,
                                          SDValue &Lo, SDValue &Hi) {
   EVT LoVT, HiVT;
   SDLoc dl(MGT);
   std::tie(LoVT, HiVT) = DAG.GetSplitDestVTs(MGT->getValueType(0));
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

std::pair<EVT, EVT> SelectionDAG::GetSplitDestVTs(const EVT &VT) const {
   if (!VT.isVector())
     LoVT = HiVT = TLI->getTypeToTransformTo(*getContext(), VT);
   else
     LoVT = HiVT = VT.getHalfNumVectorElementsVT(*getContext());
 
   return std::make_pair(LoVT, HiVT);
 }
 
-/// SplitVector - Split the vector with EXTRACT_SUBVECTOR and return the
+/// GetDependentSplitDestVTs - Compute the VTs needed for the low/hi parts of a
+/// type, dependent on an enveloping VT that has been split into two identical
+/// pieces. Sets the HiIsEmpty flag when hi type has zero storage size.
+std::pair<EVT, EVT>
+SelectionDAG::GetDependentSplitDestVTs(const EVT &VT, const EVT &EnvVT,
+                                       bool *HiIsEmpty) const {
+  EVT EltTp = VT.getVectorElementType();
+  bool IsScalable = VT.isScalableVector();
+  // Examples:
+  //   custom VL=8  with enveloping VL=8/8 yields 8/0 (hi empty)
+  //   custom VL=9  with enveloping VL=8/8 yields 8/1
+  //   custom VL=10 with enveloping VL=8/8 yields 8/2
+  //   etc.
+  unsigned VTNumElts = VT.getVectorNumElements();
+  unsigned EnvNumElts = EnvVT.getVectorNumElements();
+  EVT LoVT, HiVT;
+  if (VTNumElts > EnvNumElts) {
+    LoVT = EnvVT;
+    HiVT = EVT::getVectorVT(*getContext(), EltTp, VTNumElts - EnvNumElts,
+                            IsScalable);
+    *HiIsEmpty = false;
+  } else {
+    // Flag that hi type has zero storage size, but return split envelop type
+    // (this would be easier if vector types with zero elements were allowed).
+    LoVT = EVT::getVectorVT(*getContext(), EltTp, VTNumElts, IsScalable);
+    HiVT = EnvVT;
+    *HiIsEmpty = true;
+  }
+  return std::make_pair(LoVT, HiVT);
+}
+
+/// SplitVector - kSplit the vector with EXTRACT_SUBVECTOR and return the
 /// low/high part.
 std::pair<SDValue, SDValue>
 SelectionDAG::SplitVector(const SDValue &N, const SDLoc &DL, const EVT &LoVT,
                           const EVT &HiVT) {
   assert(LoVT.getVectorNumElements() + HiVT.getVectorNumElements() <=
          N.getValueType().getVectorNumElements() &&
          "More vector elements requested than available!");
   SDValue Lo, Hi;

llvm/test/CodeGen/X86/pr45563-2.ll

+; RUN: llc < %s -O3 -mattr=avx | FileCheck %s
+
+; Bug 45563:
+; The SplitVecRes_MLOAD method should split a extended value type
+; according to the halving of the enveloping type to avoid all sorts
+; of inconsistencies downstream. For example for a extended value type
+; with VL=14 and enveloping type VL=16 that is split 8/8, the extended
+; type should be split 8/6 and not 7/7. This also accounts for hi masked
+; load that get zero storage size (and are unused).
+
+define <9 x float> @mload_split9(<9 x i1> %mask, <9 x float>* %addr, <9 x float> %dst) {
+; CHECK-LABEL: load_split9:
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK-NOT:   vmaskmovps
+  %res = call <9 x float> @llvm.masked.load.v9f32.p0v9f32(<9 x float>* %addr, i32 4, <9 x i1>%mask, <9 x float> %dst)
+  ret <9 x float> %res
+}
+
+define <13 x float> @mload_split13(<13 x i1> %mask, <13 x float>* %addr, <13 x float> %dst) {
+; CHECK-LABEL: mload_split13:
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK-NOT:   vmaskmovps
+  %res = call <13 x float> @llvm.masked.load.v13f32.p0v13f32(<13 x float>* %addr, i32 4, <13 x i1>%mask, <13 x float> %dst)
+  ret <13 x float> %res
+}
+
+define <14 x float> @mload_split14(<14 x i1> %mask, <14 x float>* %addr, <14 x float> %dst) {
+; CHECK-LABEL: mload_split14:
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK-NOT:   vmaskmovps
+  %res = call <14 x float> @llvm.masked.load.v14f32.p0v14f32(<14 x float>* %addr, i32 4, <14 x i1>%mask, <14 x float> %dst)
+  ret <14 x float> %res
+}
+
+define <17 x float> @mload_split17(<17 x i1> %mask, <17 x float>* %addr, <17 x float> %dst) {
+; CHECK-LABEL: mload_split17:
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK-NOT:   vmaskmovps
+  %res = call <17 x float> @llvm.masked.load.v17f32.p0v17f32(<17 x float>* %addr, i32 4, <17 x i1>%mask, <17 x float> %dst)
+  ret <17 x float> %res
+}
+
+define <23 x float> @mload_split23(<23 x i1> %mask, <23 x float>* %addr, <23 x float> %dst) {
+; CHECK-LABEL: mload_split23:
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK:       vmaskmovps
+; CHECK-NOT:   vmaskmovps
+  %res = call <23 x float> @llvm.masked.load.v23f32.p0v23f32(<23 x float>* %addr, i32 4, <23 x i1>%mask, <23 x float> %dst)
+  ret <23 x float> %res
+}
+
+declare <9 x float> @llvm.masked.load.v9f32.p0v9f32(<9 x float>* %addr, i32 %align, <9 x i1> %mask, <9 x float> %dst)
+declare <13 x float> @llvm.masked.load.v13f32.p0v13f32(<13 x float>* %addr, i32 %align, <13 x i1> %mask, <13 x float> %dst)
+declare <14 x float> @llvm.masked.load.v14f32.p0v14f32(<14 x float>* %addr, i32 %align, <14 x i1> %mask, <14 x float> %dst)
+declare <17 x float> @llvm.masked.load.v17f32.p0v17f32(<17 x float>* %addr, i32 %align, <17 x i1> %mask, <17 x float> %dst)
+declare <23 x float> @llvm.masked.load.v23f32.p0v23f32(<23 x float>* %addr, i32 %align, <23 x i1> %mask, <23 x float> %dst)