[GISel]: Add Legalization/lowering code for bit counting operations

[aditya_nandakumar]

Summary

Ported Legalization expansions for CTLZ/CTTZ.. from DAG to GISel.

[aditya_nandakumar]

Updated to make use of the FileCheck API for testing all the expansions.

[rtereshin]

Hi Aditya,

Thank you for doing this!

Most of my comments are nitpicks and mild suggestions except the following ones:

https://reviews.llvm.org/D48847#inline-446232
https://reviews.llvm.org/D48847#inline-446277
https://reviews.llvm.org/D48847#inline-446317
https://reviews.llvm.org/D48847#inline-446328

include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

671

I'm not sure what's the point of having variadic templates here. Does buildICmp or buildSelect make sense with 0, 1, or let's say 4 operands (excluding destination)?

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

40

👍

1063

Do you think we could MI.setDesc(TII.get(TargetOpcode:: G_CTLZ_ZERO_UNDEF)); here in-place instead of recreating the instruction from scratch?

1083

I don't think this is going to work if Len is not a power of 2.

For the sake of a smaller example, let's say Len is 6. The shift amount is going to take values 1 and 2. Let's say the SrcReg's value is 10 00 00. The Op takes values:
10 00 00
11 00 00 = 10 00 00 | 01 00 00, i == 0
11 11 00 = 11 00 00 | 00 11 00, i == 1

The final value is 2, while it had to be 0.

This could probably be fixed by rounding Len up to the closest power of 2.

1084

Maybe slightly more meaningful name?

1093

Same here, could we change MI in-place?

1101

I think there is a little opportunity here, not sure how useful in practice: if we check

if (!isLegalOrCustom({TargetOpcode::G_CTPOP, {Ty}}) &&
    isLegalOrCustom({TargetOpcode:: G_CTLZ_ZERO_UNDEF, {Ty}}))

here and fall-through if true we could lower this to 32 - ctlz_zero_undef(~x & (x-1)) with no extra selects.

Just a thought, it almost certainly doesn't worth the effort.

1129

We already have -1 constant built just above, maybe do G_ADD with it instead?

unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

32

missing %0 at the end?

139

ditto

185

Maybe it makes sense to check here that we actually thread things correctly.

unittests/CodeGen/GlobalISel/LegalizerHelperTest.h

74

Shouldn't we call MachineModuleInfo::doInitialization(Module &M) here?

Also, if we do, MMI will be accessible as MMI.getModule() which should make createDummyModules interface simpler.

95

Why not bb.0?

98

Do you want to make these live-in?

99

The last two Twine(...)s are probably not needed (I would expect Twine::operator+ to be overloaded for every string type)

101

Any reason this is created here and passed to parseMIR by reference instead of being created and destroyed all within parseMIR?

109

It looks like if MMI is initialized properly the Module *M will be accessible from MMI.

116

Maybe swap these so an out parameter is the last?

144

doesn't look like this is used by any of the tests at the moment.

146

Not used by any of the tests?

147

So we run multiple tests over the same instance of machine function, adding new instructions w/o clearing out the machine function and the rest of the context? Also, in some not really well defined order probably (like alphabetical by test name or something).

Shouldn't we put all of this into setUp instead and clear it out properly in tearDown and run every test in a clean state?

162

Maybe call x a Block or SettingUpActionsBlock or something along these lines, then

do Block while (0);

here and at the "call site" something like

LInfoBuilder(A, {
  getActionDefinitionsBuilder(G_CTTZ_ZERO_UNDEF).legalFor({s64});
});

instead of

LInfoBuilder(A,
               getActionDefinitionsBuilder(G_CTTZ_ZERO_UNDEF).legalFor({s64}););

?

As for the name, maybe DefineLegalizerInfo is a tad better than LInfoBuilder, not sure.

164

Maybe also call verify(*ST.getInstrInfo());?

188

Do we run MachineVerifier somewhere here?

[aditya_nandakumar]

Thanks roman for the feedback. I'll update the patch shortly based on the feedback.

include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

671

This required less typing one extra argument ;) . Additionally this lets you directly buildICmp like this
auto MIB = something...
auto Reg = someReg;
auto MIBCmp = Builder.buildICmp(Pred, SomeTy, MIB, Reg); This lets Op0 and Op1 be either MachineInstrBuilder or Register (for the various 4 combinations possible).
If you need to explicitly specify the use args, it needs to be something like this.

