3 files changed, 379 insertions, 0 deletions
diff --git a/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt b/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt
index 85f5dcd6d04..60566cb2d59 100644
--- a/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt
+++ b/llvm/unittests/CodeGen/GlobalISel/CMakeLists.txt
@@ -12,4 +12,5 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_unittest(GlobalISelTests
         LegalizerInfoTest.cpp
         PatternMatchTest.cpp
+        LegalizerHelperTest.cpp
         )
diff --git a/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
new file mode 100644
index 00000000000..292281a09a5
--- /dev/null
+++ b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp
@@ -0,0 +1,188 @@
+//===- PatternMatchTest.cpp -----------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "LegalizerHelperTest.h"
+
+namespace {
+
+// Test CTTZ expansion when CTTZ_ZERO_UNDEF is legal or custom,
+// in which case it becomes CTTZ_ZERO_UNDEF with select.
+TEST_F(LegalizerHelperTest, LowerBitCountingCTTZ0) {
+  if (!TM)
+    return;
+
+  // Declare your legalization info
+  DefineLegalizerInfo(
+      A, { getActionDefinitionsBuilder(G_CTTZ_ZERO_UNDEF).legalFor({s64}); });
+  // Build Instr
+  auto MIBCTTZ = B.buildInstr(TargetOpcode::G_CTTZ, LLT::scalar(64), Copies[0]);
+  AInfo Info(MF->getSubtarget());
+  LegalizerHelper Helper(*MF, Info);
+  // Perform Legalization
+  ASSERT_TRUE(Helper.lower(*MIBCTTZ, 0, LLT::scalar(64)) ==
+              LegalizerHelper::LegalizeResult::Legalized);
+
+  auto CheckStr = R"(
+  CHECK: [[CZU:%[0-9]+]]:_(s64) = G_CTTZ_ZERO_UNDEF %0
+  CHECK: [[ZERO:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+  CHECK: [[SIXTY4:%[0-9]+]]:_(s64) = G_CONSTANT i64 64
+  CHECK: [[CMP:%[0-9]+]]:_(s1) = G_ICMP intpred(eq), %0:_(s64), [[ZERO]]
+  CHECK: [[SEL:%[0-9]+]]:_(s64) = G_SELECT [[CMP]]:_(s1), [[SIXTY4]]:_, [[CZU]]
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+
+// CTTZ expansion in terms of CTLZ
+TEST_F(LegalizerHelperTest, LowerBitCountingCTTZ1) {
+  if (!TM)
+    return;
+
+  // Declare your legalization info
+  DefineLegalizerInfo(A,
+                      { getActionDefinitionsBuilder(G_CTLZ).legalFor({s64}); });
+  // Build Instr
+  auto MIBCTTZ = B.buildInstr(TargetOpcode::G_CTTZ, LLT::scalar(64), Copies[0]);
+  AInfo Info(MF->getSubtarget());
+  LegalizerHelper Helper(*MF, Info);
+  // Perform Legalization
+  ASSERT_TRUE(Helper.lower(*MIBCTTZ, 0, LLT::scalar(64)) ==
+              LegalizerHelper::LegalizeResult::Legalized);
+
+  auto CheckStr = R"(
+  CHECK: [[NEG1:%[0-9]+]]:_(s64) = G_CONSTANT i64 -1
+  CHECK: [[NOT:%[0-9]+]]:_(s64) = G_XOR %0:_, [[NEG1]]
+  CHECK: [[SUB1:%[0-9]+]]:_(s64) = G_ADD %0:_, [[NEG1]]
+  CHECK: [[AND1:%[0-9]+]]:_(s64) = G_AND [[NOT]]:_, [[SUB1]]:_
+  CHECK: [[CST64:%[0-9]+]]:_(s64) = G_CONSTANT i64 64
+  CHECK: [[CTLZ:%[0-9]+]]:_(s64) = G_CTLZ [[AND1]]:_
+  CHECK: G_SUB [[CST64]]:_, [[CTLZ]]:_
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+
+// CTTZ expansion in terms of CTPOP
+TEST_F(LegalizerHelperTest, LowerBitCountingCTTZ2) {
+  if (!TM)
+    return;
+
+  // Declare your legalization info
+  DefineLegalizerInfo(
+      A, { getActionDefinitionsBuilder(G_CTPOP).legalFor({s64}); });
+  // Build
+  auto MIBCTTZ = B.buildInstr(TargetOpcode::G_CTTZ, LLT::scalar(64), Copies[0]);
+  AInfo Info(MF->getSubtarget());
+  LegalizerHelper Helper(*MF, Info);
+  ASSERT_TRUE(Helper.lower(*MIBCTTZ, 0, LLT::scalar(64)) ==
+              LegalizerHelper::LegalizeResult::Legalized);
+
+  auto CheckStr = R"(
+  CHECK: [[NEG1:%[0-9]+]]:_(s64) = G_CONSTANT i64 -1
+  CHECK: [[NOT:%[0-9]+]]:_(s64) = G_XOR %0:_, [[NEG1]]
+  CHECK: [[SUB1:%[0-9]+]]:_(s64) = G_ADD %0:_, [[NEG1]]
+  CHECK: [[AND1:%[0-9]+]]:_(s64) = G_AND [[NOT]]:_, [[SUB1]]:_
+  CHECK: [[POP:%[0-9]+]]:_(s64) = G_CTPOP [[AND1]]
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+
+// CTTZ_ZERO_UNDEF expansion in terms of CTTZ
+TEST_F(LegalizerHelperTest, LowerBitCountingCTTZ3) {
+  if (!TM)
+    return;
+
+  // Declare your legalization info
+  DefineLegalizerInfo(A,
+                      { getActionDefinitionsBuilder(G_CTTZ).legalFor({s64}); });
+  // Build
+  auto MIBCTTZ =
+      B.buildInstr(TargetOpcode::G_CTTZ_ZERO_UNDEF, LLT::scalar(64), Copies[0]);
+  AInfo Info(MF->getSubtarget());
+  LegalizerHelper Helper(*MF, Info);
+  ASSERT_TRUE(Helper.lower(*MIBCTTZ, 0, LLT::scalar(64)) ==
+              LegalizerHelper::LegalizeResult::Legalized);
+
+  auto CheckStr = R"(
+  CHECK: CTTZ
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+
+// CTLZ expansion in terms of CTLZ_ZERO_UNDEF
+TEST_F(LegalizerHelperTest, LowerBitCountingCTLZ0) {
+  if (!TM)
+    return;
+
+  // Declare your legalization info
+  DefineLegalizerInfo(
+      A, { getActionDefinitionsBuilder(G_CTLZ_ZERO_UNDEF).legalFor({s64}); });
+  // Build
+  auto MIBCTLZ = B.buildInstr(TargetOpcode::G_CTLZ, LLT::scalar(64), Copies[0]);
+  AInfo Info(MF->getSubtarget());
+  LegalizerHelper Helper(*MF, Info);
+  ASSERT_TRUE(Helper.lower(*MIBCTLZ, 0, LLT::scalar(64)) ==
+              LegalizerHelper::LegalizeResult::Legalized);
+
+  auto CheckStr = R"(
+  CHECK: [[CZU:%[0-9]+]]:_(s64) = G_CTLZ_ZERO_UNDEF %0
+  CHECK: [[ZERO:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+  CHECK: [[SIXTY4:%[0-9]+]]:_(s64) = G_CONSTANT i64 64
+  CHECK: [[CMP:%[0-9]+]]:_(s1) = G_ICMP intpred(eq), %0:_(s64), [[ZERO]]
+  CHECK: [[SEL:%[0-9]+]]:_(s64) = G_SELECT [[CMP]]:_(s1), [[SIXTY4]]:_, [[CZU]]
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+
+// CTLZ expansion
+TEST_F(LegalizerHelperTest, LowerBitCountingCTLZ1) {
+  if (!TM)
+    return;
+
+  // Declare your legalization info
+  DefineLegalizerInfo(A,
+                      { getActionDefinitionsBuilder(G_CTPOP).legalFor({s8}); });
+  // Build
+  // Trunc it to s8.
+  LLT s8{LLT::scalar(8)};
+  auto MIBTrunc = B.buildTrunc(s8, Copies[0]);
+  auto MIBCTLZ = B.buildInstr(TargetOpcode::G_CTLZ, s8, MIBTrunc);
+  AInfo Info(MF->getSubtarget());
+  LegalizerHelper Helper(*MF, Info);
+  ASSERT_TRUE(Helper.lower(*MIBCTLZ, 0, s8) ==
+              LegalizerHelper::LegalizeResult::Legalized);
+
+  auto CheckStr = R"(
+  CHECK: [[Trunc:%[0-9]+]]:_(s8) = G_TRUNC
+  CHECK: [[Cst1:%[0-9]+]]:_(s8) = G_CONSTANT i8 1
+  CHECK: [[Sh1:%[0-9]+]]:_(s8) = G_LSHR [[Trunc]]:_, [[Cst1]]:_
+  CHECK: [[Or1:%[0-9]+]]:_(s8) = G_OR [[Trunc]]:_, [[Sh1]]:_
+  CHECK: [[Cst2:%[0-9]+]]:_(s8) = G_CONSTANT i8 2
+  CHECK: [[Sh2:%[0-9]+]]:_(s8) = G_LSHR [[Or1]]:_, [[Cst2]]:_
+  CHECK: [[Or2:%[0-9]+]]:_(s8) = G_OR [[Or1]]:_, [[Sh2]]:_
+  CHECK: [[Cst4:%[0-9]+]]:_(s8) = G_CONSTANT i8 4
+  CHECK: [[Sh4:%[0-9]+]]:_(s8) = G_LSHR [[Or2]]:_, [[Cst4]]:_
+  CHECK: [[Or4:%[0-9]+]]:_(s8) = G_OR [[Or2]]:_, [[Sh4]]:_
+  CHECK: [[CTPOP:%[0-9]+]]:_(s8) = G_CTPOP [[Or4]]:_
+  CHECK: [[Len:%[0-9]+]]:_(s8) = G_CONSTANT i8 8
+  CHECK: [[Sub:%[0-9]+]]:_(s8) = G_SUB [[Len]]:_, [[CTPOP]]:_
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+} // namespace
diff --git a/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.h b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.h
new file mode 100644
index 00000000000..ca1aed544d2
--- /dev/null
+++ b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.h
@@ -0,0 +1,190 @@
+//===- LegalizerHelperTest.h
+//-----------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
+#include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
+#include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
+#include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
+#include "llvm/CodeGen/GlobalISel/Utils.h"
+#include "llvm/CodeGen/MIRParser/MIRParser.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/Support/FileCheck.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Target/TargetOptions.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace MIPatternMatch;
+
+void initLLVM() {
+  InitializeAllTargets();
+  InitializeAllTargetMCs();
+  InitializeAllAsmPrinters();
+  InitializeAllAsmParsers();
+
+  PassRegistry *Registry = PassRegistry::getPassRegistry();
+  initializeCore(*Registry);
+  initializeCodeGen(*Registry);
+}
+
+/// Create a TargetMachine. As we lack a dedicated always available target for
+/// unittests, we go for "AArch64".
+std::unique_ptr<TargetMachine> createTargetMachine() {
+  Triple TargetTriple("aarch64--");
+  std::string Error;
+  const Target *T = TargetRegistry::lookupTarget("", TargetTriple, Error);
+  if (!T)
+    return nullptr;
+
+  TargetOptions Options;
+  return std::unique_ptr<TargetMachine>(T->createTargetMachine(
+      "AArch64", "", "", Options, None, None, CodeGenOpt::Aggressive));
+}
+
+std::unique_ptr<Module> parseMIR(LLVMContext &Context,
+                                 std::unique_ptr<MIRParser> &MIR,
+                                 const TargetMachine &TM, StringRef MIRCode,
+                                 const char *FuncName, MachineModuleInfo &MMI) {
+  SMDiagnostic Diagnostic;
+  std::unique_ptr<MemoryBuffer> MBuffer = MemoryBuffer::getMemBuffer(MIRCode);
+  MIR = createMIRParser(std::move(MBuffer), Context);
+  if (!MIR)
+    return nullptr;
+
+  std::unique_ptr<Module> M = MIR->parseIRModule();
+  if (!M)
+    return nullptr;
+
+  M->setDataLayout(TM.createDataLayout());
+
+  if (MIR->parseMachineFunctions(*M, MMI))
+    return nullptr;
+
+  return M;
+}
+
+std::pair<std::unique_ptr<Module>, std::unique_ptr<MachineModuleInfo>>
+createDummyModule(LLVMContext &Context, const TargetMachine &TM,
+                  StringRef MIRFunc) {
+  SmallString<512> S;
+  StringRef MIRString = (Twine(R"MIR(
+---
+...
+name: func
+registers:
+  - { id: 0, class: _ }
+  - { id: 1, class: _ }
+  - { id: 2, class: _ }
+  - { id: 3, class: _ }
+body: |
+  bb.1:
+    %0(s64) = COPY $x0
+    %1(s64) = COPY $x1
+    %2(s64) = COPY $x2
+)MIR") + Twine(MIRFunc) + Twine("...\n"))
+                            .toNullTerminatedStringRef(S);
+  std::unique_ptr<MIRParser> MIR;
+  auto MMI = make_unique<MachineModuleInfo>(&TM);
+  std::unique_ptr<Module> M =
+      parseMIR(Context, MIR, TM, MIRString, "func", *MMI);
+  return make_pair(std::move(M), std::move(MMI));
+}
+
+static MachineFunction *getMFFromMMI(const Module *M,
+                                     const MachineModuleInfo *MMI) {
+  Function *F = M->getFunction("func");
+  auto *MF = MMI->getMachineFunction(*F);
+  return MF;
+}
+
+static void collectCopies(SmallVectorImpl<unsigned> &Copies,
+                          MachineFunction *MF) {
+  for (auto &MBB : *MF)
+    for (MachineInstr &MI : MBB) {
+      if (MI.getOpcode() == TargetOpcode::COPY)
+        Copies.push_back(MI.getOperand(0).getReg());
+    }
+}
+
+class LegalizerHelperTest : public ::testing::Test {
+protected:
+  LegalizerHelperTest() : ::testing::Test() {
+    TM = createTargetMachine();
+    if (!TM)
+      return;
+    ModuleMMIPair = createDummyModule(Context, *TM, "");
+    MF = getMFFromMMI(ModuleMMIPair.first.get(), ModuleMMIPair.second.get());
+    collectCopies(Copies, MF);
+    EntryMBB = &*MF->begin();
+    B.setMF(*MF);
+    MRI = &MF->getRegInfo();
+    B.setInsertPt(*EntryMBB, EntryMBB->end());
+  }
+  LLVMContext Context;
+  std::unique_ptr<TargetMachine> TM;
+  MachineFunction *MF;
+  std::pair<std::unique_ptr<Module>, std::unique_ptr<MachineModuleInfo>>
+      ModuleMMIPair;
+  SmallVector<unsigned, 4> Copies;
+  MachineBasicBlock *EntryMBB;
+  MachineIRBuilder B;
+  MachineRegisterInfo *MRI;
+};
+
+#define DefineLegalizerInfo(Name, SettingUpActionsBlock)                       \
+  class Name##Info : public LegalizerInfo {                                    \
+  public:                                                                      \
+    Name##Info(const TargetSubtargetInfo &ST) {                                \
+      using namespace TargetOpcode;                                            \
+      const LLT s8 = LLT::scalar(8);                                           \
+      (void)s8;                                                                \
+      const LLT s16 = LLT::scalar(16);                                         \
+      (void)s16;                                                               \
+      const LLT s32 = LLT::scalar(32);                                         \
+      (void)s32;                                                               \
+      const LLT s64 = LLT::scalar(64);                                         \
+      (void)s64;                                                               \
+      do                                                                       \
+        SettingUpActionsBlock while (0);                                       \
+      computeTables();                                                         \
+      verify(*ST.getInstrInfo());                                              \
+    }                                                                          \
+  };
+
+static bool CheckMachineFunction(const MachineFunction &MF,
+                                 StringRef CheckStr) {
+  SmallString<512> Msg;
+  raw_svector_ostream OS(Msg);
+  MF.print(OS);
+  auto OutputBuf = MemoryBuffer::getMemBuffer(Msg, "Output", false);
+  auto CheckBuf = MemoryBuffer::getMemBuffer(CheckStr, "");
+  SmallString<4096> CheckFileBuffer;
+  FileCheckRequest Req;
+  FileCheck FC(Req);
+  StringRef CheckFileText =
+      FC.CanonicalizeFile(*CheckBuf.get(), CheckFileBuffer);
+  SourceMgr SM;
+  SM.AddNewSourceBuffer(MemoryBuffer::getMemBuffer(CheckFileText, "CheckFile"),
+                        SMLoc());
+  Regex PrefixRE = FC.buildCheckPrefixRegex();
+  std::vector<FileCheckString> CheckStrings;
+  FC.ReadCheckFile(SM, CheckFileText, PrefixRE, CheckStrings);
+  auto OutBuffer = OutputBuf->getBuffer();
+  SM.AddNewSourceBuffer(std::move(OutputBuf), SMLoc());
+  return FC.CheckInput(SM, OutBuffer, CheckStrings);
+}