[SelectionDAG] Add expansion and promotion of [US]MUL_LOHI

[nhaehnle]

Summary

Most targets set the action for these nodes to Expand even though there
isn't actually any code for them in ExpandNode. Instead, targets simply
relied on the fact that no code generates these nodes as long as the
nodes aren't legal or custom.

However, generating these nodes can be useful e.g. for divide-by-constant
in wider integer types. Targets will want to deal with [US]MUL_LOHI
differently depending on the available instructions, e.g. targets with a
native 64-bit multiply will want to Promote the 32-bit [US]MUL_LOHI
instructions, and targets with native MULH[US] instructions will want to
Split.

This patch intends to not change the generated code, but indirect effects
are possible since expansions/promotions that were previously done in
DAGCombine may now be done in LegalizeDAG.

See D24822 for a change that actually uses the new expansion.

[nhaehnle]

Ping.

[efriedma]

include/llvm/Target/TargetLowering.h

147

It doesn't look like None is actually used anywhere... did you mean to use it somehow?

lib/CodeGen/SelectionDAG/TargetLowering.cpp

3258

Why are you using ADDC+ADDE rather than just plain ADD? It seems strange to split the operation, given that it might be legal. (I might be missing something here, though.)

lib/Target/Sparc/SparcISelLowering.cpp

1696 ↗