template<typename DstTy, typename UseArg0Ty, typename UseArg1Ty>
MachineInstrBuilder buildICmp(CmpInst::Predicate Pred, DstTy &&Dst,
                                                      UseArg0Ty &&Arg0, UseArg1Ty &&Arg1) {
     return buildICmp(Pred, getDestFromArg(Dst), getRegFromArg(Arg0), getRegFromArg(Arg1));
}

Which is a little more typing and longer.

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

1063

Definitely we should be doing that as often as we can, but in this case, if we go that route, we'd end up in a situation where the final select would need to create a new vreg and then it's probably not safe to just blindly replace the original dest with new reg. It's definitely possible to do it in a more ugly way where we replace the first inst in place, change it's destination to a newly created vreg but this sacrifices readability of code. I think CTLZ is not frequent enough to do this.

1083

Good catch. I'll replace it with the following.
NewLen = RoundUpToPow2(Len);
x = x | (x >>1);
... until NewLen/2
return Len - PopCount(x);

1101

That's something we could do. I can possibly address this in a subsequent patch.

1129

I've changed it. Thanks

unittests/CodeGen/GlobalISel/LegalizerHelperTest.h

74

This was copied over from another unit test. I'll refactor both of these in a subsequent patch.

116

Ditto for all of the above.

144

This is currently only used to set the insertion point. I left it available if any of the future tests needed that. I can remove it if required.

146

MRI was just left here in case any tests want to use it - but it's easily accessible with B.getMF().getRegInfo(). Again, I don't mind removing it.

147

I believe each googletest does not reuse the same test fixture object across multiple tests - it creates a new fixture object, constructs/calls setup, runs test, and call teardown()/destructor. I believe each test in it's own should be in a clean state when it runs.
I hope I've not misunderstood your question.

188

No we don't currently run the verifier here.

[aditya_nandakumar]

Updated patch based on feedback.

[rtereshin]

Thank you!

LGTM

unittests/CodeGen/GlobalISel/LegalizerHelperTest.h

147

True, (https://github.com/google/googletest/blob/master/googletest/docs/primer.md, https://github.com/google/googletest/blob/master/googletest/docs/faq.md#should-i-use-the-constructordestructor-of-the-test-fixture-or-the-set-uptear-down-function), maybe I got it mixed up with the python version of the library. Thanks!

188

Feels like a good idea to do that. Maybe, by default with the possibility to opt out on a test-case per test-case basis.

[aditya_nandakumar]

Thanks Roman.
Submitted in r340111

include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

enum LegalizeResult {
     Legalized,
 
     /// Some kind of error has occurred and we could not legalize this
     /// instruction.
     UnableToLegalize,
   };
 
   LegalizerHelper(MachineFunction &MF);
+  LegalizerHelper(MachineFunction &MF, const LegalizerInfo &LI);
 
   /// Replace \p MI by a sequence of legal instructions that can implement the
   /// same operation. Note that this means \p MI may be deleted, so any iterator
   /// steps should be performed before calling this function. \p Helper should
   /// be initialized to the MachineFunction containing \p MI.
   ///
   /// Considered as an opaque blob, the legal code will use and define the same
   /// registers as \p MI.
void widenScalarDst(MachineInstr &MI, LLT WideTy, unsigned OpIdx = 0,
                       unsigned TruncOpcode = TargetOpcode::G_TRUNC);
 
   /// Helper function to split a wide generic register into bitwise blocks with
   /// the given Type (which implies the number of blocks needed). The generic
   /// registers created are appended to Ops, starting at bit 0 of Reg.
   void extractParts(unsigned Reg, LLT Ty, int NumParts,
                     SmallVectorImpl<unsigned> &VRegs);
 
+  LegalizeResult lowerBitCount(MachineInstr &MI, unsigned TypeIdx, LLT Ty);
+
   MachineRegisterInfo &MRI;
   const LegalizerInfo &LI;
 };
 
 /// Helper function that creates the given libcall.
 LegalizerHelper::LegalizeResult
 createLibcall(MachineIRBuilder &MIRBuilder, RTLIB::Libcall Libcall,
               const CallLowering::ArgInfo &Result,
               ArrayRef<CallLowering::ArgInfo> Args);
 
 } // End namespace llvm.
 
 #endif

include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

public:
   /// \pre \p Op0 and Op1 must be generic virtual registers with the
   ///      same number of elements as \p Res. If \p Res is a scalar,
   ///      \p Op0 must be either a scalar or pointer.
   /// \pre \p Pred must be an integer predicate.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildICmp(CmpInst::Predicate Pred,
                                 unsigned Res, unsigned Op0, unsigned Op1);
