[AARCH64][NEON] Allow to sink operands for aarch64_neon_pmull.

[sunho]

Summary

This teaches AArch64TargetLowering::shouldSinkOperands to sink the operands of aarch64_neon_pmull intrinsic.

[sunho]

Tidy up testcases.

[fhahn]

Thanks for the patch! Can the test be stripped down? Does there need an actual loop or would just 2 different basic blocks be enough? I realise the test of a related patch is also similar, but ideally we would have more compact tests.

[sunho]

Make test case compact.

[sunho]

@fhahn I've shrinked the test case. Could you take a look?

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12188

This falls through to trying to treat the instruction like an indexed pmul, which isn't an instruction available in AArch64. This could do with a test case and possibly being guarded against.

llvm/test/CodeGen/AArch64/neon-vmull-high-p8.ll

2

Can you run update_llc_test_checks on the file. It's useful in cases like this to show all the output

27

I think you can likely remove this too.

[fhahn]

llvm/test/CodeGen/AArch64/neon-vmull-high-p8.ll

9

it would also be good to add a few negative tests, where sinking won't happen, e.g. because there are users in other blocks. Also cases where the other or both shuffles need sinking.

15

it think nowadays this should use poison here

[sunho]

llvm/test/CodeGen/AArch64/neon-vmull-high-p8.ll

9

When would sinking not happen? I've read the code, but it seems it's always sinking unless the operand is phi nod .

[sunho]

@fhahn @dmgreen I've accounted for your comments. Please take a look again!

[sunho]

.

[sunho]

Remove unused line

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12188

IIUC there's still a test case missing to guard against the original issue mentioned by @dmgreen.

llvm/test/CodeGen/AArch64/neon-vmull-high-p8.ll

9

You are right, users in different BBs do not limit sinking, because the instructions are duplicated in that case. One case where sinking won't happen is if one operand has a high mask and one a low mask for example. But it might be better to test for something like that in llvm/test/Transforms/CodeGenPrepare/AArch64/sink-free-instructions.ll

13

You could use -asm-verbose=0 to trim down the output a bit.

21

it's slightly easier to read if there's a newline after the terminators.

[sunho]

Account for comments

[sunho]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12188

I've added a test case that makes sure splat vector is not sunk for pmull in "sink-free-instructions.ll."

llvm/test/CodeGen/AArch64/neon-vmull-high-p8.ll

9

I've added some negative tests in the file you mentioned.

[dmgreen]

Looks like a good patch, from what I can see. Thanks

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12201

You don't need brackets around single statements in llvm. I would personally also invert the condition - if (thing) doathing() seems clearer in intent to me, but I do see that other cases here work the other way arounds already.

[sunho]

Thanks! Can you commit on my behalf?

Author name: sunho
Author email: [e-mail]
Commit message: [AARCH64][NEON] Allow to sink operands for aarch64_neon_pmull.

[dmgreen]

Sure thing. I'll remove the brackets whilst I do.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     switch (II->getIntrinsicID()) {
     case Intrinsic::aarch64_neon_smull:
     case Intrinsic::aarch64_neon_umull:
       if (areExtractShuffleVectors(II->getOperand(0), II->getOperand(1))) {
         Ops.push_back(&II->getOperandUse(0));
         Ops.push_back(&II->getOperandUse(1));
         return true;
       }
       LLVM_FALLTHROUGH;
 
     case Intrinsic::aarch64_neon_sqdmull:
     case Intrinsic::aarch64_neon_sqdmulh:
     case Intrinsic::aarch64_neon_sqrdmulh:
       // Sink splats for index lane variants
       if (isSplatShuffle(II->getOperand(0)))
         Ops.push_back(&II->getOperandUse(0));
       if (isSplatShuffle(II->getOperand(1)))
         Ops.push_back(&II->getOperandUse(1));
       return !Ops.empty();
 
