[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

11131

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

23

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

11131

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

23

Sure! I'll give that a shot.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11131

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

11132

This needs a bit of a clang-format.

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

4

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

11132

+1

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

4

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

11132

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

2 ↗

(On Diff #280917)

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() && !LegalOperations &&
+        Ld->getExtensionType() != ISD::LoadExtType::SEXTLOAD && N0.hasOneUse() && 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-masked-load.ll

 define arm_aapcs_vfpcc <4 x i32> @multi_user_sext(<4 x i16> *%dest, <4 x i32> %a) {
 ; CHECK-LE-LABEL: multi_user_sext:
 ; CHECK-LE:       @ %bb.0: @ %entry
 ; CHECK-LE-NEXT:    .save {r7, lr}
 ; CHECK-LE-NEXT:    push {r7, lr}
 ; CHECK-LE-NEXT:    .vsave {d8, d9}
 ; CHECK-LE-NEXT:    vpush {d8, d9}
 ; CHECK-LE-NEXT:    vpt.s32 gt, q0, zr
-; CHECK-LE-NEXT:    vldrht.u32 q4, [r0]
+; CHECK-LE-NEXT:    vldrht.s32 q4, [r0]
 ; CHECK-LE-NEXT:    vmov r0, r1, d8
 ; CHECK-LE-NEXT:    vmov r2, r3, d9
 ; CHECK-LE-NEXT:    bl foo
-; CHECK-LE-NEXT:    vmovlb.s16 q0, q4
+; CHECK-LE-NEXT:    vmov q0, q4
 ; CHECK-LE-NEXT:    vpop {d8, d9}
 ; CHECK-LE-NEXT:    pop {r7, pc}
 ;
 ; CHECK-BE-LABEL: multi_user_sext:
 ; CHECK-BE:       @ %bb.0: @ %entry
 ; CHECK-BE-NEXT:    .save {r7, lr}
 ; CHECK-BE-NEXT:    push {r7, lr}
 ; CHECK-BE-NEXT:    .vsave {d8, d9}
 ; CHECK-BE-NEXT:    vpush {d8, d9}
 ; CHECK-BE-NEXT:    vrev64.32 q1, q0
 ; CHECK-BE-NEXT:    vpt.s32 gt, q1, zr
-; CHECK-BE-NEXT:    vldrht.u32 q4, [r0]
-; CHECK-BE-NEXT:    vrev64.32 q0, q4
-; CHECK-BE-NEXT:    vmov r1, r0, d0
-; CHECK-BE-NEXT:    vmov r3, r2, d1
+; CHECK-BE-NEXT:    vldrht.s32 q0, [r0]
+; CHECK-BE-NEXT:    vrev64.32 q4, q0
+; CHECK-BE-NEXT:    vmov r1, r0, d8
+; CHECK-BE-NEXT:    vmov r3, r2, d9
 ; CHECK-BE-NEXT:    bl foo
-; CHECK-BE-NEXT:    vmovlb.s16 q1, q4
-; CHECK-BE-NEXT:    vrev64.32 q0, q1
+; CHECK-BE-NEXT:    vmov q0, q4
 ; CHECK-BE-NEXT:    vpop {d8, d9}
 ; CHECK-BE-NEXT:    pop {r7, pc}
 entry:
   %c = icmp sgt <4 x i32> %a, zeroinitializer
   %l = call <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>* %dest, i32 2, <4 x i1> %c, <4 x i16> zeroinitializer)
   call void @foo(<4 x i16> %l)
   %ext = sext <4 x i16> %l to <4 x i32>
   ret <4 x i32> %ext

llvm/test/CodeGen/Thumb2/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 -verify-machineinstrs -tail-predication=enabled -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
+}
+
+declare <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>*, i32 immarg, <4 x i1>, <4 x i16>)