+  template <typename DstTy, typename... UseArgsTy>
+  MachineInstrBuilder buildICmp(CmpInst::Predicate Pred, DstTy &&Dst,
+                                UseArgsTy &&... UseArgs) {
+    return buildICmp(Pred, getDestFromArg(Dst), getRegFromArg(UseArgs)...);
+  }
 
   /// Build and insert a \p Res = G_FCMP \p Pred\p Op0, \p Op1
   ///
   /// \pre setBasicBlock or setMI must have been called.
 
   /// \pre \p Res must be a generic virtual register with scalar or
   ///      vector type. Typically this starts as s1 or <N x s1>.
   /// \pre \p Op0 and Op1 must be generic virtual registers with the
public:
   ///      with the same type.
   /// \pre \p Tst must be a generic virtual register with scalar, pointer or
   ///      vector type. If vector then it must have the same number of
   ///      elements as the other parameters.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildSelect(unsigned Res, unsigned Tst,
                                   unsigned Op0, unsigned Op1);
+  template <typename DstTy, typename... UseArgsTy>
+  MachineInstrBuilder buildSelect(DstTy &&Dst, UseArgsTy &&... UseArgs) {
+    return buildSelect(getDestFromArg(Dst), getRegFromArg(UseArgs)...);
+  }
 
   /// Build and insert \p Res = G_INSERT_VECTOR_ELT \p Val,
   /// \p Elt, \p Idx
   ///
   /// \pre setBasicBlock or setMI must have been called.
   /// \pre \p Res and \p Val must be a generic virtual register
   //       with the same vector type.
   /// \pre \p Elt and \p Idx must be a generic virtual register

include/llvm/Target/GlobalISel/SelectionDAGCompat.td

 def : GINodeEquiv<G_FEXP2, fexp2>;
 def : GINodeEquiv<G_FLOG2, flog2>;
 def : GINodeEquiv<G_INTRINSIC, intrinsic_wo_chain>;
 // ISD::INTRINSIC_VOID can also be handled with G_INTRINSIC_W_SIDE_EFFECTS.
 def : GINodeEquiv<G_INTRINSIC_W_SIDE_EFFECTS, intrinsic_void>;
 def : GINodeEquiv<G_INTRINSIC_W_SIDE_EFFECTS, intrinsic_w_chain>;
 def : GINodeEquiv<G_BR, br>;
 def : GINodeEquiv<G_BSWAP, bswap>;
+def : GINodeEquiv<G_CTLZ, ctlz>;
+def : GINodeEquiv<G_CTTZ, cttz>;
+def : GINodeEquiv<G_CTLZ_ZERO_UNDEF, ctlz_zero_undef>;
+def : GINodeEquiv<G_CTTZ_ZERO_UNDEF, cttz_zero_undef>;
+def : GINodeEquiv<G_CTPOP, ctpop>;
 
 // Broadly speaking G_LOAD is equivalent to ISD::LOAD but there are some
 // complications that tablegen must take care of. For example, Predicates such
 // as isSignExtLoad require that this is not a perfect 1:1 mapping since a
 // sign-extending load is (G_SEXTLOAD x) in GlobalISel. Additionally,
 // G_LOAD handles both atomic and non-atomic loads where as SelectionDAG had
 // separate nodes for them. This GINodeEquiv maps the non-atomic loads to
 // G_LOAD with a non-atomic MachineMemOperand.

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

 /// primary legalization.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
 #include "llvm/CodeGen/GlobalISel/CallLowering.h"
 #include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 
 #define DEBUG_TYPE "legalizer"
 
 using namespace llvm;
 using namespace LegalizeActions;
 
 LegalizerHelper::LegalizerHelper(MachineFunction &MF)
     : MRI(MF.getRegInfo()), LI(*MF.getSubtarget().getLegalizerInfo()) {
   MIRBuilder.setMF(MF);
 }
 
+LegalizerHelper::LegalizerHelper(MachineFunction &MF, const LegalizerInfo &LI)
+    : MRI(MF.getRegInfo()), LI(LI) {
+  MIRBuilder.setMF(MF);
+}
 LegalizerHelper::LegalizeResult
 LegalizerHelper::legalizeInstrStep(MachineInstr &MI) {
   LLVM_DEBUG(dbgs() << "Legalizing: "; MI.print(dbgs()));
 
   auto Step = LI.getAction(MI, MRI);
   switch (Step.Action) {
   case Legal:
     LLVM_DEBUG(dbgs() << ".. Already legal\n");
if (DstTy.isScalar()) {
         break;
       }
       MI.eraseFromParent();
       return Legalized;
     }
 
     return UnableToLegalize;
   }
