[Legalizer] Refactoring out legalizeMachineFunction

and introducing new unittests/CodeGen/GlobalISel/LegalizerTest.cpp
relying on it to unit test the entire legalizer algorithm (including the
top-level main loop).

See also https://reviews.llvm.org/D71448

5 files changed, 159 insertions, 52 deletions

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/Legalizer.h b/llvm/include/llvm/CodeGen/GlobalISel/Legalizer.h
index 13cf3f7e694..07173b9719b 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/Legalizer.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/Legalizer.h
@@ -31,8 +31,12 @@ class Legalizer : public MachineFunctionPass {
 public:
   static char ID;
 
-private:
+  struct MFResult {
+    bool Changed;
+    const MachineInstr *FailedOn;
+  };
 
+private:
   /// Initialize the field members using \p MF.
   void init(MachineFunction &MF);
 
@@ -55,14 +59,19 @@ public:
   }
 
   MachineFunctionProperties getClearedProperties() const override {
-    return MachineFunctionProperties()
-      .set(MachineFunctionProperties::Property::NoPHIs);
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
   }
 
   bool combineExtracts(MachineInstr &MI, MachineRegisterInfo &MRI,
                        const TargetInstrInfo &TII);
 
   bool runOnMachineFunction(MachineFunction &MF) override;
+
+  static MFResult
+  legalizeMachineFunction(MachineFunction &MF, const LegalizerInfo &LI,
+                          ArrayRef<GISelChangeObserver *> AuxObservers,
+                          MachineIRBuilder &MIRBuilder);
 };
 } // End namespace llvm.
 
diff --git a/llvm/lib/CodeGen/GlobalISel/Legalizer.cpp b/llvm/lib/CodeGen/GlobalISel/Legalizer.cpp
index dbbc501e0fd..966eedd100b 100644
--- a/llvm/lib/CodeGen/GlobalISel/Legalizer.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/Legalizer.cpp
@@ -140,22 +140,13 @@ public:
 };
 } // namespace
 
-bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
-  // If the ISel pipeline failed, do not bother running that pass.
-  if (MF.getProperties().hasProperty(
-          MachineFunctionProperties::Property::FailedISel))
-    return false;
-  LLVM_DEBUG(dbgs() << "Legalize Machine IR for: " << MF.getName() << '\n');
-  init(MF);
-  const TargetPassConfig &TPC = getAnalysis<TargetPassConfig>();
-  GISelCSEAnalysisWrapper &Wrapper =
-      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
-  MachineOptimizationRemarkEmitter MORE(MF, /*MBFI=*/nullptr);
-
-  const size_t NumBlocks = MF.size();
+Legalizer::MFResult
+Legalizer::legalizeMachineFunction(MachineFunction &MF, const LegalizerInfo &LI,
+                                   ArrayRef<GISelChangeObserver *> AuxObservers,
+                                   MachineIRBuilder &MIRBuilder) {
   MachineRegisterInfo &MRI = MF.getRegInfo();
 
-  // Populate Insts
+  // Populate worklists.
   InstListTy InstList;
   ArtifactListTy ArtifactList;
   ReversePostOrderTraversal<MachineFunction *> RPOT(&MF);
@@ -178,40 +169,23 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
   }
   ArtifactList.finalize();
   InstList.finalize();
-  std::unique_ptr<MachineIRBuilder> MIRBuilder;
-  GISelCSEInfo *CSEInfo = nullptr;
-  bool EnableCSE = EnableCSEInLegalizer.getNumOccurrences()
-                       ? EnableCSEInLegalizer
-                       : TPC.isGISelCSEEnabled();
 
-  if (EnableCSE) {
-    MIRBuilder = std::make_unique<CSEMIRBuilder>();
-    CSEInfo = &Wrapper.get(TPC.getCSEConfig());
-    MIRBuilder->setCSEInfo(CSEInfo);
-  } else
-    MIRBuilder = std::make_unique<MachineIRBuilder>();
-  // This observer keeps the worklist updated.
+  // This observer keeps the worklists updated.
   LegalizerWorkListManager WorkListObserver(InstList, ArtifactList);