(On Diff #72620)

This FIXME doesn't make sense here; Promote is just the right thing to do here.

lib/Target/X86/X86ISelLowering.cpp

32532 ↗

(On Diff #72620)

I'm confused; you're returning "Scalarize" for scalar types?

[nhaehnle]

Thank you for taking a look. This patch should address all your comments,
plus I'm adding a guard against incorrectly using the known bits
optimizations with vector types.

[nhaehnle]

include/llvm/Target/TargetLowering.h

147

I used it at some point to preserve the original behaviour. I'll remove it.

lib/CodeGen/SelectionDAG/TargetLowering.cpp

3258

That's just how the code happened to evolve because AMDGPU (which is where this all started) doesn't have 64-bit additions, but it's a good point. I'm changing it to use additions in the larger type.

lib/Target/Sparc/SparcISelLowering.cpp

1696 ↗

(On Diff #72620)

FIXME removed.

lib/Target/X86/X86ISelLowering.cpp

32532 ↗

(On Diff #72620)

On X86 this only gets called for vector types. I'm adding an assertion to make that explicit.

[jyknight]

I don't really understand why this needs to be a new target hook. Can't it have similar code as the MUL expansion code at LegalizeDAG.cpp:3316 (in the base version): if MULH and MUL are LegalOrCustom, then you'd split the MUL_LOHI into those two operations. Otherwise, you'd do half-width multiplies?

Also, I've fixed the sparc backend to properly implement [US]MUL_LOHI in r283381, removing the FIXME here.

[nhaehnle]

In D24956#562725, @jyknight wrote:

I don't really understand why this needs to be a new target hook. Can't it have similar code as the MUL expansion code at LegalizeDAG.cpp:3316 (in the base version): if MULH and MUL are LegalOrCustom, then you'd split the MUL_LOHI into those two operations. Otherwise, you'd do half-width multiplies?

IMHO, explicit is better than implicit. Also, how would you choose between scalarizing and half-width multiplies for vector types? You might be able to come up with a magic heuristic that works now, only to be broken when a new backend comes along.

Also, I've fixed the sparc backend to properly implement [US]MUL_LOHI in r283381, removing the FIXME here.

Thanks!

[jyknight]

IMHO, explicit is better than implicit. Also, how would you choose between scalarizing and half-width multiplies for vector types? You might be able to come up with a magic heuristic that works now, only to be broken when a new backend comes along.

I don't think it makes much sense to require backends to decide between Split and Half; whenever you can split into MULH and MUL, you should, and do 4 half-width multiplies otherwise.

I don't know much about the vector question -- how is that done for other operations? Surely the question about which axis to "expand" when expanding a vector operation isn't only an question for multiplication, so a MUL-specific solution to this issue feels strange -- I'd expect that either there's already a solution in use elsewhere, or a wider problem.

[RKSimon]

lib/CodeGen/SelectionDAG/TargetLowering.cpp

3112

If this is going to start being used by vectors, shouldn't these getSizeInBits() calls be replaced with getScalarSizeInBits()?

[efriedma]

Legalization for certain vector operations uses "PROMOTE" to indicate that the operation should be performed in a different vector type. That really stretches the meaning if we're using it to mean "split into multiple vector multiplies", but I guess it could work. See VectorLegalizer::Promote in LegalizeVectorOps.cpp.

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

3341

Sorry, I really haven't worked with this code in a while; I should have spotted the issue here earlier. :(

We can't ever scalarize an operation in LegalizeDAG: scalarizing can produce illegal types, and LegalizeDAG can't handle them. We have to handle this a bit earlier, in LegalizeVectorOps.

[nhaehnle]

Thank you all for looking into this and sorry for the delay. I think I've
addressed all your comments:

• getMulExpansion is gone
• vector scalarization is done in LegalizeVectorOps
• use getScalarSizeInBits() throughout to prepare for vector use

There's also the question of "promoting" / using half-width multiplies for
vectors. That could become useful to have more aggressive replacement of
divide-by-constant with the equivalent multiply+shift sequence also for
vector types, i.e. pushing D24822 even further.

However, there are more problems there (X86 gets unnecessarily worse looking
instruction sequences in some tests), so that should be a separate change.

[nhaehnle]

Simplify the Expand part in LegalizeDAG, since vectors are now expanded
earlier.

[nhaehnle]

Ping.

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

3341

Okay, fixing this.

lib/CodeGen/SelectionDAG/TargetLowering.cpp

3112

Fixed, thanks.

[nhaehnle]

Ping^2

[efriedma]

Sorry about the delay...

Have you seen https://reviews.llvm.org/D26628 ? Does that interact with the patch in any way?

TargetLowering::expandMUL_LOHI is kind of difficult to review with all the different unrelated changes involved; would it be possible to separate out the non-functional changes?

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

3340

Maybe add an assertion here that HalfType is a legal type? Some architectures could trip over this...

lib/CodeGen/SelectionDAG/TargetLowering.cpp

3305

Should there be an assertion here? If expandMUL_LOHI isn't returning exactly two results, something went wrong.

[nhaehnle]

Thank you for taking a look.

D26628 does not directly interact, since it affected a custom sequence used
by DAGTypeLegalizer if the call to TargetLowering has failed. I suppose now
that TargetLowering no longer insists on only using legal instructions one
could think about using that mode instead of the custom sequence in
DAGTypeLegalizer, but that sequence is also qualitatively different (using a
regular NxN-bit multiply on values where everything but the low N/2 bits
have been masked out).

I've addressed your minor comments and split this up into three patches,
the first two being D27063 and D27064.

I'm going to add them as dependencies as well. Unfortunately, handling
patch series with Phabricator seems really awkward :(

[nhaehnle]

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

3340

I missed this one, sorry.

[nhaehnle]

Add assertion that HalfType is legal.

[nhaehnle]

Ping.

[efriedma]

include/llvm/Target/TargetLowering.h

3060

Please make OnlyLegalOrCustom an enum; otherwise it's completely impossible to tell what the parameter does without looking at the header.

lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

2212

Can you replace this code with a call to expandMUL?

[nhaehnle]

• Use an enum instead of OnlyLegalOrCustom
• Handle a corner case where shift amounts are too large for the native shift type

[nhaehnle]

lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

2212

I played around with that, but I'd rather not do it in this patch. The generated code sequences are different, and in particular we'd generate non-legal instructions that at least the X86 target doesn't expect and doesn't handle properly. This is visible in test/CodeGen/X86/imul-256/512/1024.ll.

[efriedma]

LGTM with minor comments addressed.

lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

2212

Okay, that's fine.

lib/CodeGen/SelectionDAG/TargetLowering.cpp

3155

MaskedValueIsZero should do the right thing for vectors, I think, with recent changes to computeKnownBits.

3181

Hmm...

I was going to say that ShiftAmount always fits into ShiftAmountTy because it's a defined shift (ShiftAmount < OuterBitSize < maximum unsigned value), but getShiftAmountTy doesn't always return a sensible result for illegal types. Please just note this problem with a FIXME.

[chapuni]

I pruned three of \param(s) in r289057. Could you recheck them?

llvm/trunk/include/llvm/Target/TargetLowering.h

3051 ↗

(On Diff #80754)

HalfVT and OnlyLegalOrCustom cannot be seen here.

3067 ↗

(On Diff #80754)

ditto.

include/llvm/Target/TargetLowering.h

enum class AtomicExpansionKind {
     None,    // Don't expand the instruction.
     LLSC,    // Expand the instruction into loadlinked/storeconditional; used
              // by ARM/AArch64.
     LLOnly,  // Expand the (load) instruction into just a load-linked, which has
              // greater atomic guarantees than a normal load.
     CmpXChg, // Expand the instruction into cmpxchg; used by at least X86.
   };
 
+  /// Enum that specifies when a multiplication should be expanded.
+  enum class MulExpansionKind {
+    Always,            // Always expand the instruction.
+    OnlyLegalOrCustom, // Only expand when the resulting instructions are legal
+                       // or custom.
+  };
+
   static ISD::NodeType getExtendForContent(BooleanContent Content) {
     switch (Content) {
     case UndefinedBooleanContent:
       // Extend by adding rubbish bits.
       return ISD::ANY_EXTEND;
     case ZeroOrOneBooleanContent:
       // Extend by adding zero bits.
       return ISD::ZERO_EXTEND;
virtual SDValue getRecipEstimate(SDValue Operand, SelectionDAG &DAG,
                                    int Enabled, int &RefinementSteps) const {
     return SDValue();
   }
 
   //===--------------------------------------------------------------------===//
   // Legalization utility functions
   //
 
+  /// Expand a MUL or [US]MUL_LOHI of n-bit values into two or four nodes,
+  /// respectively, each computing an n/2-bit part of the result.
+  /// \param Result A vector that will be filled with the parts of the result
+  ///        in little-endian order.
+  /// \param HalfVT The value type to use for the result nodes.
+  /// \param OnlyLegalOrCustom Only legal or custom instructions are used.
+  /// \param LL Low bits of the LHS of the MUL.  You can use this parameter
+  ///        if you want to control how low bits are extracted from the LHS.
+  /// \param LH High bits of the LHS of the MUL.  See LL for meaning.
+  /// \param RL Low bits of the RHS of the MUL.  See LL for meaning
+  /// \param RH High bits of the RHS of the MUL.  See LL for meaning.
+  /// \returns true if the node has been expanded, false if it has not
+  bool expandMUL_LOHI(unsigned Opcode, EVT VT, SDLoc dl, SDValue LHS,
+                      SDValue RHS, SmallVectorImpl<SDValue> &Result, EVT HiLoVT,
+                      SelectionDAG &DAG, MulExpansionKind Kind,
+                      SDValue LL = SDValue(), SDValue LH = SDValue(),
+                      SDValue RL = SDValue(), SDValue RH = SDValue()) const;
+
   /// Expand a MUL into two nodes.  One that computes the high bits of
   /// the result and one that computes the low bits.
   /// \param HiLoVT The value type to use for the Lo and Hi nodes.
+  /// \param OnlyLegalOrCustom Only legal or custom instructions are used.
   /// \param LL Low bits of the LHS of the MUL.  You can use this parameter
   ///        if you want to control how low bits are extracted from the LHS.
   /// \param LH High bits of the LHS of the MUL.  See LL for meaning.
   /// \param RL Low bits of the RHS of the MUL.  See LL for meaning
   /// \param RH High bits of the RHS of the MUL.  See LL for meaning.
   /// \returns true if the node has been expanded. false if it has not
   bool expandMUL(SDNode *N, SDValue &Lo, SDValue &Hi, EVT HiLoVT,
-                 SelectionDAG &DAG, SDValue LL = SDValue(),
-                 SDValue LH = SDValue(), SDValue RL = SDValue(),
-                 SDValue RH = SDValue()) const;
+                 SelectionDAG &DAG, MulExpansionKind Kind,
+                 SDValue LL = SDValue(), SDValue LH = SDValue(),
+                 SDValue RL = SDValue(), SDValue RH = SDValue()) const;
 
   /// Expand float(f32) to SINT(i64) conversion
   /// \param N Node to expand
   /// \param Result output after conversion
   /// \returns True, if the expansion was successful, false otherwise
   bool expandFP_TO_SINT(SDNode *N, SDValue &Result, SelectionDAG &DAG) const;
 
   /// Turn load of vector type into a load of the individual elements.

lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

if (TLI.isOperationLegalOrCustom(DivRemOpc, VT)) {
       Tmp1 = DAG.getNode(DivRemOpc, dl, VTs, Node->getOperand(0),
                          Node->getOperand(1));
       Results.push_back(Tmp1);
     }
     break;
   }
   case ISD::MULHU:
   case ISD::MULHS: {
-    unsigned ExpandOpcode = Node->getOpcode() == ISD::MULHU ? ISD::UMUL_LOHI :
-                                                              ISD::SMUL_LOHI;
+    unsigned ExpandOpcode =
+        Node->getOpcode() == ISD::MULHU ? ISD::UMUL_LOHI : ISD::SMUL_LOHI;
     EVT VT = Node->getValueType(0);
     SDVTList VTs = DAG.getVTList(VT, VT);
-    assert(TLI.isOperationLegalOrCustom(ExpandOpcode, VT) &&
-           "If this wasn't legal, it shouldn't have been created!");
+
     Tmp1 = DAG.getNode(ExpandOpcode, dl, VTs, Node->getOperand(0),
                        Node->getOperand(1));
     Results.push_back(Tmp1.getValue(1));
     break;
   }
+  case ISD::UMUL_LOHI:
+  case ISD::SMUL_LOHI: {
+    SDValue LHS = Node->getOperand(0);
+    SDValue RHS = Node->getOperand(1);
+    MVT VT = LHS.getSimpleValueType();
+    unsigned MULHOpcode =
+        Node->getOpcode() == ISD::UMUL_LOHI ? ISD::MULHU : ISD::MULHS;
+
+    if (TLI.isOperationLegalOrCustom(MULHOpcode, VT)) {
+      Results.push_back(DAG.getNode(ISD::MUL, dl, VT, LHS, RHS));
+      Results.push_back(DAG.getNode(MULHOpcode, dl, VT, LHS, RHS));
+      break;
+    }
+
+    SmallVector<SDValue, 4> Halves;
+    EVT HalfType = EVT(VT).getHalfSizedIntegerVT(*DAG.getContext());
+    assert(TLI.isTypeLegal(HalfType));
+    if (TLI.expandMUL_LOHI(Node->getOpcode(), VT, Node, LHS, RHS, Halves,
+                           HalfType, DAG,
+                           TargetLowering::MulExpansionKind::Always)) {
+      for (unsigned i = 0; i < 2; ++i) {
+        SDValue Lo = DAG.getNode(ISD::ZERO_EXTEND, dl, VT, Halves[2 * i]);
+        SDValue Hi = DAG.getNode(ISD::ANY_EXTEND, dl, VT, Halves[2 * i + 1]);
+        SDValue Shift = DAG.getConstant(
+            HalfType.getScalarSizeInBits(), dl,
+            TLI.getShiftAmountTy(HalfType, DAG.getDataLayout()));
+        Hi = DAG.getNode(ISD::SHL, dl, VT, Hi, Shift);
+        Results.push_back(DAG.getNode(ISD::OR, dl, VT, Lo, Hi));
+      }
+      break;
+    }
+    break;
+  }
   case ISD::MUL: {
     EVT VT = Node->getValueType(0);
     SDVTList VTs = DAG.getVTList(VT, VT);
     // See if multiply or divide can be lowered using two-result operations.
     // We just need the low half of the multiply; try both the signed
     // and unsigned forms. If the target supports both SMUL_LOHI and
     // UMUL_LOHI, form a preference by checking which forms of plain
     // MULH it supports.
case ISD::MUL: {
     }
 
     SDValue Lo, Hi;
     EVT HalfType = VT.getHalfSizedIntegerVT(*DAG.getContext());
     if (TLI.isOperationLegalOrCustom(ISD::ZERO_EXTEND, VT) &&
         TLI.isOperationLegalOrCustom(ISD::ANY_EXTEND, VT) &&
         TLI.isOperationLegalOrCustom(ISD::SHL, VT) &&
         TLI.isOperationLegalOrCustom(ISD::OR, VT) &&
-        TLI.expandMUL(Node, Lo, Hi, HalfType, DAG)) {
+        TLI.expandMUL(Node, Lo, Hi, HalfType, DAG,
+                      TargetLowering::MulExpansionKind::OnlyLegalOrCustom)) {
       Lo = DAG.getNode(ISD::ZERO_EXTEND, dl, VT, Lo);
       Hi = DAG.getNode(ISD::ANY_EXTEND, dl, VT, Hi);
       SDValue Shift =
           DAG.getConstant(HalfType.getSizeInBits(), dl,
                           TLI.getShiftAmountTy(HalfType, DAG.getDataLayout()));
       Hi = DAG.getNode(ISD::SHL, dl, VT, Hi, Shift);
       Results.push_back(DAG.getNode(ISD::OR, dl, VT, Lo, Hi));
     }
case ISD::XOR: {
     // Promote each of the values to the new type.
     Tmp1 = DAG.getNode(ExtOp, dl, NVT, Node->getOperand(0));
     Tmp2 = DAG.getNode(ExtOp, dl, NVT, Node->getOperand(1));
     // Perform the larger operation, then convert back
     Tmp1 = DAG.getNode(Node->getOpcode(), dl, NVT, Tmp1, Tmp2);
     Results.push_back(DAG.getNode(TruncOp, dl, OVT, Tmp1));
     break;
   }
+  case ISD::UMUL_LOHI:
+  case ISD::SMUL_LOHI: {
+    // Promote to a multiply in a wider integer type.
+    unsigned ExtOp = Node->getOpcode() == ISD::UMUL_LOHI ? ISD::ZERO_EXTEND
+                                                         : ISD::SIGN_EXTEND;
+    Tmp1 = DAG.getNode(ExtOp, dl, NVT, Node->getOperand(0));
+    Tmp2 = DAG.getNode(ExtOp, dl, NVT, Node->getOperand(1));
+    Tmp1 = DAG.getNode(ISD::MUL, dl, NVT, Tmp1, Tmp2);
+
+    auto &DL = DAG.getDataLayout();
+    unsigned OriginalSize = OVT.getScalarSizeInBits();
+    Tmp2 = DAG.getNode(
+        ISD::SRL, dl, NVT, Tmp1,
+        DAG.getConstant(OriginalSize, dl, TLI.getScalarShiftAmountTy(DL, NVT)));
+    Results.push_back(DAG.getNode(ISD::TRUNCATE, dl, OVT, Tmp1));
+    Results.push_back(DAG.getNode(ISD::TRUNCATE, dl, OVT, Tmp2));
+    break;
+  }
   case ISD::SELECT: {
     unsigned ExtOp, TruncOp;
     if (Node->getValueType(0).isVector() ||
         Node->getValueType(0).getSizeInBits() == NVT.getSizeInBits()) {
       ExtOp   = ISD::BITCAST;
       TruncOp = ISD::BITCAST;
     } else if (Node->getValueType(0).isInteger()) {
       ExtOp   = ISD::ANY_EXTEND;

lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

void DAGTypeLegalizer::ExpandIntRes_MUL(SDNode *N,
   EVT VT = N->getValueType(0);
   EVT NVT = TLI.getTypeToTransformTo(*DAG.getContext(), VT);
   SDLoc dl(N);
 
   SDValue LL, LH, RL, RH;
   GetExpandedInteger(N->getOperand(0), LL, LH);
   GetExpandedInteger(N->getOperand(1), RL, RH);
 
-  if (TLI.expandMUL(N, Lo, Hi, NVT, DAG, LL, LH, RL, RH))
+  if (TLI.expandMUL(N, Lo, Hi, NVT, DAG,
+                    TargetLowering::MulExpansionKind::OnlyLegalOrCustom,
+                    LL, LH, RL, RH))
     return;
 
   // If nothing else, we can make a libcall.
   RTLIB::Libcall LC = RTLIB::UNKNOWN_LIBCALL;
   if (VT == MVT::i16)
     LC = RTLIB::MUL_I16;
   else if (VT == MVT::i32)
     LC = RTLIB::MUL_I32;
   else if (VT == MVT::i64)
     LC = RTLIB::MUL_I64;
   else if (VT == MVT::i128)
     LC = RTLIB::MUL_I128;
 
   if (LC == RTLIB::UNKNOWN_LIBCALL) {
     // We'll expand the multiplication by brute force because we have no other
     // options. This is a trivially-generalized version of the code from
     // Hacker's Delight (itself derived from Knuth's Algorithm M from section
     // 4.3.1).
     unsigned Bits = NVT.getSizeInBits();
     unsigned HalfBits = Bits >> 1;
     SDValue Mask = DAG.getConstant(APInt::getLowBitsSet(Bits, HalfBits), dl,
                                    NVT);
     SDValue LLL = DAG.getNode(ISD::AND, dl, NVT, LL, Mask);
     SDValue RLL = DAG.getNode(ISD::AND, dl, NVT, RL, Mask);
 
     SDValue T = DAG.getNode(ISD::MUL, dl, NVT, LLL, RLL);

lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

SDValue VectorLegalizer::LegalizeOp(SDValue Op) {
   case ISD::SIGN_EXTEND_INREG:
   case ISD::ANY_EXTEND_VECTOR_INREG:
   case ISD::SIGN_EXTEND_VECTOR_INREG:
   case ISD::ZERO_EXTEND_VECTOR_INREG:
   case ISD::SMIN:
   case ISD::SMAX:
   case ISD::UMIN:
   case ISD::UMAX:
+  case ISD::SMUL_LOHI:
+  case ISD::UMUL_LOHI:
     QueryType = Node->getValueType(0);
     break;
   case ISD::FP_ROUND_INREG:
     QueryType = cast<VTSDNode>(Node->getOperand(1))->getVT();
     break;
   case ISD::SINT_TO_FP:
   case ISD::UINT_TO_FP:
     QueryType = Node->getOperand(0).getValueType();

lib/CodeGen/SelectionDAG/TargetLowering.cpp

verifyReturnAddressArgumentIsConstant(SDValue Op, SelectionDAG &DAG) const {
 
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Legalization Utilities
 //===----------------------------------------------------------------------===//
 
-bool TargetLowering::expandMUL(SDNode *N, SDValue &Lo, SDValue &Hi, EVT HiLoVT,
-                               SelectionDAG &DAG, SDValue LL, SDValue LH,
-                               SDValue RL, SDValue RH) const {
-  EVT VT = N->getValueType(0);
-  SDLoc dl(N);
-
-  bool HasMULHS = isOperationLegalOrCustom(ISD::MULHS, HiLoVT);
-  bool HasMULHU = isOperationLegalOrCustom(ISD::MULHU, HiLoVT);
-  bool HasSMUL_LOHI = isOperationLegalOrCustom(ISD::SMUL_LOHI, HiLoVT);
-  bool HasUMUL_LOHI = isOperationLegalOrCustom(ISD::UMUL_LOHI, HiLoVT);
+bool TargetLowering::expandMUL_LOHI(unsigned Opcode, EVT VT, SDLoc dl,
+                                    SDValue LHS, SDValue RHS,
+                                    SmallVectorImpl<SDValue> &Result,
+                                    EVT HiLoVT, SelectionDAG &DAG,
+                                    MulExpansionKind Kind, SDValue LL,
+                                    SDValue LH, SDValue RL, SDValue RH) const {
+  assert(Opcode == ISD::MUL || Opcode == ISD::UMUL_LOHI ||
+         Opcode == ISD::SMUL_LOHI);
+
+  bool HasMULHS = (Kind == MulExpansionKind::Always) ||
+                  isOperationLegalOrCustom(ISD::MULHS, HiLoVT);
+  bool HasMULHU = (Kind == MulExpansionKind::Always) ||
+                  isOperationLegalOrCustom(ISD::MULHU, HiLoVT);
+  bool HasSMUL_LOHI = (Kind == MulExpansionKind::Always) ||
+                      isOperationLegalOrCustom(ISD::SMUL_LOHI, HiLoVT);
+  bool HasUMUL_LOHI = (Kind == MulExpansionKind::Always) ||
+                      isOperationLegalOrCustom(ISD::UMUL_LOHI, HiLoVT);
 
   if (!HasMULHU && !HasMULHS && !HasUMUL_LOHI && !HasSMUL_LOHI)
     return false;
 
   unsigned OuterBitSize = VT.getScalarSizeInBits();
   unsigned InnerBitSize = HiLoVT.getScalarSizeInBits();
-  SDValue LHS = N->getOperand(0);
-  SDValue RHS = N->getOperand(1);
   unsigned LHSSB = DAG.ComputeNumSignBits(LHS);
   unsigned RHSSB = DAG.ComputeNumSignBits(RHS);
 
   // LL, LH, RL, and RH must be either all NULL or all set to a value.
   assert((LL.getNode() && LH.getNode() && RL.getNode() && RH.getNode()) ||
          (!LL.getNode() && !LH.getNode() && !RL.getNode() && !RH.getNode()));
 
   SDVTList VTs = DAG.getVTList(HiLoVT, HiLoVT);
   auto MakeMUL_LOHI = [&](SDValue L, SDValue R, SDValue &Lo, SDValue &Hi,
                           bool Signed) -> bool {
     if ((Signed && HasSMUL_LOHI) || (!Signed && HasUMUL_LOHI)) {
       Lo = DAG.getNode(Signed ? ISD::SMUL_LOHI : ISD::UMUL_LOHI, dl, VTs, L, R);
       Hi = SDValue(Lo.getNode(), 1);
       return true;
     }
     if ((Signed && HasMULHS) || (!Signed && HasMULHU)) {
       Lo = DAG.getNode(ISD::MUL, dl, HiLoVT, L, R);
       Hi = DAG.getNode(Signed ? ISD::MULHS : ISD::MULHU, dl, HiLoVT, L, R);
       return true;
     }
     return false;
   };
 
+  SDValue Lo, Hi;
+
   if (!LL.getNode() && !RL.getNode() &&
       isOperationLegalOrCustom(ISD::TRUNCATE, HiLoVT)) {
     LL = DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, LHS);
     RL = DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, RHS);
   }
 
   if (!LL.getNode())
     return false;
 
   APInt HighMask = APInt::getHighBitsSet(OuterBitSize, InnerBitSize);
-  if (DAG.MaskedValueIsZero(N->getOperand(0), HighMask) &&
-      DAG.MaskedValueIsZero(N->getOperand(1), HighMask)) {
+  if (!VT.isVector() && DAG.MaskedValueIsZero(LHS, HighMask) &&
+      DAG.MaskedValueIsZero(RHS, HighMask)) {
     // The inputs are both zero-extended.
-    if (MakeMUL_LOHI(LL, RL, Lo, Hi, false))
+    if (MakeMUL_LOHI(LL, RL, Lo, Hi, false)) {
+      Result.push_back(Lo);
+      Result.push_back(Hi);
+      if (Opcode != ISD::MUL) {
+        SDValue Zero = DAG.getConstant(0, dl, HiLoVT);
+        Result.push_back(Zero);
+        Result.push_back(Zero);
+      }
       return true;
     }
-  if (LHSSB > InnerBitSize && RHSSB > InnerBitSize) {
+  }
+
+  if (!VT.isVector() && Opcode == ISD::MUL && LHSSB > InnerBitSize &&
+      RHSSB > InnerBitSize) {
     // The input values are both sign-extended.
-    if (MakeMUL_LOHI(LL, RL, Lo, Hi, true))
+    // TODO non-MUL case?
+    if (MakeMUL_LOHI(LL, RL, Lo, Hi, true)) {
+      Result.push_back(Lo);
+      Result.push_back(Hi);
       return true;
     }
+  }
 
-  SDValue Shift = DAG.getConstant(OuterBitSize - InnerBitSize, dl,
-                                  getShiftAmountTy(VT, DAG.getDataLayout()));
+  unsigned ShiftAmount = OuterBitSize - InnerBitSize;
+  EVT ShiftAmountTy = getShiftAmountTy(VT, DAG.getDataLayout());
+  if (APInt::getMaxValue(ShiftAmountTy.getSizeInBits()).ult(ShiftAmount)) {
+    // The type from TLI is too small to fit the shift amount we want.
+    // Override it with i32. The shift will have to be legalized.
+    ShiftAmountTy = MVT::i32;
+  }
+  SDValue Shift = DAG.getConstant(ShiftAmount, dl, ShiftAmountTy);
 
   if (!LH.getNode() && !RH.getNode() &&
       isOperationLegalOrCustom(ISD::SRL, VT) &&
       isOperationLegalOrCustom(ISD::TRUNCATE, HiLoVT)) {
     LH = DAG.getNode(ISD::SRL, dl, VT, LHS, Shift);
     LH = DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, LH);
     RH = DAG.getNode(ISD::SRL, dl, VT, RHS, Shift);
     RH = DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, RH);
   }
 
   if (!LH.getNode())
     return false;
 
-  if (MakeMUL_LOHI(LL, RL, Lo, Hi, false)) {
+  if (!MakeMUL_LOHI(LL, RL, Lo, Hi, false))
+    return false;
+
+  Result.push_back(Lo);
+
+  if (Opcode == ISD::MUL) {
     RH = DAG.getNode(ISD::MUL, dl, HiLoVT, LL, RH);
     LH = DAG.getNode(ISD::MUL, dl, HiLoVT, LH, RL);
     Hi = DAG.getNode(ISD::ADD, dl, HiLoVT, Hi, RH);
     Hi = DAG.getNode(ISD::ADD, dl, HiLoVT, Hi, LH);
+    Result.push_back(Hi);
     return true;
   }
 
+  // Compute the full width result.
+  auto Merge = [&](SDValue Lo, SDValue Hi) -> SDValue {
+    Lo = DAG.getNode(ISD::ZERO_EXTEND, dl, VT, Lo);
+    Hi = DAG.getNode(ISD::ZERO_EXTEND, dl, VT, Hi);
+    Hi = DAG.getNode(ISD::SHL, dl, VT, Hi, Shift);
+    return DAG.getNode(ISD::OR, dl, VT, Lo, Hi);
+  };
+
+  SDValue Next = DAG.getNode(ISD::ZERO_EXTEND, dl, VT, Hi);
+  if (!MakeMUL_LOHI(LL, RH, Lo, Hi, false))
     return false;
+
+  // This is effectively the add part of a multiply-add of half-sized operands,
+  // so it cannot overflow.
+  Next = DAG.getNode(ISD::ADD, dl, VT, Next, Merge(Lo, Hi));
+
+  if (!MakeMUL_LOHI(LH, RL, Lo, Hi, false))
+    return false;
+
+  Next = DAG.getNode(ISD::ADDC, dl, DAG.getVTList(VT, MVT::Glue), Next,
+                     Merge(Lo, Hi));
+
+  SDValue Carry = Next.getValue(1);
+  Result.push_back(DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, Next));
+  Next = DAG.getNode(ISD::SRL, dl, VT, Next, Shift);
+
+  if (!MakeMUL_LOHI(LH, RH, Lo, Hi, Opcode == ISD::SMUL_LOHI))
+    return false;
+
+  SDValue Zero = DAG.getConstant(0, dl, HiLoVT);
+  Hi = DAG.getNode(ISD::ADDE, dl, DAG.getVTList(HiLoVT, MVT::Glue), Hi, Zero,
+                   Carry);
+  Next = DAG.getNode(ISD::ADD, dl, VT, Next, Merge(Lo, Hi));
+
+  if (Opcode == ISD::SMUL_LOHI) {
+    SDValue NextSub = DAG.getNode(ISD::SUB, dl, VT, Next,
+                                  DAG.getNode(ISD::ZERO_EXTEND, dl, VT, RL));
+    Next = DAG.getSelectCC(dl, LH, Zero, NextSub, Next, ISD::SETLT);
+
+    NextSub = DAG.getNode(ISD::SUB, dl, VT, Next,
+                          DAG.getNode(ISD::ZERO_EXTEND, dl, VT, LL));
+    Next = DAG.getSelectCC(dl, RH, Zero, NextSub, Next, ISD::SETLT);
+  }
+
+  Result.push_back(DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, Next));
+  Next = DAG.getNode(ISD::SRL, dl, VT, Next, Shift);
+  Result.push_back(DAG.getNode(ISD::TRUNCATE, dl, HiLoVT, Next));
+  return true;
+}
+
+bool TargetLowering::expandMUL(SDNode *N, SDValue &Lo, SDValue &Hi, EVT HiLoVT,
+                               SelectionDAG &DAG, MulExpansionKind Kind,
+                               SDValue LL, SDValue LH, SDValue RL,
+                               SDValue RH) const {
+  SmallVector<SDValue, 2> Result;
+  bool Ok = expandMUL_LOHI(N->getOpcode(), N->getValueType(0), N,
+                           N->getOperand(0), N->getOperand(1), Result, HiLoVT,
+                           DAG, Kind, LL, LH, RL, RH);
+  if (Ok) {
+    assert(Result.size() == 2);
+    Lo = Result[0];
+    Hi = Result[1];
+  }
+  return Ok;
 }
 
 bool TargetLowering::expandFP_TO_SINT(SDNode *Node, SDValue &Result,
                                SelectionDAG &DAG) const {
   EVT VT = Node->getOperand(0).getValueType();
   EVT NVT = Node->getValueType(0);
   SDLoc dl(SDValue(Node, 0));
 
bool TargetLowering::expandFP_TO_SINT(SDNode *Node, SDValue &Result,
   SDValue ExponentLoBit = DAG.getConstant(23, dl, IntVT);
   SDValue Bias = DAG.getConstant(127, dl, IntVT);
   SDValue SignMask = DAG.getConstant(APInt::getSignBit(VT.getSizeInBits()), dl,
                                      IntVT);
   SDValue SignLowBit = DAG.getConstant(VT.getSizeInBits() - 1, dl, IntVT);
   SDValue MantissaMask = DAG.getConstant(0x007FFFFF, dl, IntVT);
 
   SDValue Bits = DAG.getNode(ISD::BITCAST, dl, IntVT, Node->getOperand(0));
 
   auto &DL = DAG.getDataLayout();
   SDValue ExponentBits = DAG.getNode(
       ISD::SRL, dl, IntVT, DAG.getNode(ISD::AND, dl, IntVT, Bits, ExponentMask),
       DAG.getZExtOrTrunc(ExponentLoBit, dl, getShiftAmountTy(IntVT, DL)));
   SDValue Exponent = DAG.getNode(ISD::SUB, dl, IntVT, ExponentBits, Bias);
 
   SDValue Sign = DAG.getNode(
       ISD::SRA, dl, IntVT, DAG.getNode(ISD::AND, dl, IntVT, Bits, SignMask),