+  case TargetOpcode::G_CTLZ_ZERO_UNDEF:
+  case TargetOpcode::G_CTTZ_ZERO_UNDEF:
+  case TargetOpcode::G_CTLZ:
+  case TargetOpcode::G_CTTZ:
+  case TargetOpcode::G_CTPOP:
+    return lowerBitCount(MI, TypeIdx, Ty);
   }
 }
 
 LegalizerHelper::LegalizeResult
 LegalizerHelper::fewerElementsVector(MachineInstr &MI, unsigned TypeIdx,
                                      LLT NarrowTy) {
   // FIXME: Don't know how to handle secondary types yet.
   if (TypeIdx != 0)
case TargetOpcode::G_ADD: {
     }
 
     MIRBuilder.buildMerge(DstReg, DstRegs);
     MI.eraseFromParent();
     return Legalized;
   }
   }
 }
+
+LegalizerHelper::LegalizeResult
+LegalizerHelper::lowerBitCount(MachineInstr &MI, unsigned TypeIdx, LLT Ty) {
+  unsigned Opc = MI.getOpcode();
+  auto &TII = *MI.getMF()->getSubtarget().getInstrInfo();
+  auto isLegalOrCustom = [this](const LegalityQuery &Q) {
+    auto QAction = LI.getAction(Q).Action;
+    return QAction == Legal || QAction == Custom;
+  };
+  switch (Opc) {
+  default:
+    return UnableToLegalize;
+  case TargetOpcode::G_CTLZ_ZERO_UNDEF: {
+    // This trivially expands to CTLZ.
+    MI.setDesc(TII.get(TargetOpcode::G_CTLZ));
+    MIRBuilder.recordInsertion(&MI);
+    return Legalized;
+  }
+  case TargetOpcode::G_CTLZ: {
+    unsigned SrcReg = MI.getOperand(1).getReg();
+    unsigned Len = Ty.getSizeInBits();
+    if (isLegalOrCustom({TargetOpcode::G_CTLZ_ZERO_UNDEF, {Ty}})) {
+      // If CTLZ_ZERO_UNDEF is legal or custom, emit that and a select with
+      // zero.
+      auto MIBCtlzZU =
+          MIRBuilder.buildInstr(TargetOpcode::G_CTLZ_ZERO_UNDEF, Ty, SrcReg);
+      auto MIBZero = MIRBuilder.buildConstant(Ty, 0);
+      auto MIBLen = MIRBuilder.buildConstant(Ty, Len);
+      auto MIBICmp = MIRBuilder.buildICmp(CmpInst::ICMP_EQ, LLT::scalar(1),
+                                          SrcReg, MIBZero);
+      MIRBuilder.buildSelect(MI.getOperand(0).getReg(), MIBICmp, MIBLen,
+                             MIBCtlzZU);
+      MI.eraseFromParent();
+      return Legalized;
+    }
+    // for now, we do this:
+    // x = x | (x >> 1);
+    // x = x | (x >> 2);
+    // ...
+    // x = x | (x >>16);
+    // x = x | (x >>32); // for 64-bit input
+    // return popcount(~x);
+    //
+    // Ref: "Hacker's Delight" by Henry Warren
+    unsigned Op = SrcReg;
+    for (unsigned i = 0; (1U << i) <= (Len / 2); ++i) {
+      auto MIBTmp3 = MIRBuilder.buildConstant(Ty, 1ULL << i);
+      auto MIBOp = MIRBuilder.buildInstr(
+          TargetOpcode::G_OR, Ty, Op,
+          MIRBuilder.buildInstr(TargetOpcode::G_LSHR, Ty, Op, MIBTmp3));
+      Op = MIBOp->getOperand(0).getReg();
+    }
+    auto MIBNot = MIRBuilder.buildInstr(TargetOpcode::G_XOR, Ty, Op,
+                                        MIRBuilder.buildConstant(Ty, -1));
+    MIRBuilder.buildInstr(TargetOpcode::G_CTPOP, MI.getOperand(0).getReg(),
+                          MIBNot);
+    MI.eraseFromParent();
+    return Legalized;
+  }
+  case TargetOpcode::G_CTTZ_ZERO_UNDEF: {
+    // This trivially expands to CTTZ.
+    MI.setDesc(TII.get(TargetOpcode::G_CTTZ));
+    MIRBuilder.recordInsertion(&MI);
+    return Legalized;
+  }
+  case TargetOpcode::G_CTTZ: {
+    unsigned SrcReg = MI.getOperand(1).getReg();
+    unsigned Len = Ty.getSizeInBits();
+    if (isLegalOrCustom({TargetOpcode::G_CTTZ_ZERO_UNDEF, {Ty}})) {
+      // If CTTZ_ZERO_UNDEF is legal or custom, emit that and a select with
+      // zero.
+      auto MIBCttzZU =
+          MIRBuilder.buildInstr(TargetOpcode::G_CTTZ_ZERO_UNDEF, Ty, SrcReg);
+      auto MIBZero = MIRBuilder.buildConstant(Ty, 0);
+      auto MIBLen = MIRBuilder.buildConstant(Ty, Len);
+      auto MIBICmp = MIRBuilder.buildICmp(CmpInst::ICMP_EQ, LLT::scalar(1),
+                                          SrcReg, MIBZero);
+      MIRBuilder.buildSelect(MI.getOperand(0).getReg(), MIBICmp, MIBLen,
+                             MIBCttzZU);
+      MI.eraseFromParent();
+      return Legalized;
+    }
+    // for now, we use: { return popcount(~x & (x - 1)); }
+    // unless the target has ctlz but not ctpop, in which case we use:
+    // { return 32 - nlz(~x & (x-1)); }
+    // Ref: "Hacker's Delight" by Henry Warren
+    auto MIBNot = MIRBuilder.buildInstr(TargetOpcode::G_XOR, Ty, SrcReg,
+                                        MIRBuilder.buildConstant(Ty, -1));
+    auto MIBTmp = MIRBuilder.buildInstr(
+        TargetOpcode::G_AND, Ty, MIBNot,
+        MIRBuilder.buildInstr(TargetOpcode::G_SUB, Ty, SrcReg,
+                              MIRBuilder.buildConstant(Ty, 1)));
+    if (!isLegalOrCustom({TargetOpcode::G_CTPOP, {Ty}}) &&
+        isLegalOrCustom({TargetOpcode::G_CTLZ, {Ty}})) {
+      MIRBuilder.buildInstr(
+          TargetOpcode::G_SUB, MI.getOperand(0).getReg(),
+          MIRBuilder.buildConstant(Ty, Len),
+          MIRBuilder.buildInstr(TargetOpcode::G_CTLZ, Ty, MIBTmp));
+      MI.eraseFromParent();
+      return Legalized;
+    }
+    MIRBuilder.buildInstr(TargetOpcode::G_CTPOP, MI.getOperand(0).getReg(),
+                          MIBTmp);
+    MI.eraseFromParent();
+    return Legalized;
+  }
+  }
+}

