[SVE][CodeGen] Fix incorrect legalisation of zero-extended masked loads

[kmclaughlin]

Summary

PromoteIntRes_MLOAD always sets the extension type to EXTLOAD, which
results in a sign-extended load. If the type returned by getExtensionType()
for the load being promoted is something other than NON_EXTLOAD, we
should instead pass this to getMaskedLoad() as the extension type.

[CarolineConcatto]

Hey @kmclaughlin,
Thank you for adding me as a reviewer of the patch.
Do we know what was the logic to set ExtType to ISD::EXTLOAD?
I don't see that on other parts of the legalisation for Masked load.

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

791

As we are replacing ExtType, do we need to pass isExpanding?
I can see in void DAGTypeLegalizer::SplitVecRes_MLOAD and DAGTypeLegalizer::WidenVecRes_MLOAD
they have a
XLD->isExpandingLoad()

[kmclaughlin]

• Pass isExpandingLoad() to getMaskedLoad

[kmclaughlin]

In D112320#3080828, @CarolineConcatto wrote:

Do we know what was the logic to set ExtType to ISD::EXTLOAD?
I don't see that on other parts of the legalisation for Masked load.

Hi @CarolineConcatto, thank you for taking a look at this. We are promoting the masked load by extending the return type, which is why we need to set ExtType when we create the new masked load with the promoted type. The original masked load which had an illegal return type may not have been extending, e.g.

nxv2i16,ch = masked_load<(load (s32) from %ir.in, align 2)> t0, t2, undef:i64, t4, undef:nxv2i16
->
nxv2i64,ch = masked_load<(load (s32) from %ir.in, align 2), anyext from nxv2i16> t0, t2, undef:i64, t4, undef:nxv2i64

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

791

I think we should be passing isExpandingLoad here, I can see in PromoteIntOp_MSTORE we similarly pass N->isCompressingStore()

[CarolineConcatto]

Thank you @kmclaughlin

[david-arm]

LGTM, moi aussi!

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

796

nit: Perhaps just address the formatting issue before merging? Thanks!

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

SDValue DAGTypeLegalizer::PromoteIntRes_LOAD(LoadSDNode *N) {
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
   return Res;
 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_MLOAD(MaskedLoadSDNode *N) {
   EVT NVT = TLI.getTypeToTransformTo(*DAG.getContext(), N->getValueType(0));
   SDValue ExtPassThru = GetPromotedInteger(N->getPassThru());
 
+  ISD::LoadExtType ExtType = N->getExtensionType();
+  if (ExtType == ISD::NON_EXTLOAD)
+    ExtType = ISD::EXTLOAD;
+
   SDLoc dl(N);
   SDValue Res = DAG.getMaskedLoad(NVT, dl, N->getChain(), N->getBasePtr(),
                                   N->getOffset(), N->getMask(), ExtPassThru,
                                   N->getMemoryVT(), N->getMemOperand(),
-                                  N->getAddressingMode(), ISD::EXTLOAD);
+                                  N->getAddressingMode(), ExtType,
+                                  N->isExpandingLoad());
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
   return Res;
 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_MGATHER(MaskedGatherSDNode *N) {
   EVT NVT = TLI.getTypeToTransformTo(*DAG.getContext(), N->getValueType(0));

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

 ; CHECK-LABEL: masked_load_passthru:
 ; CHECK-NEXT: ld1w { z1.s }, p0/z, [x0]
 ; CHECK-NEXT: mov z0.s, p0/m, z1.s
 ; CHECK-NEXT: ret
   %load = call <vscale x 4 x i32> @llvm.masked.load.nxv4i32(<vscale x 4 x i32> *%a, i32 4, <vscale x 4 x i1> %mask, <vscale x 4 x i32> %passthru)
   ret <vscale x 4 x i32> %load
 }
 
+; Masked load requires promotion
+define <vscale x 2 x i16> @masked_load_nxv2i16(<vscale x 2 x i16>* noalias %in, <vscale x 2 x i1> %mask) {
+; CHECK-LABEL: masked_load_nxv2i16
+; CHECK:       ld1sh { z0.d }, p0/z, [x0]
+; CHECK-NEXT:  ret
+  %wide.load = call <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>* %in, i32 2, <vscale x 2 x i1> %mask, <vscale x 2 x i16> undef)
+  ret <vscale x 2 x i16> %wide.load
+}
+
 ;
 ; Masked Stores
 ;
 
 define void @masked_store_nxv2i64(<vscale x 2 x i64> *%a, <vscale x 2 x i64> %val, <vscale x 2 x i1> %mask) nounwind {
 ; CHECK-LABEL: masked_store_nxv2i64:
 ; CHECK-NEXT: st1d { z0.d }, p0, [x0]
 ; CHECK-NEXT: ret

 ; CHECK-NEXT:    ret
   call void @llvm.masked.store.nxv2p0s_struct.p0nxv2p0s_struct(<vscale x 2 x %struct*> %data, <vscale x 2 x %struct*>* %vector_ptr, i32 8, <vscale x 2 x i1> %mask)
   ret void
 }
 
 
 declare <vscale x 2 x i64> @llvm.masked.load.nxv2i64(<vscale x 2 x i64>*, i32, <vscale x 2 x i1>, <vscale x 2 x i64>)
 declare <vscale x 4 x i32> @llvm.masked.load.nxv4i32(<vscale x 4 x i32>*, i32, <vscale x 4 x i1>, <vscale x 4 x i32>)
