1 files changed, 31 insertions, 23 deletions

diff --git a/llvm/lib/CodeGen/GlobalISel/MachineLegalizer.cpp b/llvm/lib/CodeGen/GlobalISel/MachineLegalizer.cpp
index fa23627de7c..dd32b1fe10f 100644
--- a/llvm/lib/CodeGen/GlobalISel/MachineLegalizer.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/MachineLegalizer.cpp
@@ -35,12 +35,12 @@ MachineLegalizer::MachineLegalizer() : TablesInitialized(false) {
 
 void MachineLegalizer::computeTables() {
   for (auto &Op : Actions) {
-    LLT Ty = Op.first.second;
+    LLT Ty = Op.first.Type;
     if (!Ty.isVector())
       continue;
 
-    auto &Entry =
-        MaxLegalVectorElts[std::make_pair(Op.first.first, Ty.getElementType())];
+    auto &Entry = MaxLegalVectorElts[std::make_pair(Op.first.Opcode,
+                                                    Ty.getElementType())];
     Entry = std::max(Entry, Ty.getNumElements());
   }
 
@@ -52,27 +52,30 @@ void MachineLegalizer::computeTables() {
 // we have any hope of doing well with something like <13 x i3>. Even the common
 // cases should do better than what we have now.
 std::pair<MachineLegalizer::LegalizeAction, LLT>
-MachineLegalizer::getAction(unsigned Opcode, LLT Ty) const {
+MachineLegalizer::getAction(const InstrAspect &Aspect) const {
   assert(TablesInitialized && "backend forgot to call computeTables");
   // These *have* to be implemented for now, they're the fundamental basis of
   // how everything else is transformed.
 
   // FIXME: the long-term plan calls for expansion in terms of load/store (if
   // they're not legal).
-  if (Opcode == TargetOpcode::G_SEQUENCE || Opcode == TargetOpcode::G_EXTRACT)
-    return std::make_pair(Legal, Ty);
+  if (Aspect.Opcode == TargetOpcode::G_SEQUENCE ||
+      Aspect.Opcode == TargetOpcode::G_EXTRACT)
+    return std::make_pair(Legal, Aspect.Type);
 
-  auto ActionIt = Actions.find(std::make_pair(Opcode, Ty));
+  auto ActionIt = Actions.find(Aspect);
   if (ActionIt != Actions.end())
-    return findLegalAction(Opcode, Ty, ActionIt->second);
+    return findLegalAction(Aspect, ActionIt->second);
 
+  unsigned Opcode = Aspect.Opcode;
+  LLT Ty = Aspect.Type;
   if (!Ty.isVector()) {
-    auto DefaultAction = DefaultActions.find(Opcode);
+    auto DefaultAction = DefaultActions.find(Aspect.Opcode);
     if (DefaultAction != DefaultActions.end() && DefaultAction->second == Legal)
       return std::make_pair(Legal, Ty);
 
     assert(DefaultAction->second == NarrowScalar && "unexpected default");
-    return findLegalAction(Opcode, Ty, NarrowScalar);
+    return findLegalAction(Aspect, NarrowScalar);
   }
 
   LLT EltTy = Ty.getElementType();
@@ -81,13 +84,13 @@ MachineLegalizer::getAction(unsigned Opcode, LLT Ty) const {
   auto ScalarAction = ScalarInVectorActions.find(std::make_pair(Opcode, EltTy));
   if (ScalarAction != ScalarInVectorActions.end() &&
       ScalarAction->second != Legal)
-    return findLegalAction(Opcode, EltTy, ScalarAction->second);
+    return findLegalAction(Aspect, ScalarAction->second);
 
   // The element type is legal in principle, but the number of elements is
   // wrong.
   auto MaxLegalElts = MaxLegalVectorElts.lookup(std::make_pair(Opcode, EltTy));
   if (MaxLegalElts > NumElts)
-    return findLegalAction(Opcode, Ty, MoreElements);
+    return findLegalAction(Aspect, MoreElements);
 
   if (MaxLegalElts == 0) {
     // Scalarize if there's no legal vector type, which is just a special case
@@ -95,41 +98,46 @@ MachineLegalizer::getAction(unsigned Opcode, LLT Ty) const {
     return std::make_pair(FewerElements, EltTy);
   }
 
-  return findLegalAction(Opcode, Ty, FewerElements);
+  return findLegalAction(Aspect, FewerElements);
 }
 
-std::pair<MachineLegalizer::LegalizeAction, LLT>
+std::tuple<MachineLegalizer::LegalizeAction, unsigned, LLT>
 MachineLegalizer::getAction(const MachineInstr &MI) const {
-  return getAction(MI.getOpcode(), MI.getType());
+  for (unsigned i = 0; i < MI.getNumTypes(); ++i) {
+    auto Action = getAction({MI.getOpcode(), i, MI.getType(i)});
+    if (Action.first != Legal)
+      return {Action.first, i, Action.second};
+  }
+  return {Legal, 0, LLT{}};
 }
 
 bool MachineLegalizer::isLegal(const MachineInstr &MI) const {
-  return getAction(MI).first == Legal;
+  return std::get<0>(getAction(MI)) == Legal;
 }
 
-LLT MachineLegalizer::findLegalType(unsigned Opcode, LLT Ty,
+LLT MachineLegalizer::findLegalType(const InstrAspect &Aspect,
                                     LegalizeAction Action) const {
   switch(Action) {
   default:
     llvm_unreachable("Cannot find legal type");
   case Legal:
-    return Ty;
+    return Aspect.Type;
   case NarrowScalar: {
-    return findLegalType(Opcode, Ty,
+    return findLegalType(Aspect,
                          [&](LLT Ty) -> LLT { return Ty.halfScalarSize(); });
   }
   case WidenScalar: {
-    return findLegalType(Opcode, Ty, [&](LLT Ty) -> LLT {
+    return findLegalType(Aspect, [&](LLT Ty) -> LLT {
       return Ty.getSizeInBits() < 8 ? LLT::scalar(8) : Ty.doubleScalarSize();
     });
   }
   case FewerElements: {
-    return findLegalType(Opcode, Ty,
+    return findLegalType(Aspect,
                          [&](LLT Ty) -> LLT { return Ty.halfElements(); });
   }
   case MoreElements: {
-    return findLegalType(
-        Opcode, Ty, [&](LLT Ty) -> LLT { return Ty.doubleElements(); });
+    return findLegalType(Aspect,
+                         [&](LLT Ty) -> LLT { return Ty.doubleElements(); });
   }
   }
 }