-  // We want both WorkListObserver as well as CSEInfo to observe all changes.
-  // Use the wrapper observer.
+  // We want both WorkListObserver as well as all the auxiliary observers (e.g.
+  // CSEInfo) to observe all changes. Use the wrapper observer.
   GISelObserverWrapper WrapperObserver(&WorkListObserver);
-  if (EnableCSE && CSEInfo)
-    WrapperObserver.addObserver(CSEInfo);
+  for (GISelChangeObserver *Observer : AuxObservers)
+    WrapperObserver.addObserver(Observer);
+
   // Now install the observer as the delegate to MF.
   // This will keep all the observers notified about new insertions/deletions.
   RAIIDelegateInstaller DelInstall(MF, &WrapperObserver);
-  LegalizerHelper Helper(MF, WrapperObserver, *MIRBuilder.get());
-  const LegalizerInfo &LInfo(Helper.getLegalizerInfo());
-  LegalizationArtifactCombiner ArtCombiner(*MIRBuilder.get(), MF.getRegInfo(),
-                                           LInfo);
+  LegalizerHelper Helper(MF, LI, WrapperObserver, MIRBuilder);
+  LegalizationArtifactCombiner ArtCombiner(MIRBuilder, MRI, LI);
   auto RemoveDeadInstFromLists = [&WrapperObserver](MachineInstr *DeadMI) {
     WrapperObserver.erasingInstr(*DeadMI);
   };
-  auto stopLegalizing = [&](MachineInstr &MI) {
-    Helper.MIRBuilder.stopObservingChanges();
-    reportGISelFailure(MF, TPC, MORE, "gisel-legalize",
-                       "unable to legalize instruction", MI);
-  };
   bool Changed = false;
   SmallVector<MachineInstr *, 128> RetryList;
   do {
@@ -220,7 +194,8 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
     unsigned NumArtifacts = ArtifactList.size();
     while (!InstList.empty()) {
       MachineInstr &MI = *InstList.pop_back_val();
-      assert(isPreISelGenericOpcode(MI.getOpcode()) && "Expecting generic opcode");
+      assert(isPreISelGenericOpcode(MI.getOpcode()) &&
+             "Expecting generic opcode");
       if (isTriviallyDead(MI, MRI)) {
         LLVM_DEBUG(dbgs() << MI << "Is dead; erasing.\n");
         MI.eraseFromParentAndMarkDBGValuesForRemoval();
@@ -240,8 +215,8 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
           RetryList.push_back(&MI);
           continue;
         }
-        stopLegalizing(MI);
-        return false;
+        Helper.MIRBuilder.stopObservingChanges();
+        return {Changed, &MI};
       }
       WorkListObserver.printNewInstrs();
       Changed |= Res == LegalizerHelper::Legalized;
@@ -254,14 +229,14 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
           ArtifactList.insert(RetryList.pop_back_val());
       } else {
         LLVM_DEBUG(dbgs() << "No new artifacts created, not retrying!\n");
-        MachineInstr *MI = *RetryList.begin();
-        stopLegalizing(*MI);
-        return false;
+        Helper.MIRBuilder.stopObservingChanges();
+        return {Changed, RetryList.front()};
       }
     }
     while (!ArtifactList.empty()) {
       MachineInstr &MI = *ArtifactList.pop_back_val();
-      assert(isPreISelGenericOpcode(MI.getOpcode()) && "Expecting generic opcode");
+      assert(isPreISelGenericOpcode(MI.getOpcode()) &&
+             "Expecting generic opcode");
       if (isTriviallyDead(MI, MRI)) {
         LLVM_DEBUG(dbgs() << MI << "Is dead\n");
         RemoveDeadInstFromLists(&MI);
@@ -291,8 +266,51 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
     }
   } while (!InstList.empty());
 
