[GlobalISel] Make multi-step legalization work.

[kristof.beyls]

Summary

In r301116, a custom lowering needed to be introduced to be able to
legalize 8 and 16-bit divisions on ARM targets without a division
instruction, since 2-step legalization (WidenScalar from 8 bit to 32
bit, then Libcall the 32-bit division) doesn't work.

This fixes this and makes this kind of multi-step legalization, where
first the size of the type needs to be changed and then some action is
needed that doesn't require changing the size of the type,
straightforward to specify.

[rovka]

Thanks for splitting this out! The ARM part looks excellent :)
I think you could add some target-independent tests for this in unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp, to make sure it works to WidenScalar/MoreElements etc to a type that has a Libcall or Custom action. That way we won't have to depend on ARM for testing this functionality.

[kristof.beyls]

Adding a unit test.

[kristof.beyls]

In D32529#738192, @rovka wrote:

Thanks for splitting this out! The ARM part looks excellent :)
I think you could add some target-independent tests for this in unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp, to make sure it works to WidenScalar/MoreElements etc to a type that has a Libcall or Custom action. That way we won't have to depend on ARM for testing this functionality.

A yes, that makes sense. Now added.

[kristof.beyls]

Ping.

Thanks for the review on the ARM-specific part, Diana!
@t.p.northover : would you have time for a quick glance at the target-independent parts in LegalizerInfo.h, LegalizerInfo.cpp and LegalizerInfoTest.cpp?

[kristof.beyls]

Rebased to ToT & ping.

[t.p.northover]

Looks fine to me. Just one tiny nit (sorry).

include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

107

Functions should start with lower-case letters unless they need to match surrounding code.

[kristof.beyls]

include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

107

D'oh! Of course.

Thanks for the review!

include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

public:
   LegalizerInfo();
   virtual ~LegalizerInfo() = default;
 
   /// Compute any ancillary tables needed to quickly decide how an operation
   /// should be handled. This must be called after all "set*Action"methods but
   /// before any query is made or incorrect results may be returned.
   void computeTables();
 