+    case Intrinsic::aarch64_neon_pmull:
+      if (!areExtractShuffleVectors(II->getOperand(0), II->getOperand(1)))
+        return false;
+      Ops.push_back(&II->getOperandUse(0));
+      Ops.push_back(&II->getOperandUse(1));
+      return true;
     case Intrinsic::aarch64_neon_pmull64:
       if (!areOperandsOfVmullHighP64(II->getArgOperand(0),
                                      II->getArgOperand(1)))
         return false;
       Ops.push_back(&II->getArgOperandUse(0));
       Ops.push_back(&II->getArgOperandUse(1));
       return true;
 

llvm/test/CodeGen/AArch64/neon-vmull-high-p8.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -asm-verbose=0 -mtriple=aarch64-none-linux-gnu -mattr=+neon | FileCheck %s
+
+; Check that pmull2 instruction is used for vmull_high_p8 intrinsic
+; even if shufflevector instructions are located in different basic blocks,
+; which can happen when vmull_high_p8 is used inside a loop body.
+;
+
+define <8 x i16> @test_pmull2_sink(<16 x i8> %a, <16 x i8> %b, <8 x i16> %c, i1 %t) {
+; CHECK-LABEL: test_pmull2_sink:
+; CHECK:         tbz w0, #0, .LBB0_2
+; CHECK-NEXT:    pmull2 v0.8h, v0.16b, v1.16b
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB0_2:
+; CHECK-NEXT:    mov v0.16b, v2.16b
+; CHECK-NEXT:    ret
+entry:
+  %0 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  br i1 %t, label %if.then, label %cleanup
+
+if.then:
+  %1 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %res = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %0, <8 x i8> %1)
+  br label %cleanup
+
+cleanup:
+  %retval = phi <8 x i16> [ %res, %if.then ], [ %c, %entry ]
+  ret <8 x i16> %retval
+}
+
+define <8 x i16> @test_pmull2_sink2(<16 x i8> %a, <16 x i8> %b, <8 x i16> %c, i1 %t) {
+; CHECK-LABEL: test_pmull2_sink2:
+; CHECK:         tbz w0, #0, .LBB1_2
+; CHECK-NEXT:    pmull2 v0.8h, v0.16b, v0.16b
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB1_2:
+; CHECK-NEXT:    mov v0.16b, v2.16b
+; CHECK-NEXT:    ret
+entry:
+  %0 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %1 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  br i1 %t, label %if.then, label %cleanup
+
+if.then:
+  %res = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %0, <8 x i8> %1)
+  br label %cleanup
+
+cleanup:
+  %retval = phi <8 x i16> [ %res, %if.then ], [ %c, %entry ]
+  ret <8 x i16> %retval
+}
+
+define <8 x i16> @test_pmull2_sink3(<16 x i8> %a, <16 x i8> %b, <8 x i16> %c, i1 %t, i1 %t2) {
+; CHECK-LABEL: test_pmull2_sink3:
+; CHECK:         tbz w0, #0, .LBB2_2
+; CHECK-NEXT:    tbz w1, #0, .LBB2_3
+; CHECK-NEXT:  .LBB2_2:
+; CHECK-NEXT:    pmull2 v0.8h, v0.16b, v1.16b
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB2_3:
+; CHECK-NEXT:    pmull2 v0.8h, v0.16b, v0.16b
+; CHECK-NEXT:    ret
+entry:
+  %0 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  br i1 %t, label %if.then, label %if.then.2
+
+if.then:
+  %res = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %0, <8 x i8> %0)
+  br i1 %t2, label %if.then.2, label %cleanup
+
+if.then.2:
+  %1 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %res2 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %0, <8 x i8> %1)
+  br label %cleanup
+
+cleanup:
+  %retval = phi <8 x i16> [ %res2, %if.then.2 ], [ %res, %if.then ]
+  ret <8 x i16> %retval
+}
+
+declare <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8>, <8 x i8>)

llvm/test/Transforms/CodeGenPrepare/AArch64/sink-free-instructions.ll

for.cond4.preheader.us.preheader:                 ; preds = %for.cond4.preheader.lr.ph
   %e2 = sext <4 x i16> %s2 to <4 x i32>
   %e1 = sext <4 x i16> %s1 to <4 x i32>
   %0 = add <4 x i32> %e1, %e2
   ret <4 x i32> %0
 
 for.cond4.preheader.preheader:                    ; preds = %for.cond4.preheader.lr.ph
   ret <4 x i32> zeroinitializer
 }
