[DAGCombiner] Fold sext_inreg of a masked load into a signed extended masked load

[samtebbs]

Summary

This patch adds a DAG combine fold for a sext(masked_load) into a sign extended masked load.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11135

I think you need to check for hasOneUse. And maybe with more of the checks from the above fold (sext_inreg (zextload x)) case too, like checking types match and maybe for legal masked loads.

llvm/test/CodeGen/Thumb2/sext-masked-load.ll

22 ↗

(On Diff #279839)

Can you simplify this down to just a llvm.masked.load.v4i16.p0v4i16 and a sitofp ?
Ideally they would look something like the tests in llvm/test/CodeGen/Thumb2/mve-masked-load.ll.

If so can you also add some other tests for different types and for uitofp.

[samtebbs]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11135

Thanks, I do need to add those. But for checking if masked loads are legal, is that necessary if there is already a masked load in the DAG?

llvm/test/CodeGen/Thumb2/sext-masked-load.ll

22 ↗

(On Diff #279839)

Sure! I'll give that a shot.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11135

OK. tryToFoldExtOfMaskedLoad and the load fold above both have a isLoadExtLegal check. That does seem to make some assumptions, but they seem to be valid so far.

It would be a strange architecture that included zext/anyext masked loads, but didn't include sext masked loads. But even if just for consistency it sounds like it would be worth included that isLoadExtLegal check too.

[samtebbs]

Simplify test and add extra checks to sext masked load generation.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11136

This needs a bit of a clang-format.

llvm/test/CodeGen/Thumb2/sext-masked-load.ll

4 ↗

(On Diff #280479)

Can you add tests for these legal types:
8xi8 -> half
4xi8 -> float
4xi16 -> float

I think it's worth having some illegal types too, Maybe something like a 4xi32 -> double?

And can you then replicate them for uitofp too. If I'm right that might need some fixup too, but it may be better to handle that in a different patch.

Oh, and do mind moving the test to llvm/test/CodeGen/Thumb2/mve-sext-masked-load.ll to keep it with the others MVE test.

[samtebbs]

Add more tests and clean up checks in visitSIGN_EXTEND_INREG

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11136

+1

llvm/test/CodeGen/Thumb2/sext-masked-load.ll

4 ↗

(On Diff #280479)

I've added those tests, thanks for the suggestions (I already had a 4xi16 -> float test, if I understand correctly).

I'd prefer to handle the uitofp cases in a separate patch too.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11136

Hmm. One last thing... If the extension type is a SEXTLOAD, should we be removing the sign_ext_inreg anyway?
I guess with NON_EXTLOAD the ExtVT and MemoryVT would not match?

llvm/test/CodeGen/Thumb2/mve-sext-masked-load.ll

3

You likely don't need -tail-predication=enabled any more.

[samtebbs]

Remove -tail-predication option from test and change SEXTLOAD to NON_EXTLOAD in DAGCombiner check.

[dmgreen]

Thanks. LGTM, but please give others a day in case they have comments.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue ExtLoad = DAG.getExtLoad(ISD::SEXTLOAD, SDLoc(N), VT,
                                      LN0->getChain(),
                                      LN0->getBasePtr(), ExtVT,
                                      LN0->getMemOperand());
     CombineTo(N, ExtLoad);
     CombineTo(N0.getNode(), ExtLoad, ExtLoad.getValue(1));
     return SDValue(N, 0);   // Return N so it doesn't get rechecked!
   }
 
+  // fold (sext_inreg (masked_load x)) -> (sext_masked_load x)
+  // ignore it if the masked load is already sign extended
+  if (MaskedLoadSDNode *Ld = dyn_cast<MaskedLoadSDNode>(N0)) {
+    if (ExtVT == Ld->getMemoryVT() && N0.hasOneUse() &&
+        Ld->getExtensionType() != ISD::LoadExtType::NON_EXTLOAD &&
+        TLI.isLoadExtLegal(ISD::SEXTLOAD, VT, ExtVT)) {
+      SDValue ExtMaskedLoad = DAG.getMaskedLoad(
+          VT, SDLoc(N), Ld->getChain(), Ld->getBasePtr(), Ld->getOffset(),
+          Ld->getMask(), Ld->getPassThru(), ExtVT, Ld->getMemOperand(),
+          Ld->getAddressingMode(), ISD::SEXTLOAD, Ld->isExpandingLoad());
+      CombineTo(N, ExtMaskedLoad);
+      CombineTo(N0.getNode(), ExtMaskedLoad, ExtMaskedLoad.getValue(1));
+      return SDValue(N, 0); // Return N so it doesn't get rechecked!
+    }
+  }
+
   // Form (sext_inreg (bswap >> 16)) or (sext_inreg (rotl (bswap) 16))
   if (ExtVTBits <= 16 && N0.getOpcode() == ISD::OR) {
     if (SDValue BSwap = MatchBSwapHWordLow(N0.getNode(), N0.getOperand(0),
                                            N0.getOperand(1), false))
       return DAG.getNode(ISD::SIGN_EXTEND_INREG, SDLoc(N), VT, BSwap, N1);
   }
 
   return SDValue();

llvm/test/CodeGen/Thumb2/mve-sext-masked-load.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+mve.fp,+fp64 -verify-machineinstrs -o - %s | FileCheck %s
+
+define arm_aapcs_vfpcc <4 x float> @foo_v4i16(<4 x i16>* nocapture readonly %pSrc, i32 %blockSize, <4 x i16> %a) {
+; CHECK-LABEL: foo_v4i16:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmovlb.s16 q0, q0
+; CHECK-NEXT:    vpt.s32 lt, q0, zr
+; CHECK-NEXT:    vldrht.s32 q0, [r0]
+; CHECK-NEXT:    vcvt.f32.s32 q0, q0
+; CHECK-NEXT:    bx lr
+entry:
+  %active.lane.mask = icmp slt <4 x i16> %a, zeroinitializer
+  %wide.masked.load = call <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>* %pSrc, i32 2, <4 x i1> %active.lane.mask, <4 x i16> undef)
+  %0 = sitofp <4 x i16> %wide.masked.load to <4 x float>
+  ret <4 x float> %0
+}
+
+define arm_aapcs_vfpcc <8 x half> @foo_v8i8(<8 x i8>* nocapture readonly %pSrc, i32 %blockSize, <8 x i8> %a) {
+; CHECK-LABEL: foo_v8i8:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmovlb.s8 q0, q0
+; CHECK-NEXT:    vpt.s16 lt, q0, zr
+; CHECK-NEXT:    vldrbt.s16 q0, [r0]
+; CHECK-NEXT:    vcvt.f16.s16 q0, q0
+; CHECK-NEXT:    bx lr
+entry:
+  %active.lane.mask = icmp slt <8 x i8> %a, zeroinitializer
+  %wide.masked.load = call <8 x i8> @llvm.masked.load.v8i8.p0v8i8(<8 x i8>* %pSrc, i32 1, <8 x i1> %active.lane.mask, <8 x i8> undef)
+  %0 = sitofp <8 x i8> %wide.masked.load to <8 x half>
+  ret <8 x half> %0
+}
+
+define arm_aapcs_vfpcc <4 x float> @foo_v4i8(<4 x i8>* nocapture readonly %pSrc, i32 %blockSize, <4 x i8> %a) {
+; CHECK-LABEL: foo_v4i8:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmovlb.s8 q0, q0
+; CHECK-NEXT:    vmovlb.s16 q0, q0
+; CHECK-NEXT:    vpt.s32 lt, q0, zr
+; CHECK-NEXT:    vldrbt.s32 q0, [r0]
+; CHECK-NEXT:    vcvt.f32.s32 q0, q0
+; CHECK-NEXT:    bx lr
+entry:
+  %active.lane.mask = icmp slt <4 x i8> %a, zeroinitializer
+  %wide.masked.load = call <4 x i8> @llvm.masked.load.v4i8.p0v4i8(<4 x i8>* %pSrc, i32 1, <4 x i1> %active.lane.mask, <4 x i8> undef)
+  %0 = sitofp <4 x i8> %wide.masked.load to <4 x float>
+  ret <4 x float> %0
+}
+
+define arm_aapcs_vfpcc <4 x double> @foo_v4i32(<4 x i32>* nocapture readonly %pSrc, i32 %blockSize, <4 x i32> %a) {
+; CHECK-LABEL: foo_v4i32:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    .save {r7, lr}
+; CHECK-NEXT:    push {r7, lr}
+; CHECK-NEXT:    .vsave {d8, d9, d10, d11, d12, d13}
+; CHECK-NEXT:    vpush {d8, d9, d10, d11, d12, d13}
+; CHECK-NEXT:    vpt.s32 lt, q0, zr
+; CHECK-NEXT:    vldrwt.u32 q5, [r0]
+; CHECK-NEXT:    vmov.f64 d8, d10
+; CHECK-NEXT:    vmov.f32 s18, s21
+; CHECK-NEXT:    vmov r0, s18
+; CHECK-NEXT:    asrs r1, r0, #31
+; CHECK-NEXT:    bl __aeabi_l2d
+; CHECK-NEXT:    vmov r2, s16
+; CHECK-NEXT:    vmov d9, r0, r1
+; CHECK-NEXT:    asrs r3, r2, #31
+; CHECK-NEXT:    mov r0, r2
+; CHECK-NEXT:    mov r1, r3
+; CHECK-NEXT:    bl __aeabi_l2d
+; CHECK-NEXT:    vmov.f64 d12, d11
+; CHECK-NEXT:    vmov.f32 s26, s23
+; CHECK-NEXT:    vmov d8, r0, r1
+; CHECK-NEXT:    vmov r2, s26
+; CHECK-NEXT:    asrs r3, r2, #31
+; CHECK-NEXT:    mov r0, r2
+; CHECK-NEXT:    mov r1, r3
+; CHECK-NEXT:    bl __aeabi_l2d
+; CHECK-NEXT:    vmov r2, s24
+; CHECK-NEXT:    vmov d11, r0, r1
+; CHECK-NEXT:    asrs r3, r2, #31
+; CHECK-NEXT:    mov r0, r2
+; CHECK-NEXT:    mov r1, r3
+; CHECK-NEXT:    bl __aeabi_l2d
+; CHECK-NEXT:    vmov d10, r0, r1
+; CHECK-NEXT:    vmov q0, q4
+; CHECK-NEXT:    vmov q1, q5
+; CHECK-NEXT:    vpop {d8, d9, d10, d11, d12, d13}
+; CHECK-NEXT:    pop {r7, pc}
+entry:
+  %active.lane.mask = icmp slt <4 x i32> %a, zeroinitializer
+  %wide.masked.load = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %pSrc, i32 4, <4 x i1> %active.lane.mask, <4 x i32> undef)
+  %0 = sitofp <4 x i32> %wide.masked.load to <4 x double>
+  ret <4 x double> %0
+}
+
+declare <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>*, i32 immarg, <4 x i1>, <4 x i16>)
+
+declare <8 x i8> @llvm.masked.load.v8i8.p0v8i8(<8 x i8>*, i32 immarg, <8 x i1>, <8 x i8>)
+
+declare <4 x i8> @llvm.masked.load.v4i8.p0v4i8(<4 x i8>*, i32 immarg, <4 x i1>, <4 x i8>)
+
+declare <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>*, i32 immarg, <4 x i1>, <4 x i32>)