+  static bool NeedsLegalizingToDifferentSize(const LegalizeAction Action) {
+    switch (Action) {
+    case NarrowScalar:
+    case WidenScalar:
+    case FewerElements:
+    case MoreElements:
+    case Unsupported:
+      return true;
+    default:
+      return false;
+    }
+  }
+
   /// More friendly way to set an action for common types that have an LLT
   /// representation.
   void setAction(const InstrAspect &Aspect, LegalizeAction Action) {
     TablesInitialized = false;
     unsigned Opcode = Aspect.Opcode - FirstOp;
     if (Actions[Opcode].size() <= Aspect.Idx)
       Actions[Opcode].resize(Aspect.Idx + 1);
     Actions[Aspect.Opcode - FirstOp][Aspect.Idx][Aspect.Type] = Action;
public:
   ///
   /// \returns a tuple consisting of the LegalizeAction that should be
   /// performed, the type-index it should be performed on and the destination
   /// type.
   std::tuple<LegalizeAction, unsigned, LLT>
   getAction(const MachineInstr &MI, const MachineRegisterInfo &MRI) const;
 
   /// Iterate the given function (typically something like doubling the width)
-  /// on Ty until we find a legal type for this operation.
-  LLT findLegalType(const InstrAspect &Aspect,
+  /// on Ty until we find a legalizable size for this operation.
+  LLT findLegalizableSize(const InstrAspect &Aspect,
                           function_ref<LLT(LLT)> NextType) const {
     LegalizeAction Action;
     const TypeMap &Map = Actions[Aspect.Opcode - FirstOp][Aspect.Idx];
     LLT Ty = Aspect.Type;
     do {
       Ty = NextType(Ty);
       auto ActionIt = Map.find(Ty);
       if (ActionIt == Map.end())
         Action = DefaultActions.find(Aspect.Opcode)->second;
       else
         Action = ActionIt->second;
-    } while(Action != Legal);
+    } while(NeedsLegalizingToDifferentSize(Action));
     return Ty;
   }
 
   /// Find what type it's actually OK to perform the given operation on, given
   /// the general approach we've decided to take.
   LLT findLegalType(const InstrAspect &Aspect, LegalizeAction Action) const;
 
   std::pair<LegalizeAction, LLT> findLegalAction(const InstrAspect &Aspect,

lib/CodeGen/GlobalISel/LegalizerInfo.cpp

LLT LegalizerInfo::findLegalType(const InstrAspect &Aspect,
   default:
     llvm_unreachable("Cannot find legal type");
   case Legal:
   case Lower:
   case Libcall:
   case Custom:
     return Aspect.Type;
   case NarrowScalar: {
-    return findLegalType(Aspect,
-                         [&](LLT Ty) -> LLT { return Ty.halfScalarSize(); });
+    return findLegalizableSize(
+        Aspect, [&](LLT Ty) -> LLT { return Ty.halfScalarSize(); });
   }
   case WidenScalar: {
-    return findLegalType(Aspect, [&](LLT Ty) -> LLT {
+    return findLegalizableSize(Aspect, [&](LLT Ty) -> LLT {
       return Ty.getSizeInBits() < 8 ? LLT::scalar(8) : Ty.doubleScalarSize();
     });
   }
   case FewerElements: {
-    return findLegalType(Aspect,
-                         [&](LLT Ty) -> LLT { return Ty.halfElements(); });
+    return findLegalizableSize(
+        Aspect, [&](LLT Ty) -> LLT { return Ty.halfElements(); });
   }
   case MoreElements: {
-    return findLegalType(Aspect,
-                         [&](LLT Ty) -> LLT { return Ty.doubleElements(); });
+    return findLegalizableSize(
+        Aspect, [&](LLT Ty) -> LLT { return Ty.doubleElements(); });
   }
   }
 }
 
 bool LegalizerInfo::legalizeCustom(MachineInstr &MI,
                                    MachineRegisterInfo &MRI,
                                    MachineIRBuilder &MIRBuilder) const {
   return false;
 }

lib/Target/ARM/ARMLegalizerInfo.h

 namespace llvm {
 
 class ARMSubtarget;
 
 /// This class provides the information for the target register banks.
 class ARMLegalizerInfo : public LegalizerInfo {
 public:
   ARMLegalizerInfo(const ARMSubtarget &ST);
-
-  bool legalizeCustom(MachineInstr &MI, MachineRegisterInfo &MRI,
-                      MachineIRBuilder &MIRBuilder) const override;
 };
 } // End llvm namespace.
 #endif

lib/Target/ARM/ARMLegalizerInfo.cpp

ARMLegalizerInfo::ARMLegalizerInfo(const ARMSubtarget &ST) {
   }
 
   for (unsigned Op : {G_ADD, G_SUB, G_MUL})
     for (auto Ty : {s1, s8, s16, s32})
       setAction({Op, Ty}, Legal);
 
   for (unsigned Op : {G_SDIV, G_UDIV}) {
     for (auto Ty : {s8, s16})
-      // FIXME: We need WidenScalar here, but in the case of targets with
-      // software division we'll also need Libcall afterwards. Treat as Custom
-      // until we have better support for chaining legalization actions.
-      setAction({Op, Ty}, Custom);
+      setAction({Op, Ty}, WidenScalar);
     if (ST.hasDivideInARMMode())
       setAction({Op, s32}, Legal);
     else
       setAction({Op, s32}, Libcall);
   }
 
   for (unsigned Op : {G_SEXT, G_ZEXT}) {
     setAction({Op, s32}, Legal);
ARMLegalizerInfo::ARMLegalizerInfo(const ARMSubtarget &ST) {
   }
 
   for (unsigned Op : {G_FREM, G_FPOW})
     for (auto Ty : {s32, s64})
       setAction({Op, Ty}, Libcall);
 
   computeTables();
 }
-
-bool ARMLegalizerInfo::legalizeCustom(MachineInstr &MI,
-                                      MachineRegisterInfo &MRI,
-                                      MachineIRBuilder &MIRBuilder) const {
-  using namespace TargetOpcode;
-
-  switch (MI.getOpcode()) {
-  default:
-    return false;
-  case G_SDIV:
-  case G_UDIV: {
-    LLT Ty = MRI.getType(MI.getOperand(0).getReg());
-    if (Ty != LLT::scalar(16) && Ty != LLT::scalar(8))
-      return false;
-
-    // We need to widen to 32 bits and then maybe, if the target requires,
-    // transform into a libcall.
-    LegalizerHelper Helper(MIRBuilder.getMF());
-
-    MachineInstr *NewMI = nullptr;
-    Helper.MIRBuilder.recordInsertions([&](MachineInstr *MI) {
-      // Store the new, 32-bit div instruction.
-      if (MI->getOpcode() == G_SDIV || MI->getOpcode() == G_UDIV)
-        NewMI = MI;
-    });
-
-    auto Result = Helper.widenScalar(MI, 0, LLT::scalar(32));
-    Helper.MIRBuilder.stopRecordingInsertions();
-    if (Result == LegalizerHelper::UnableToLegalize) {
-      return false;
-    }
-    assert(NewMI && "Couldn't find widened instruction");
-    assert((NewMI->getOpcode() == G_SDIV || NewMI->getOpcode() == G_UDIV) &&
-           "Unexpected widened instruction");
-    assert(MRI.getType(NewMI->getOperand(0).getReg()).getSizeInBits() == 32 &&
-           "Unexpected type for the widened instruction");
-
-    Result = Helper.legalizeInstrStep(*NewMI);
-    if (Result == LegalizerHelper::UnableToLegalize) {
-      return false;
-    }
-    return true;
-  }
-  }
-}

unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp

TEST(LegalizerInfoTest, MultipleTypes) {
   L.computeTables();
 
   // Check we infer the correct types and actually do what we're told.
   ASSERT_EQ(L.getAction({G_PTRTOINT, 0, s64}),
             std::make_pair(LegalizerInfo::Legal, s64));
   ASSERT_EQ(L.getAction({G_PTRTOINT, 1, p0}),
             std::make_pair(LegalizerInfo::Legal, p0));
 }
+
+TEST(LegalizerInfoTest, MultipleSteps) {
+  using namespace TargetOpcode;
+  LegalizerInfo L;
+  LLT s16 = LLT::scalar(16);
+  LLT s32 = LLT::scalar(32);
+  LLT s64 = LLT::scalar(64);
+
+  L.setAction({G_UREM, 0, s16}, LegalizerInfo::WidenScalar);
+  L.setAction({G_UREM, 0, s32}, LegalizerInfo::Lower);
+  L.setAction({G_UREM, 0, s64}, LegalizerInfo::Lower);
+
+  L.computeTables();
+
+  ASSERT_EQ(L.getAction({G_UREM, LLT::scalar(16)}),
+            std::make_pair(LegalizerInfo::WidenScalar, LLT::scalar(32)));
+  ASSERT_EQ(L.getAction({G_UREM, LLT::scalar(32)}),
+            std::make_pair(LegalizerInfo::Lower, LLT::scalar(32)));
+}
 }