[AArch64] Preserve chain when lowering fixed length load to SVE (PR55281)

[nikic]

Summary

When a fixed length load is lowered to an SVE masked load, the result chain is currently set to the input chain of the old load, rather than the result chain of the new load. This may cause stores to be incorrectly reordered.

Fixes https://github.com/llvm/llvm-project/issues/55281.

[DavidSpickett]

Thanks for the fix!

Don't pay much attention to the SVE VLS bot still being red, we have multiple issues going on at the same time that will mask this change for a while. I tested this locally, it fixes my reproducer and the unit tests that were failing on the bot at the time.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerFixedLengthVectorLoadToSVE(
 
   auto Pg = getPredicateForFixedLengthVector(DAG, DL, VT);
 
   if (VT.isFloatingPoint() && Load->getExtensionType() == ISD::EXTLOAD) {
     LoadVT = ContainerVT.changeTypeToInteger();
     MemVT = MemVT.changeTypeToInteger();
   }
 
-  auto NewLoad = DAG.getMaskedLoad(
+  SDValue NewLoad = DAG.getMaskedLoad(
       LoadVT, DL, Load->getChain(), Load->getBasePtr(), Load->getOffset(), Pg,
       DAG.getUNDEF(LoadVT), MemVT, Load->getMemOperand(),
       Load->getAddressingMode(), Load->getExtensionType());
 
+  SDValue Result = NewLoad;
   if (VT.isFloatingPoint() && Load->getExtensionType() == ISD::EXTLOAD) {
     EVT ExtendVT = ContainerVT.changeVectorElementType(
         Load->getMemoryVT().getVectorElementType());
 
-    NewLoad = getSVESafeBitCast(ExtendVT, NewLoad, DAG);
-    NewLoad = DAG.getNode(AArch64ISD::FP_EXTEND_MERGE_PASSTHRU, DL, ContainerVT,
-                          Pg, NewLoad, DAG.getUNDEF(ContainerVT));
+    Result = getSVESafeBitCast(ExtendVT, Result, DAG);
+    Result = DAG.getNode(AArch64ISD::FP_EXTEND_MERGE_PASSTHRU, DL, ContainerVT,
+                         Pg, Result, DAG.getUNDEF(ContainerVT));
   }
 
-  auto Result = convertFromScalableVector(DAG, VT, NewLoad);
-  SDValue MergedValues[2] = {Result, Load->getChain()};
+  Result = convertFromScalableVector(DAG, VT, Result);
+  SDValue MergedValues[2] = {Result, NewLoad.getValue(1)};
   return DAG.getMergeValues(MergedValues, DL);
 }
 
 static SDValue convertFixedMaskToScalableVector(SDValue Mask,
                                                 SelectionDAG &DAG) {
   SDLoc DL(Mask);
   EVT InVT = Mask.getValueType();
   EVT ContainerVT = getContainerForFixedLengthVector(DAG, InVT);
if (Load->getPassThru()->isUndef()) {
     if (ContainerVT.isInteger())
       PassThru = DAG.getConstant(0, DL, ContainerVT);
     else
       PassThru = DAG.getConstantFP(0, DL, ContainerVT);
     if (isZerosVector(Load->getPassThru().getNode()))
       IsPassThruZeroOrUndef = true;
   }
 
-  auto NewLoad = DAG.getMaskedLoad(
+  SDValue NewLoad = DAG.getMaskedLoad(
       ContainerVT, DL, Load->getChain(), Load->getBasePtr(), Load->getOffset(),
       Mask, PassThru, Load->getMemoryVT(), Load->getMemOperand(),
       Load->getAddressingMode(), Load->getExtensionType());
 
+  SDValue Result = NewLoad;
   if (!IsPassThruZeroOrUndef) {
     SDValue OldPassThru =
         convertToScalableVector(DAG, ContainerVT, Load->getPassThru());
-    NewLoad = DAG.getSelect(DL, ContainerVT, Mask, NewLoad, OldPassThru);
+    Result = DAG.getSelect(DL, ContainerVT, Mask, Result, OldPassThru);
   }
 
-  auto Result = convertFromScalableVector(DAG, VT, NewLoad);
-  SDValue MergedValues[2] = {Result, Load->getChain()};
+  Result = convertFromScalableVector(DAG, VT, Result);
+  SDValue MergedValues[2] = {Result, NewLoad.getValue(1)};
   return DAG.getMergeValues(MergedValues, DL);
 }
 
 // Convert all fixed length vector stores larger than NEON to masked_stores.
 SDValue AArch64TargetLowering::LowerFixedLengthVectorStoreToSVE(
     SDValue Op, SelectionDAG &DAG) const {
   auto Store = cast<StoreSDNode>(Op);
 

llvm/test/CodeGen/AArch64/sve-fixed-length-frame-offests.ll

 define void @foo(<8 x i64>* %a) #0 {
 ; CHECK-LABEL: foo:
 ; CHECK:       SelectionDAG has 14 nodes:
 ; CHECK-NEXT:    t0: ch = EntryToken
 ; CHECK-NEXT:    t12: nxv2i1 = PTRUE_D TargetConstant:i32<31>
 ; CHECK-NEXT:    t2: i64,ch = CopyFromReg t0, Register:i64 %0
 ; CHECK-NEXT:    t18: nxv2i64,ch = LD1D_IMM<Mem:(volatile load (s512) from %ir.a)> t12, t2, TargetConstant:i64<0>, t0
 ; CHECK-NEXT:    t8: i64 = ADDXri TargetFrameIndex:i64<1>, TargetConstant:i32<0>, TargetConstant:i32<0>
-; CHECK-NEXT:    t17: ch = ST1D_IMM<Mem:(volatile store (s512) into %ir.r0)> t18, t12, TargetFrameIndex:i64<0>, TargetConstant:i64<0>, t0
+; CHECK-NEXT:    t17: ch = ST1D_IMM<Mem:(volatile store (s512) into %ir.r0)> t18, t12, TargetFrameIndex:i64<0>, TargetConstant:i64<0>, t18:1
 ; CHECK-NEXT:    t16: ch = ST1D_IMM<Mem:(volatile store (s512) into %ir.r1)> t18, t12, t8, TargetConstant:i64<0>, t17
 ; CHECK-NEXT:    t10: ch = RET_ReallyLR t16
 ; CHECK-EMPTY:
 entry:
   %r0 = alloca <8 x i64>
   %r1 = alloca <8 x i64>
   %r = load volatile <8 x i64>, <8 x i64>* %a
   store volatile <8 x i64> %r, <8 x i64>* %r0
   store volatile <8 x i64> %r, <8 x i64>* %r1
   ret void
 }
 
 attributes #0 = { nounwind "target-features"="+sve" vscale_range(4,4) }