+  return {Changed, /*FailedOn*/ nullptr};
+}
+
+bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
+  // If the ISel pipeline failed, do not bother running that pass.
+  if (MF.getProperties().hasProperty(
+          MachineFunctionProperties::Property::FailedISel))
+    return false;
+  LLVM_DEBUG(dbgs() << "Legalize Machine IR for: " << MF.getName() << '\n');
+  init(MF);
+  const TargetPassConfig &TPC = getAnalysis<TargetPassConfig>();
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  MachineOptimizationRemarkEmitter MORE(MF, /*MBFI=*/nullptr);
+
+  const size_t NumBlocks = MF.size();
+
+  std::unique_ptr<MachineIRBuilder> MIRBuilder;
+  GISelCSEInfo *CSEInfo = nullptr;
+  bool EnableCSE = EnableCSEInLegalizer.getNumOccurrences()
+                       ? EnableCSEInLegalizer
+                       : TPC.isGISelCSEEnabled();
+  if (EnableCSE) {
+    MIRBuilder = std::make_unique<CSEMIRBuilder>();
+    CSEInfo = &Wrapper.get(TPC.getCSEConfig());
+    MIRBuilder->setCSEInfo(CSEInfo);
+  } else
+    MIRBuilder = std::make_unique<MachineIRBuilder>();
+
+  SmallVector<GISelChangeObserver *, 1> AuxObservers;
+  if (EnableCSE && CSEInfo) {
+    // We want CSEInfo in addition to WorkListObserver to observe all changes.
+    AuxObservers.push_back(CSEInfo);
+  }
+
+  const LegalizerInfo &LI = *MF.getSubtarget().getLegalizerInfo();
+  MFResult Result = legalizeMachineFunction(MF, LI, AuxObservers, *MIRBuilder);
+
+  if (Result.FailedOn) {
+    reportGISelFailure(MF, TPC, MORE, "gisel-legalize",
+                       "unable to legalize instruction", *Result.FailedOn);
+    return false;
+  }
   // For now don't support if new blocks are inserted - we would need to fix the
-  // outerloop for that.
+  // outer loop for that.
   if (MF.size() != NumBlocks) {
     MachineOptimizationRemarkMissed R("gisel-legalize", "GISelFailure",
                                       MF.getFunction().getSubprogram(),
@@ -301,6 +319,5 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
     reportGISelFailure(MF, TPC, MORE, R);
     return false;
   }
-
-  return Changed;
+  return Result.Changed;
 }