unittests/CodeGen/GlobalISel/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   ${LLVM_TARGETS_TO_BUILD}
   CodeGen
   Core
   GlobalISel
   MC
   MIRParser
   Support
   Target
   )
 
 add_llvm_unittest(GlobalISelTests
         LegalizerInfoTest.cpp
         PatternMatchTest.cpp
+        LegalizerHelperTest.cpp
         )

unittests/CodeGen/GlobalISel/LegalizerHelperTest.h

+//===- 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 LInfoBuilder(Name, x)                                                  \
+  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;                                                               \
+      x;                                                                       \
+      computeTables();                                                         \
+    }                                                                          \
+  };
+
+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);
+}

unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

+//===- 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
+  LInfoBuilder(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
+  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
+  LInfoBuilder(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: [[ONE:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+  CHECK: [[SUB1:%[0-9]+]]:_(s64) = G_SUB %0:_, [[ONE]]
+  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
+  LInfoBuilder(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: [[ONE:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+  CHECK: [[SUB1:%[0-9]+]]:_(s64) = G_SUB %0:_, [[ONE]]
+  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
+  LInfoBuilder(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
+  LInfoBuilder(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
+  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
+  LInfoBuilder(A, getActionDefinitionsBuilder(G_CTPOP).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: G_CONSTANT i64 1
+  CHECK: G_LSHR
+  CHECK: G_OR
+  CHECK: G_CONSTANT i64 2
+  CHECK: G_LSHR
+  CHECK: G_OR
+  CHECK: G_CONSTANT i64 4
+  CHECK: G_LSHR
+  CHECK: G_OR
+  CHECK: G_CONSTANT i64 8
+  CHECK: G_LSHR
+  CHECK: G_OR
+  CHECK: G_CONSTANT i64 16
+  CHECK: G_LSHR
+  CHECK: G_OR
+  CHECK: G_CONSTANT i64 32
+  CHECK: G_LSHR
+  CHECK: G_OR
+  CHECK: G_CONSTANT i64 -1
+  CHECK: G_XOR
+  CHECK: G_CTPOP
+  )";
+
+  // Check
+  ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+}
+} // namespace