+declare <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>*, i32, <vscale x 2 x i1>, <vscale x 2 x i16>)
 declare <vscale x 8 x i16> @llvm.masked.load.nxv8i16(<vscale x 8 x i16>*, i32, <vscale x 8 x i1>, <vscale x 8 x i16>)
 declare <vscale x 16 x i8> @llvm.masked.load.nxv16i8(<vscale x 16 x i8>*, i32, <vscale x 16 x i1>, <vscale x 16 x i8>)
 
 declare <vscale x 2 x double> @llvm.masked.load.nxv2f64(<vscale x 2 x double>*, i32, <vscale x 2 x i1>, <vscale x 2 x double>)
 declare <vscale x 2 x float> @llvm.masked.load.nxv2f32(<vscale x 2 x float>*, i32, <vscale x 2 x i1>, <vscale x 2 x float>)
 declare <vscale x 2 x half> @llvm.masked.load.nxv2f16(<vscale x 2 x half>*, i32, <vscale x 2 x i1>, <vscale x 2 x half>)
 declare <vscale x 4 x float> @llvm.masked.load.nxv4f32(<vscale x 4 x float>*, i32, <vscale x 4 x i1>, <vscale x 4 x float>)
 declare <vscale x 4 x half> @llvm.masked.load.nxv4f16(<vscale x 4 x half>*, i32, <vscale x 4 x i1>, <vscale x 4 x half>)

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

 ; CHECK-NEXT:    ld1sb { z2.s }, p2/z, [x0, #2, mul vl]
 ; CHECK-NEXT:    ld1sb { z3.s }, p0/z, [x0, #3, mul vl]
 ; CHECK-NEXT:    ret
   %load = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8(<vscale x 16 x i8>* %a, i32 2, <vscale x 16 x i1> %mask, <vscale x 16 x i8> undef)
   %ext = sext <vscale x 16 x i8> %load to <vscale x 16 x i32>
   ret <vscale x 16 x i32> %ext
 }
 
+; Masked load requires promotion
+define <vscale x 4 x double> @masked_sload_4i8_4f32(<vscale x 4 x i8>* noalias %in, <vscale x 4 x i1> %mask) {
+; CHECK-LABEL: masked_sload_4i8_4f32:
+; CHECK:       punpkhi p2.h, p0.b
+; CHECK-NEXT:  punpklo p0.h, p0.b
+; CHECK-NEXT:  ld1sb { z1.d }, p2/z, [x0, #1, mul vl]
+; CHECK-NEXT:  ld1sb { z0.d }, p0/z, [x0]
+; CHECK-NEXT:  ptrue p1.d
+; CHECK-NEXT:  scvtf z0.d, p1/m, z0.d
+; CHECK-NEXT:  scvtf z1.d, p1/m, z1.d
+; CHECK-NEXT:  ret
+  %wide.load = call <vscale x 4 x i8> @llvm.masked.load.nxv4i8(<vscale x 4 x i8>* %in, i32 2, <vscale x 4 x i1> %mask, <vscale x 4 x i8> undef)
+  %sext = sext <vscale x 4 x i8> %wide.load to <vscale x 4 x i64>
+  %res = sitofp <vscale x 4 x i64> %sext to <vscale x 4 x double>
+  ret <vscale x 4 x double> %res
+}
+
 declare <vscale x 2 x i8> @llvm.masked.load.nxv2i8(<vscale x 2 x i8>*, i32, <vscale x 2 x i1>, <vscale x 2 x i8>)
 declare <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>*, i32, <vscale x 2 x i1>, <vscale x 2 x i16>)
 declare <vscale x 2 x i32> @llvm.masked.load.nxv2i32(<vscale x 2 x i32>*, i32, <vscale x 2 x i1>, <vscale x 2 x i32>)
 declare <vscale x 4 x i8> @llvm.masked.load.nxv4i8(<vscale x 4 x i8>*, i32, <vscale x 4 x i1>, <vscale x 4 x i8>)
 declare <vscale x 4 x i16> @llvm.masked.load.nxv4i16(<vscale x 4 x i16>*, i32, <vscale x 4 x i1>, <vscale x 4 x i16>)
 declare <vscale x 8 x i8> @llvm.masked.load.nxv8i8(<vscale x 8 x i8>*, i32, <vscale x 8 x i1>, <vscale x 8 x i8>)
 declare <vscale x 16 x i8> @llvm.masked.load.nxv16i8(<vscale x 16 x i8>*, i32, <vscale x 16 x i1>, <vscale x 16 x i8>)

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

 ; CHECK-NEXT:  ld1h { z2.d }, p2/z, [x0, #2, mul vl]
 ; CHECK-NEXT:  ld1h { z3.d }, p0/z, [x0, #3, mul vl]
 ; CHECK-NEXT:  ret
   %load = call <vscale x 8 x i16> @llvm.masked.load.nxv8i16(<vscale x 8 x i16>* %a, i32 2, <vscale x 8 x i1> %mask, <vscale x 8 x i16> undef)
   %ext = zext <vscale x 8 x i16> %load to <vscale x 8 x i64>
   ret <vscale x 8 x i64> %ext
 }
 
+; Masked load requires promotion
+define <vscale x 2 x double> @masked_zload_2i16_2f64(<vscale x 2 x i16>* noalias %in, <vscale x 2 x i1> %mask) {
+; CHECK-LABEL: masked_zload_2i16_2f64:
+; CHECK:       ld1h { z0.d }, p0/z, [x0]
+; CHECK-NEXT:  ptrue p0.d
+; CHECK-NEXT:  ucvtf z0.d, p0/m, z0.s
+; CHECK-NEXT:  ret
+  %wide.load = call <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>* %in, i32 2, <vscale x 2 x i1> %mask, <vscale x 2 x i16> undef)
+  %zext = zext <vscale x 2 x i16> %wide.load to <vscale x 2 x i32>
+  %res = uitofp <vscale x 2 x i32> %zext to <vscale x 2 x double>
+  ret <vscale x 2 x double> %res
+}
+
 declare <vscale x 2 x i8> @llvm.masked.load.nxv2i8(<vscale x 2 x i8>*, i32, <vscale x 2 x i1>, <vscale x 2 x i8>)
 declare <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>*, i32, <vscale x 2 x i1>, <vscale x 2 x i16>)
 declare <vscale x 2 x i32> @llvm.masked.load.nxv2i32(<vscale x 2 x i32>*, i32, <vscale x 2 x i1>, <vscale x 2 x i32>)
 declare <vscale x 4 x i8> @llvm.masked.load.nxv4i8(<vscale x 4 x i8>*, i32, <vscale x 4 x i1>, <vscale x 4 x i8>)
 declare <vscale x 4 x i16> @llvm.masked.load.nxv4i16(<vscale x 4 x i16>*, i32, <vscale x 4 x i1>, <vscale x 4 x i16>)
 declare <vscale x 8 x i8> @llvm.masked.load.nxv8i8(<vscale x 8 x i8>*, i32, <vscale x 8 x i1>, <vscale x 8 x i8>)
 declare <vscale x 8 x i16> @llvm.masked.load.nxv8i16(<vscale x 8 x i16>*, i32, <vscale x 8 x i1>, <vscale x 8 x i16>)