[AArch64] Custom Lower MULLH{S,U} for v16i8, v8i16, and v4i32

[zatrazz]

Summary

This patch adds a custom lowering for ISD::MULH{S,U} used on divide by
constant optimization (DAGCombiner::BuildSDIV and DAGCombiner::BuildUDIV).

New patterns for smull and umull are added, so AArch64ISD::{S,U}MULL
can be correctly lowered to smull2 and umull2.

[SjoerdMeijer]

lib/Target/AArch64/AArch64ISelLowering.cpp

2557

Do you need to pass the 'sign' boolean? Can you not look at the opcode and check for ISD::MULHS or ISD::MULHU?

2728

If you check the opcode inside LowerMULH, this can be simplified a bit, both these cases can fallthrough and call LowerMULH.

test/CodeGen/AArch64/arm64-neon-mul-div-cte.ll

4

nit: this is mul16xi8

15

nit: this is mul8xi16

[zatrazz]

lib/Target/AArch64/AArch64ISelLowering.cpp

2557

Yes, you can divise the correct ISD from SDValue opcode, I will change it.

test/CodeGen/AArch64/arm64-neon-mul-div-cte.ll

4

Ack.

15

Ack.

[zatrazz]

Update patch from previous comments.

[SjoerdMeijer]

Sorry, I needed to get up to speed here, also with this hacker's delight trick.
I think it would be good if you add some comments to LowerMULH what we are trying to do here.

I think I have some more comments, mainly about the tests:

• you're not checking the magic number (e.g. in the first test "movi v1.16b, #57"), but I think that is quite relevant here.
• you're only using constant vectors with value 9. I think it would be good to vary here.
• I think you should also define the operands of the smull2 and smull instructions with regexps, and use them.
• nit: some functions still have inconsistent names, e.g. function "umul8xi16" tests "udiv <16 x i8>"

[zatrazz]

In D46009#1083782, @SjoerdMeijer wrote:

Sorry, I needed to get up to speed here, also with this hacker's delight trick.
I think it would be good if you add some comments to LowerMULH what we are trying to do here.

Ack.

I think I have some more comments, mainly about the tests:

• you're not checking the magic number (e.g. in the first test "movi v1.16b, #57"), but I think that is quite relevant here.

Right, but the change itself is to correct lower the multiply high part of the divide by constant optimization. I have adjusted the next patch version to match the expected constant as well.

• you're only using constant vectors with value 9. I think it would be good to vary here.

Ack.

• I think you should also define the operands of the smull2 and smull instructions with regexps, and use them.

Ack.

• nit: some functions still have inconsistent names, e.g. function "umul8xi16" tests "udiv <16 x i8>"

Ack, I have changed to div*.

[zatrazz]

Update patch from previous comments.

[SjoerdMeijer]

Thanks. This looks OK to me.

lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasNEON()) {
 
     setOperationAction(ISD::ANY_EXTEND, MVT::v4i32, Legal);
     setTruncStoreAction(MVT::v2i32, MVT::v2i16, Expand);
     // Likewise, narrowing and extending vector loads/stores aren't handled
     // directly.
     for (MVT VT : MVT::vector_valuetypes()) {
       setOperationAction(ISD::SIGN_EXTEND_INREG, VT, Expand);
 
+      if (VT == MVT::v16i8 || VT == MVT::v8i16 || VT == MVT::v4i32) {
+        setOperationAction(ISD::MULHS, VT, Custom);
+        setOperationAction(ISD::MULHU, VT, Custom);
+      } else {
         setOperationAction(ISD::MULHS, VT, Expand);
-      setOperationAction(ISD::SMUL_LOHI, VT, Expand);
         setOperationAction(ISD::MULHU, VT, Expand);
+      }
+      setOperationAction(ISD::SMUL_LOHI, VT, Expand);
       setOperationAction(ISD::UMUL_LOHI, VT, Expand);
 
       setOperationAction(ISD::BSWAP, VT, Expand);
 
       for (MVT InnerVT : MVT::vector_valuetypes()) {
         setTruncStoreAction(VT, InnerVT, Expand);
         setLoadExtAction(ISD::SEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::ZEXTLOAD, VT, InnerVT, Expand);
static SDValue LowerMUL(SDValue Op, SelectionDAG &DAG) {
   EVT Op1VT = Op1.getValueType();
   return DAG.getNode(N0->getOpcode(), DL, VT,
                      DAG.getNode(NewOpc, DL, VT,
                                DAG.getNode(ISD::BITCAST, DL, Op1VT, N00), Op1),
                      DAG.getNode(NewOpc, DL, VT,
                                DAG.getNode(ISD::BITCAST, DL, Op1VT, N01), Op1));
 }
 
+static SDValue LowerMULH(SDValue Op, SelectionDAG &DAG) {
+  EVT VT = Op.getValueType();
+  assert(VT.is128BitVector() && VT.isInteger() &&
+         "unexpected type for custom-lowering ISD::MULH{U,S}");
+
+  SDValue V0 = Op.getOperand(0);
+  SDValue V1 = Op.getOperand(1);
+
+  SDLoc DL(Op);
+
+  EVT ExtractVT = VT.getHalfNumVectorElementsVT(*DAG.getContext());
+
+  unsigned Mull2VectorIdx = VT.getVectorNumElements () / 2;
+  SDValue VMullIdx = DAG.getConstant(0, DL, MVT::i64);
+  SDValue VMull2Idx = DAG.getConstant(Mull2VectorIdx, DL, MVT::i64);
+
+  SDValue VMullV0 =
+    DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ExtractVT, V0, VMullIdx);
+  SDValue VMullV1 =
+    DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ExtractVT, V1, VMullIdx);
+
+  SDValue VMull2V0 =
+    DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ExtractVT, V0, VMull2Idx);
+  SDValue VMull2V1 =
+    DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ExtractVT, V1, VMull2Idx);
+
+  unsigned MullOpc = Op.getOpcode() == ISD::MULHS ? AArch64ISD::SMULL
+                                                  : AArch64ISD::UMULL;
+
+  EVT MullVT = ExtractVT.widenIntegerVectorElementType(*DAG.getContext());
+  SDValue Mull  = DAG.getNode(MullOpc, DL, MullVT, VMullV0, VMullV1);
+  SDValue Mull2 = DAG.getNode(MullOpc, DL, MullVT, VMull2V0, VMull2V1);
+
+  Mull  = DAG.getNode(ISD::BITCAST, DL, VT, Mull);
+  Mull2 = DAG.getNode(ISD::BITCAST, DL, VT, Mull2);
+
+  return DAG.getNode(AArch64ISD::UZP2, DL, VT, Mull, Mull2);
+}
+
 SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
                                                      SelectionDAG &DAG) const {
   unsigned IntNo = cast<ConstantSDNode>(Op.getOperand(0))->getZExtValue();
   SDLoc dl(Op);
   switch (IntNo) {
   default: return SDValue();    // Don't custom lower most intrinsics.
   case Intrinsic::thread_pointer: {
     EVT PtrVT = getPointerTy(DAG.getDataLayout());
case ISD::UINT_TO_FP:
     return LowerINT_TO_FP(Op, DAG);
   case ISD::FP_TO_SINT:
   case ISD::FP_TO_UINT:
     return LowerFP_TO_INT(Op, DAG);
   case ISD::FSINCOS:
     return LowerFSINCOS(Op, DAG);
   case ISD::MUL:
     return LowerMUL(Op, DAG);
+  case ISD::MULHS:
+  case ISD::MULHU:
+    return LowerMULH(Op, DAG);
   case ISD::INTRINSIC_WO_CHAIN:
     return LowerINTRINSIC_WO_CHAIN(Op, DAG);
   case ISD::VECREDUCE_ADD:
   case ISD::VECREDUCE_SMAX:
   case ISD::VECREDUCE_SMIN:
   case ISD::VECREDUCE_UMAX:
   case ISD::VECREDUCE_UMIN:
   case ISD::VECREDUCE_FMAX:

lib/Target/AArch64/AArch64InstrInfo.td

def : Pat<(v2i64 (opnode (v2i32 V64:$Rn), (v2i32 V64:$Rm))),
             (INST2S V64:$Rn, V64:$Rm)>;
 }
 
 defm : Neon_mul_widen_patterns<AArch64smull, SMULLv8i8_v8i16,
   SMULLv4i16_v4i32, SMULLv2i32_v2i64>;
 defm : Neon_mul_widen_patterns<AArch64umull, UMULLv8i8_v8i16,
   UMULLv4i16_v4i32, UMULLv2i32_v2i64>;
 
+// Patterns for smull2/umull2.
+multiclass Neon_mul_high_patterns<SDPatternOperator opnode,
+  Instruction INST8B, Instruction INST4H, Instruction INST2S> {
+  def : Pat<(v8i16 (opnode (extract_high_v16i8 V128:$Rn),
+                           (extract_high_v16i8 V128:$Rm))),
+             (INST8B V128:$Rn, V128:$Rm)>;
+  def : Pat<(v4i32 (opnode (extract_high_v8i16 V128:$Rn),
+                           (extract_high_v8i16 V128:$Rm))),
+             (INST4H V128:$Rn, V128:$Rm)>;
+  def : Pat<(v2i64 (opnode (extract_high_v4i32 V128:$Rn),
+                           (extract_high_v4i32 V128:$Rm))),
+             (INST2S V128:$Rn, V128:$Rm)>;
+}
+
+defm : Neon_mul_high_patterns<AArch64smull, SMULLv16i8_v8i16,
+  SMULLv8i16_v4i32, SMULLv4i32_v2i64>;
+defm : Neon_mul_high_patterns<AArch64umull, UMULLv16i8_v8i16,
+  UMULLv8i16_v4i32, UMULLv4i32_v2i64>;
+
 // Additional patterns for SMLAL/SMLSL and UMLAL/UMLSL
 multiclass Neon_mulacc_widen_patterns<SDPatternOperator opnode,
   Instruction INST8B, Instruction INST4H, Instruction INST2S> {
   def : Pat<(v8i16 (opnode (v8i16 V128:$Rd), (v8i8 V64:$Rn), (v8i8 V64:$Rm))),
             (INST8B V128:$Rd, V64:$Rn, V64:$Rm)>;
   def : Pat<(v4i32 (opnode (v4i32 V128:$Rd), (v4i16 V64:$Rn), (v4i16 V64:$Rm))),
             (INST4H V128:$Rd, V64:$Rn, V64:$Rm)>;
   def : Pat<(v2i64 (opnode (v2i64 V128:$Rd), (v2i32 V64:$Rn), (v2i32 V64:$Rm))),

test/CodeGen/AArch64/arm64-neon-mul-div-cte.ll

+; RUN: llc < %s -verify-machineinstrs -mtriple=arm64-none-linux-gnu -mattr=+neon | FileCheck %s
+
+define <16 x i8> @mul16xi8(<16 x i8> %x) {
+; CHECK-LABEL: mul16xi8:
+; CHECK: smull2 [[SMULL2:(v[0-9]+)]].8h, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
+; CHECK: smull  [[SMULL:(v[0-9]+)]].8h, {{v[0-9]+}}.8b, {{v[0-9]+}}.8b
+; CHECK: uzp2   [[UZP2:(v[0-9]+).16b]], [[SMULL]].16b, [[SMULL2]].16b
+; CHECK: sshr   [[SSHR:(v[0-9]+.16b)]], [[UZP2]], #1
+; CHECK: usra   v0.16b, [[SSHR]], #7
+  %div = sdiv <16 x i8> %x, <i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9>
+  ret <16 x i8> %div
+}
+
+define <8 x i16> @mul8xi16(<8 x i16> %x) {
+; CHECK-LABEL: mul8xi16:
+; CHECK: smull2 [[SMULL2:(v[0-9]+)]].4s, {{v[0-9]+}}.8h, {{v[0-9]+}}.8h
+; CHECK: smull  [[SMULL:(v[0-9]+)]].4s, {{v[0-9]+}}.4h, {{v[0-9]+}}.4h
+; CHECK: uzp2   [[UZP2:(v[0-9]+).8h]], [[SMULL]].8h, [[SMULL2]].8h
+; CHECK: usra   v0.8h, [[UZP2]], #15
+  %div = sdiv <8 x i16> %x, <i16 9, i16 9, i16 9, i16 9, i16 9, i16 9, i16 9, i16 9>
+  ret <8 x i16> %div
+}
+
+define <4 x i32> @mul32xi4(<4 x i32> %x) {
+; CHECK-LABEL: mul32xi4:
+; CHECK: smull2 [[SMULL2:(v[0-9]+)]].2d, {{v[0-9]+}}.4s, {{v[0-9]+}}.4s
+; CHECK: smull  [[SMULL:(v[0-9]+)]].2d, {{v[0-9]+}}.2s, {{v[0-9]+}}.2s
+; CHECK: uzp2   [[UZP2:(v[0-9]+).4s]], [[SMULL]].4s, [[SMULL2]].4s
+; CHECK: sshr   [[SSHR:(v[0-9]+.4s)]], [[UZP2]], #1
+; CHECK: usra   v0.4s, [[UZP2]], #31
+  %div = sdiv <4 x i32> %x, <i32 9, i32 9, i32 9, i32 9>
+  ret <4 x i32> %div
+}
+
+define <16 x i8> @umul8xi16(<16 x i8> %x) {
+; CHECK-LABEL: umul8xi16:
+; CHECK: umull2 [[SMULL2:(v[0-9]+)]].8h, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
+; CHECK: umull  [[SMULL:(v[0-9]+)]].8h, {{v[0-9]+}}.8b, {{v[0-9]+}}.8b
+; CHECK: uzp2   [[UZP2:(v[0-9]+).16b]], [[SMULL]].16b, [[SMULL2]].16b
+; CHECK: ushr   v0.16b, [[UZP2]], #1
+  %div = udiv <16 x i8> %x, <i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9, i8 9>
+  ret <16 x i8> %div
+}
+
+define <8 x i16> @umul16xi8(<8 x i16> %x) {
+; CHECK-LABEL: umul16xi8:
+; CHECK: umull2 [[SMULL2:(v[0-9]+)]].4s, {{v[0-9]+}}.8h, {{v[0-9]+}}.8h
+; CHECK: umull  [[SMULL:(v[0-9]+)]].4s, {{v[0-9]+}}.4h, {{v[0-9]+}}.4h
+; CHECK: uzp2   [[UZP2:(v[0-9]+).8h]], [[SMULL]].8h, [[SMULL2]].8h
+; CHECK: ushr   v0.8h, [[UZP2]], #3
+  %div = udiv <8 x i16> %x, <i16 9, i16 9, i16 9, i16 9, i16 9, i16 9, i16 9, i16 9>
+  ret <8 x i16> %div
+}
+
+define <4 x i32> @umul32xi4(<4 x i32> %x) {
+; CHECK-LABEL: umul32xi4:
+; CHECK: umull2 [[SMULL2:(v[0-9]+)]].2d, {{v[0-9]+}}.4s, {{v[0-9]+}}.4s
+; CHECK: umull  [[SMULL:(v[0-9]+)]].2d, {{v[0-9]+}}.2s, {{v[0-9]+}}.2s
+; CHECK: uzp2   [[UZP2:(v[0-9]+).4s]], [[SMULL]].4s, [[SMULL2]].4s
+; CHECK: ushr   v0.4s, [[UZP2]], #1
+  %div = udiv <4 x i32> %x, <i32 9, i32 9, i32 9, i32 9>
+  ret <4 x i32> %div
+}

test/CodeGen/AArch64/neon-idiv.ll

 ; RUN: llc -mtriple=aarch64-none-linux-gnu < %s -mattr=+neon | FileCheck %s
 
 define <4 x i32> @test1(<4 x i32> %a) {
   %rem = srem <4 x i32> %a, <i32 7, i32 7, i32 7, i32 7>
   ret <4 x i32> %rem
-; CHECK-LABEL: test1
-; FIXME: Can we lower this more efficiently?
-; CHECK: mul
-; CHECK: mul
-; CHECK: mul
-; CHECK: mul
+; For C constant X/C is simplified to X-X/C*C. The X/C division is lowered
+; to MULHS due the simplification by multiplying by a magic number
+; (TargetLowering::BuildSDIV).
+; CHECK-LABEL: test1:
+; CHECK: smull2 [[SMULL2:(v[0-9]+)]].2d, {{v[0-9]+}}.4s, {{v[0-9]+}}.4s
+; CHECK: smull  [[SMULL:(v[0-9]+)]].2d, {{v[0-9]+}}.2s, {{v[0-9]+}}.2s
+; CHECK: uzp2   [[UZP2:(v[0-9]+).4s]], [[SMULL]].4s, [[SMULL2]].4s
+; CHECK: add    [[ADD:(v[0-9]+.4s)]], [[UZP2]], v0.4s
+; CHECK: sshr   [[SSHR:(v[0-9]+.4s)]], [[ADD]], #2
 }
-