[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

22

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

22

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

5

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

5

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) {
+      SDValue ExtMaskedLoad =
+          DAG.getMaskedLoad(VT, SDLoc(N), Ld->getChain(), Ld->getBasePtr(),
+                            Ld->getOffset(), Ld->getMask(), Ld->getPassThru(),
+                            ExtVT, Ld->getMemOperand(), Ld->getAddressingMode(),
+                            ISD::LoadExtType::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 void @arm_q15_to_float(i16* nocapture readonly %pSrc, float* nocapture %pDst, i32 %blockSize) {
+; CHECK-LABEL: arm_q15_to_float:
+; CHECK:    vldrw.u32 q1, [r3]
+; CHECK-NEXT:    vdup.32 q2, r2
+; CHECK-NEXT:    vpt.u32 cs, q2, q1
+; CHECK-NEXT:    vldrht.s32 q1, [r0]
+; CHECK-NEXT:    vmov.i32 q0, #0x38000000
+; CHECK-NEXT:    vcvt.f32.s32 q1, q1
+; CHECK-NEXT:    vmul.f32 q0, q1, q0
+entry:
+  %cmp.not5 = icmp eq i32 %blockSize, 0
+  br i1 %cmp.not5, label %return, label %body
+
+body:
+  %trip.count.minus.1 = add i32 %blockSize, -1
+  %next.gep = getelementptr i16, i16* %pSrc, i32 0
+  %next.gep15 = getelementptr float, float* %pDst, i32 0
+  %active.lane.mask = call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 %trip.count.minus.1)
+  %0 = bitcast i16* %next.gep to <4 x i16>*
+  %wide.masked.load = call <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>* %0, i32 2, <4 x i1> %active.lane.mask, <4 x i16> undef)
+  %1 = sitofp <4 x i16> %wide.masked.load to <4 x float>
+  %2 = fmul fast <4 x float> %1, <float 0x3F00000000000000, float 0x3F00000000000000, float 0x3F00000000000000, float 0x3F00000000000000>
+  %3 = bitcast float* %next.gep15 to <4 x float>*
+  call void @llvm.masked.store.v4f32.p0v4f32(<4 x float> %2, <4 x float>* %3, i32 4, <4 x i1> %active.lane.mask)
+  br label %return
+
+return:
+  ret void
+}
+
+declare <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32, i32)
+
+declare <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>*, i32 immarg, <4 x i1>, <4 x i16>)
+
+declare void @llvm.masked.store.v4f32.p0v4f32(<4 x float>, <4 x float>*, i32 immarg, <4 x i1>)