+
+declare <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8>, <8 x i8>)
+
+define <8 x i16> @sink_shufflevector_pmull(<16 x i8> %a, <16 x i8> %b) {
+; CHECK-LABEL: @sink_shufflevector_pmull(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 poison, label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[TMP0:%.*]] = shufflevector <16 x i8> [[A:%.*]], <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[S2:%.*]] = shufflevector <16 x i8> [[B:%.*]], <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[VMULL0:%.*]] = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> [[TMP0]], <8 x i8> [[S2]])
+; CHECK-NEXT:    ret <8 x i16> [[VMULL0]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <16 x i8> [[A]], <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    [[S4:%.*]] = shufflevector <16 x i8> [[B]], <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    [[VMULL1:%.*]] = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> [[TMP1]], <8 x i8> [[S4]])
+; CHECK-NEXT:    ret <8 x i16> [[VMULL1]]
+;
+entry:
+  %s1 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  %s3 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  br i1 poison, label %if.then, label %if.else
+
+if.then:
+  %s2 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  %vmull0 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s1, <8 x i8> %s2)
+  ret <8 x i16> %vmull0
+
+if.else:
+  %s4 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %vmull1 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s3, <8 x i8> %s4)
+  ret <8 x i16> %vmull1
+}
+
+; Indexed pmull is not available on aarch64. Shuffle vector should not be sunk here.
+define <8 x i16> @no_sink_splatvector_pmull(<16 x i8> %a, <16 x i8> %b) {
+; CHECK-LABEL: @no_sink_splatvector_pmull(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[S1:%.*]] = shufflevector <16 x i8> [[A:%.*]], <16 x i8> poison, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    [[S3:%.*]] = shufflevector <16 x i8> [[A]], <16 x i8> poison, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    br i1 poison, label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[S2:%.*]] = shufflevector <16 x i8> [[B:%.*]], <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[VMULL0:%.*]] = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> [[S1]], <8 x i8> [[S2]])
+; CHECK-NEXT:    ret <8 x i16> [[VMULL0]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[S4:%.*]] = shufflevector <16 x i8> [[B]], <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    [[VMULL1:%.*]] = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> [[S3]], <8 x i8> [[S4]])
+; CHECK-NEXT:    ret <8 x i16> [[VMULL1]]
+;
+entry:
+  %s1 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
+  %s3 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
+  br i1 poison, label %if.then, label %if.else
+
+if.then:
+  %s2 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  %vmull0 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s1, <8 x i8> %s2)
+  ret <8 x i16> %vmull0
+
+if.else:
+  %s4 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %vmull1 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s3, <8 x i8> %s4)
+  ret <8 x i16> %vmull1
+}
+
+; Mask used are not suitable for pmull. Shuffle vector should not be sunk here.
+define <8 x i16> @no_sink_shufflevector_pmull(<16 x i8> %a, <16 x i8> %b) {
+; CHECK-LABEL: @no_sink_shufflevector_pmull(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[S1:%.*]] = shufflevector <16 x i8> [[A:%.*]], <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    [[S3:%.*]] = shufflevector <16 x i8> [[A]], <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    br i1 poison, label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[S2:%.*]] = shufflevector <16 x i8> [[B:%.*]], <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[VMULL0:%.*]] = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> [[S1]], <8 x i8> [[S2]])
+; CHECK-NEXT:    ret <8 x i16> [[VMULL0]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[S4:%.*]] = shufflevector <16 x i8> [[B]], <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    [[VMULL1:%.*]] = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> [[S3]], <8 x i8> [[S4]])
+; CHECK-NEXT:    ret <8 x i16> [[VMULL1]]
+;
+entry:
+  %s1 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %s3 = shufflevector <16 x i8> %a, <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  br i1 poison, label %if.then, label %if.else
+
+if.then:
+  %s2 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  %vmull0 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s1, <8 x i8> %s2)
+  ret <8 x i16> %vmull0
+
+if.else:
+  %s4 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %vmull1 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s3, <8 x i8> %s4)
+  ret <8 x i16> %vmull1
+}