diff --git a/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt b/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt
index 01c8f4ecdec..43a0302449c 100644
--- a/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt
+++ b/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt
@@ -12,6 +12,7 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_unittest(GlobalISelTests
   ConstantFoldingTest.cpp
   CSETest.cpp
+  LegalizerTest.cpp
   LegalizerHelperTest.cpp
   LegalizerInfoTest.cpp
   MachineIRBuilderTest.cpp
diff --git a/llvm/unittests/CodeGen/GlobalISel/LegalizerTest.cpp b/llvm/unittests/CodeGen/GlobalISel/LegalizerTest.cpp
new file mode 100644
index 00000000000..7bb348d3046
--- /dev/null
+++ b/llvm/unittests/CodeGen/GlobalISel/LegalizerTest.cpp
@@ -0,0 +1,79 @@
+//===- LegalizerTest.cpp --------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "GISelMITest.h"
+#include "llvm/CodeGen/GlobalISel/Legalizer.h"
+
+using namespace LegalizeActions;
+using namespace LegalizeMutations;
+using namespace LegalityPredicates;
+
+namespace {
+
+::testing::AssertionResult isNullMIPtr(const MachineInstr *MI) {
+  if (MI == nullptr)
+    return ::testing::AssertionSuccess();
+  std::string MIBuffer;
+  raw_string_ostream MISStream(MIBuffer);
+  MI->print(MISStream, /*IsStandalone=*/true, /*SkipOpers=*/false,
+            /*SkipDebugLoc=*/false, /*AddNewLine=*/false);
+  return ::testing::AssertionFailure()
+         << "unable to legalize instruction: " << MISStream.str();
+}
+
+TEST_F(GISelMITest, BasicLegalizerTest) {
+  StringRef MIRString = R"(
+    %vptr:_(p0) = COPY $x4
+    %v:_(<2 x s8>) = G_LOAD %vptr:_(p0) :: (load 2, align 1)
+    $h4 = COPY %v:_(<2 x s8>)
+  )";
+  setUp(MIRString.rtrim(' '));
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(ALegalizer, {
+    auto p0 = LLT::pointer(0, 64);
+    auto v2s8 = LLT::vector(2, 8);
+    auto v2s16 = LLT::vector(2, 16);
+    getActionDefinitionsBuilder(G_LOAD)
+        .legalForTypesWithMemDesc({{s16, p0, 8, 8}})
+        .scalarize(0)
+        .clampScalar(0, s16, s16);
+    getActionDefinitionsBuilder(G_PTR_ADD).legalFor({{p0, s64}});
+    getActionDefinitionsBuilder(G_CONSTANT).legalFor({s64});
+    getActionDefinitionsBuilder(G_BUILD_VECTOR)
+        .legalFor({{v2s16, s16}})
+        .clampScalar(1, s16, s16);
+    getActionDefinitionsBuilder(G_BUILD_VECTOR_TRUNC).legalFor({{v2s8, s16}});
+    getActionDefinitionsBuilder(G_ANYEXT).legalFor({{s32, s16}});
+  });
+
+  ALegalizerInfo LI(MF->getSubtarget());
+
+  Legalizer::MFResult Result =
+      Legalizer::legalizeMachineFunction(*MF, LI, {}, B);
+
+  EXPECT_TRUE(isNullMIPtr(Result.FailedOn));
+  EXPECT_TRUE(Result.Changed);
+
+  StringRef CheckString = R"(
+    CHECK:      %vptr:_(p0) = COPY $x4
+    CHECK-NEXT: [[LOAD_0:%[0-9]+]]:_(s16) = G_LOAD %vptr:_(p0) :: (load 1)
+    CHECK-NEXT: [[OFFSET_1:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+    CHECK-NEXT: [[VPTR_1:%[0-9]+]]:_(p0) = G_PTR_ADD %vptr:_, [[OFFSET_1]]:_(s64)
+    CHECK-NEXT: [[LOAD_1:%[0-9]+]]:_(s16) = G_LOAD [[VPTR_1]]:_(p0) :: (load 1)
+    CHECK-NEXT: [[V0:%[0-9]+]]:_(s16) = COPY [[LOAD_0]]:_(s16)
+    CHECK-NEXT: [[V1:%[0-9]+]]:_(s16) = COPY [[LOAD_1]]:_(s16)
+    CHECK-NEXT: %v:_(<2 x s8>) = G_BUILD_VECTOR_TRUNC [[V0]]:_(s16), [[V1]]:_(s16)
+    CHECK-NEXT: $h4 = COPY %v:_(<2 x s8>)
+  )";
+
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckString)) << *MF;
+}
+
+} // namespace
diff --git a/llvm/utils/gn/secondary/llvm/unittests/CodeGen/GlobalISel/BUILD.gn b/llvm/utils/gn/secondary/llvm/unittests/CodeGen/GlobalISel/BUILD.gn
index 6b222134977..4a042fdf189 100644
--- a/llvm/utils/gn/secondary/llvm/unittests/CodeGen/GlobalISel/BUILD.gn
+++ b/llvm/utils/gn/secondary/llvm/unittests/CodeGen/GlobalISel/BUILD.gn
@@ -16,6 +16,7 @@ unittest("GlobalISelTests") {
     "ConstantFoldingTest.cpp",
     "GISelMITest.cpp",
     "KnownBitsTest.cpp",
+    "LegalizerTest.cpp",
     "LegalizerHelperTest.cpp",
     "LegalizerInfoTest.cpp",
     "MachineIRBuilderTest.cpp",