[MCA] Moving the target specific CustomBehaviour impl. from /tools/llvm-mca/ to /lib/Target/.

Differential Revision: https://reviews.llvm.org/D106775

llvm/CMakeLists.txt

@@ -718,6 +718,7 @@ set(LLVM_ENUM_TARGETS "")
 set(LLVM_ENUM_ASM_PRINTERS "")
 set(LLVM_ENUM_ASM_PARSERS "")
 set(LLVM_ENUM_DISASSEMBLERS "")
+set(LLVM_ENUM_TARGETMCAS "")
 foreach(t ${LLVM_TARGETS_TO_BUILD})
   set( td ${LLVM_MAIN_SRC_DIR}/lib/Target/${t} )
 
@@ -747,6 +748,10 @@ foreach(t ${LLVM_TARGETS_TO_BUILD})
     set(LLVM_ENUM_DISASSEMBLERS
       "${LLVM_ENUM_DISASSEMBLERS}LLVM_DISASSEMBLER(${t})\n")
   endif()
+    if( EXISTS ${td}/MCA/CMakeLists.txt )
+    set(LLVM_ENUM_TARGETMCAS
+      "${LLVM_ENUM_TARGETMCAS}LLVM_TARGETMCA(${t})\n")
+  endif()
 endforeach(t)
 
 # Produce the target definition files, which provide a way for clients to easily
@@ -767,6 +772,10 @@ configure_file(
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Config/Targets.def.in
   ${LLVM_INCLUDE_DIR}/llvm/Config/Targets.def
   )
+configure_file(
+  ${LLVM_MAIN_INCLUDE_DIR}/llvm/Config/TargetMCAs.def.in
+  ${LLVM_INCLUDE_DIR}/llvm/Config/TargetMCAs.def
+  )
 
 # They are not referenced. See set_output_directory().
 set( CMAKE_RUNTIME_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/bin )

llvm/cmake/modules/LLVM-Config.cmake

@@ -201,6 +201,13 @@ function(llvm_expand_pseudo_components out_components)
           list(APPEND expanded_components "${t}Info")
         endif()
       endforeach(t)
+    elseif( c STREQUAL "AllTargetsMCAs" )
+      # Link all the TargetMCAs from all the targets
+      foreach(t ${LLVM_TARGETS_TO_BUILD})
+        if( TARGET LLVM${t}TargetMCA )
+          list(APPEND expanded_components "${t}TargetMCA")
+        endif()
+      endforeach(t)
     else()
       list(APPEND expanded_components "${c}")
     endif()

llvm/docs/CommandGuide/llvm-mca.rst

@@ -217,8 +217,10 @@ option specifies "``-``", then the output will also be sent to standard output.
 
 .. option:: -disable-cb
 
-  Force usage of the generic CustomBehaviour class rather than using the target
-  specific class. The generic class never detects any custom hazards.
+  Force usage of the generic CustomBehaviour and InstrPostProcess classes rather
+  than using the target specific implementation. The generic classes never
+  detect any custom hazards or make any post processing modifications to
+  instructions.
 
 
 EXIT STATUS
@@ -1013,4 +1015,5 @@ if you don't know the exact number and a value of 0 represents no stall).
 
 If you'd like to add a CustomBehaviour class for a target that doesn't
 already have one, refer to an existing implementation to see how to set it
-up. Remember to look at (and add to) `/llvm-mca/lib/CMakeLists.txt`.
+up. The classes are implemented within the target specific backend (for
+example `/llvm/lib/Target/AMDGPU/MCA/`) so that they can access backend symbols.

llvm/include/llvm/Config/TargetMCAs.def.in

@@ -0,0 +1,29 @@
+/*===------ llvm/Config/TargetMCAs.def - LLVM Target MCAs -------*- C++ -*-===*\
+|*                                                                            *|
+|* 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                    *|
+|*                                                                            *|
+|*===----------------------------------------------------------------------===*|
+|*                                                                            *|
+|* This file enumerates all of the target MCAs                  *|
+|* supported by this build of LLVM. Clients of this file should define        *|
+|* the LLVM_TARGETMCA macro to be a function-like macro with a             *|
+|* single parameter (the name of the target whose assembly can be             *|
+|* generated); including this file will then enumerate all of the             *|
+|* targets with target MCAs.                                             *|
+|*                                                                            *|
+|* The set of targets supported by LLVM is generated at configuration         *|
+|* time, at which point this header is generated. Do not modify this          *|
+|* header directly.                                                           *|
+|*                                                                            *|
+\*===----------------------------------------------------------------------===*/
+
+#ifndef LLVM_TARGETMCA
+#  error Please define the macro LLVM_TARGETMCA(TargetName)
+#endif
+
+@LLVM_ENUM_TARGETMCAS@
+
+#undef LLVM_TARGETMCA

llvm/include/llvm/Config/llvm-config.h.cmake

@@ -50,6 +50,9 @@
 /* LLVM name for the native target MC init function, if available */
 #cmakedefine LLVM_NATIVE_TARGETMC LLVMInitialize${LLVM_NATIVE_ARCH}TargetMC
 
+/* LLVM name for the native target MCA init function, if available */
+#cmakedefine LLVM_NATIVE_TARGETMCA LLVMInitialize${LLVM_NATIVE_ARCH}TargetMCA
+
 /* Define if this is Unixish platform */
 #cmakedefine LLVM_ON_UNIX ${LLVM_ON_UNIX}
 

llvm/include/llvm/MCA/CustomBehaviour.h

@@ -55,11 +55,11 @@ public:
 class CustomBehaviour {
 protected:
   const MCSubtargetInfo &STI;
-  const SourceMgr &SrcMgr;
+  const mca::SourceMgr &SrcMgr;
   const MCInstrInfo &MCII;
 
 public:
-  CustomBehaviour(const MCSubtargetInfo &STI, const SourceMgr &SrcMgr,
+  CustomBehaviour(const MCSubtargetInfo &STI, const mca::SourceMgr &SrcMgr,
                   const MCInstrInfo &MCII)
       : STI(STI), SrcMgr(SrcMgr), MCII(MCII) {}
 

llvm/include/llvm/Support/TargetRegistry.h

@@ -59,6 +59,11 @@ class raw_ostream;
 class raw_pwrite_stream;
 class TargetMachine;
 class TargetOptions;
+namespace mca {
+class CustomBehaviour;
+class InstrPostProcess;
+class SourceMgr;
+} // namespace mca
 
 MCStreamer *createNullStreamer(MCContext &Ctx);
 // Takes ownership of \p TAB and \p CE.
@@ -114,6 +119,13 @@ MCSymbolizer *createMCSymbolizer(const Triple &TT, LLVMOpInfoCallback GetOpInfo,
                                  void *DisInfo, MCContext *Ctx,
                                  std::unique_ptr<MCRelocationInfo> &&RelInfo);
 
+mca::CustomBehaviour *createCustomBehaviour(const MCSubtargetInfo &STI,
+                                            const mca::SourceMgr &SrcMgr,
+                                            const MCInstrInfo &MCII);
+
+mca::InstrPostProcess *createInstrPostProcess(const MCSubtargetInfo &STI,
+                                              const MCInstrInfo &MCII);
+
 /// Target - Wrapper for Target specific information.
 ///
 /// For registration purposes, this is a POD type so that targets can be
@@ -206,6 +218,15 @@ public:
       LLVMSymbolLookupCallback SymbolLookUp, void *DisInfo, MCContext *Ctx,
       std::unique_ptr<MCRelocationInfo> &&RelInfo);
 
+  using CustomBehaviourCtorTy =
+      mca::CustomBehaviour *(*)(const MCSubtargetInfo &STI,
+                                const mca::SourceMgr &SrcMgr,
+                                const MCInstrInfo &MCII);
+
+  using InstrPostProcessCtorTy =
+      mca::InstrPostProcess *(*)(const MCSubtargetInfo &STI,
+                                 const MCInstrInfo &MCII);
+
 private:
   /// Next - The next registered target in the linked list, maintained by the
   /// TargetRegistry.
@@ -305,6 +326,14 @@ private:
   /// MCSymbolizer, if registered (default = llvm::createMCSymbolizer)
   MCSymbolizerCtorTy MCSymbolizerCtorFn = nullptr;
 
+  /// CustomBehaviourCtorFn - Construction function for this target's
+  /// CustomBehaviour, if registered (default = nullptr).
+  CustomBehaviourCtorTy CustomBehaviourCtorFn = nullptr;
+
+  /// InstrPostProcessCtorFn - Construction function for this target's
+  /// InstrPostProcess, if registered (default = nullptr).
+  InstrPostProcessCtorTy InstrPostProcessCtorFn = nullptr;
+
 public:
   Target() = default;
 
@@ -623,6 +652,25 @@ public:
               std::move(RelInfo));
   }
 
+  /// createCustomBehaviour - Create a target specific CustomBehaviour.
+  /// This class is used by llvm-mca and requires backend functionality.
+  mca::CustomBehaviour *createCustomBehaviour(const MCSubtargetInfo &STI,
+                                              const mca::SourceMgr &SrcMgr,
+                                              const MCInstrInfo &MCII) const {
+    if (CustomBehaviourCtorFn)
+      return CustomBehaviourCtorFn(STI, SrcMgr, MCII);
+    return nullptr;
+  }
+
+  /// createInstrPostProcess - Create a target specific InstrPostProcess.
+  /// This class is used by llvm-mca and requires backend functionality.
+  mca::InstrPostProcess *createInstrPostProcess(const MCSubtargetInfo &STI,
+                                                const MCInstrInfo &MCII) const {
+    if (InstrPostProcessCtorFn)
+      return InstrPostProcessCtorFn(STI, MCII);
+    return nullptr;
+  }
+
   /// @}
 };
 
@@ -959,6 +1007,34 @@ struct TargetRegistry {
     T.MCSymbolizerCtorFn = Fn;
   }
 
+  /// RegisterCustomBehaviour - Register a CustomBehaviour
+  /// implementation for the given target.
+  ///
+  /// Clients are responsible for ensuring that registration doesn't occur
+  /// while another thread is attempting to access the registry. Typically
+  /// this is done by initializing all targets at program startup.
+  ///
+  /// @param T - The target being registered.
+  /// @param Fn - A function to construct a CustomBehaviour for the target.
+  static void RegisterCustomBehaviour(Target &T,
+                                      Target::CustomBehaviourCtorTy Fn) {
+    T.CustomBehaviourCtorFn = Fn;
+  }
+
+  /// RegisterInstrPostProcess - Register an InstrPostProcess
+  /// implementation for the given target.
+  ///
+  /// Clients are responsible for ensuring that registration doesn't occur
+  /// while another thread is attempting to access the registry. Typically
+  /// this is done by initializing all targets at program startup.
+  ///
+  /// @param T - The target being registered.
+  /// @param Fn - A function to construct an InstrPostProcess for the target.
+  static void RegisterInstrPostProcess(Target &T,
+                                       Target::InstrPostProcessCtorTy Fn) {
+    T.InstrPostProcessCtorFn = Fn;
+  }
+
   /// @}
 };
 

llvm/include/llvm/Support/TargetSelect.h

@@ -41,6 +41,10 @@ extern "C" {
 #define LLVM_DISASSEMBLER(TargetName) \
   void LLVMInitialize##TargetName##Disassembler();
 #include "llvm/Config/Disassemblers.def"
+
+// Declare all of the available TargetMCA initialization functions.
+#define LLVM_TARGETMCA(TargetName) void LLVMInitialize##TargetName##TargetMCA();
+#include "llvm/Config/TargetMCAs.def"
 }
 
 namespace llvm {
@@ -159,6 +163,14 @@ namespace llvm {
     return true;
 #endif
   }
+
+  /// InitializeAllTargetMCAs - The main program should call
+  /// this function to initialize the target CustomBehaviour and
+  /// InstrPostProcess classes.
+  inline void InitializeAllTargetMCAs() {
+#define LLVM_TARGETMCA(TargetName) LLVMInitialize##TargetName##TargetMCA();
+#include "llvm/Config/TargetMCAs.def"
+  }
 }
 
 #endif

llvm/lib/Target/AMDGPU/CMakeLists.txt

@@ -175,6 +175,7 @@ add_llvm_target(AMDGPUCodeGen
 
 add_subdirectory(AsmParser)
 add_subdirectory(Disassembler)
+add_subdirectory(MCA)
 add_subdirectory(MCTargetDesc)
 add_subdirectory(TargetInfo)
 add_subdirectory(Utils)

llvm/lib/Target/AMDGPU/MCA/AMDGPUCustomBehaviour.cpp

@@ -0,0 +1,66 @@
+//===------------------ AMDGPUCustomBehaviour.cpp ---------------*-C++ -* -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file implements methods from the AMDGPUCustomBehaviour class.
+///
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUCustomBehaviour.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIInstrInfo.h"
+#include "TargetInfo/AMDGPUTargetInfo.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/WithColor.h"
+
+namespace llvm {
+namespace mca {
+
+AMDGPUCustomBehaviour::AMDGPUCustomBehaviour(const MCSubtargetInfo &STI,
+                                             const mca::SourceMgr &SrcMgr,
+                                             const MCInstrInfo &MCII)
+    : CustomBehaviour(STI, SrcMgr, MCII) {}
+
+unsigned
+AMDGPUCustomBehaviour::checkCustomHazard(ArrayRef<mca::InstRef> IssuedInst,
+                                         const mca::InstRef &IR) {
+  return 0;
+}
+
+} // namespace mca
+} // namespace llvm
+
+using namespace llvm;
+using namespace mca;
+
+static CustomBehaviour *
+createAMDGPUCustomBehaviour(const MCSubtargetInfo &STI,
+                            const mca::SourceMgr &SrcMgr,
+                            const MCInstrInfo &MCII) {
+  return new AMDGPUCustomBehaviour(STI, SrcMgr, MCII);
+}
+
+static InstrPostProcess *
+createAMDGPUInstrPostProcess(const MCSubtargetInfo &STI,
+                             const MCInstrInfo &MCII) {
+  return new AMDGPUInstrPostProcess(STI, MCII);
+}
+
+/// Extern function to initialize the targets for the AMDGPU backend
+
+extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTargetMCA() {
+  TargetRegistry::RegisterCustomBehaviour(getTheAMDGPUTarget(),
+                                          createAMDGPUCustomBehaviour);
+  TargetRegistry::RegisterInstrPostProcess(getTheAMDGPUTarget(),
+                                           createAMDGPUInstrPostProcess);
+
+  TargetRegistry::RegisterCustomBehaviour(getTheGCNTarget(),
+                                          createAMDGPUCustomBehaviour);
+  TargetRegistry::RegisterInstrPostProcess(getTheGCNTarget(),
+                                           createAMDGPUInstrPostProcess);
+}

llvm/lib/Target/AMDGPU/MCA/AMDGPUCustomBehaviour.h

@@ -8,12 +8,14 @@
 /// \file
 ///
 /// This file defines the AMDGPUCustomBehaviour class which inherits from
-/// CustomBehaviour.
+/// CustomBehaviour. This class is used by the tool llvm-mca to enforce
+/// target specific behaviour that is not expressed well enough in the
+/// scheduling model for mca to enforce it automatically.
 ///
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_TOOLS_LLVM_MCA_LIB_AMDGPU_AMDGPUCUSTOMBEHAVIOUR_H
-#define LLVM_TOOLS_LLVM_MCA_LIB_AMDGPU_AMDGPUCUSTOMBEHAVIOUR_H
+#ifndef LLVM_LIB_TARGET_AMDGPU_MCA_AMDGPUCUSTOMBEHAVIOUR_H
+#define LLVM_LIB_TARGET_AMDGPU_MCA_AMDGPUCUSTOMBEHAVIOUR_H
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/MCA/CustomBehaviour.h"
@@ -29,14 +31,14 @@ public:
 
   ~AMDGPUInstrPostProcess() {}
 
-  void postProcessInstruction(std::unique_ptr<Instruction> &Inst,
+  void postProcessInstruction(std::unique_ptr<mca::Instruction> &Inst,
                               const MCInst &MCI) override {}
 };
 
 class AMDGPUCustomBehaviour : public CustomBehaviour {
 public:
-  AMDGPUCustomBehaviour(const MCSubtargetInfo &STI, const SourceMgr &SrcMgr,
-                        const MCInstrInfo &MCII);
+  AMDGPUCustomBehaviour(const MCSubtargetInfo &STI,
+                        const mca::SourceMgr &SrcMgr, const MCInstrInfo &MCII);
 
   ~AMDGPUCustomBehaviour() {}
 
@@ -47,11 +49,11 @@ public:
   /// register and hardware dependencies so this method should only
   /// implement custom behaviour and dependencies that are not picked up
   /// by MCA naturally.
-  unsigned checkCustomHazard(ArrayRef<InstRef> IssuedInst,
-                             const InstRef &IR) override;
+  unsigned checkCustomHazard(ArrayRef<mca::InstRef> IssuedInst,
+                             const mca::InstRef &IR) override;
 };
 
 } // namespace mca
 } // namespace llvm
 
-#endif /* LLVM_TOOLS_LLVM_MCA_LIB_AMDGPU_AMDGPUCUSTOMBEHAVIOUR_H */
+#endif

llvm/lib/Target/AMDGPU/MCA/CMakeLists.txt

@@ -0,0 +1,17 @@
+add_llvm_component_library(LLVMAMDGPUTargetMCA
+  AMDGPUCustomBehaviour.cpp
+
+  LINK_COMPONENTS
+  MC
+  MCParser
+  AMDGPUDesc
+  AMDGPUInfo
+  AMDGPUUtils
+  Support
+  MCA
+
+  ADD_TO_COMPONENT
+  AMDGPU
+  )
+
+add_dependencies(LLVMAMDGPUTargetMCA LLVMAMDGPUUtils)

llvm/tools/llvm-mca/CMakeLists.txt

@@ -1,9 +1,8 @@
 include_directories(include)
 
-add_subdirectory(lib)
-
 set(LLVM_LINK_COMPONENTS
   AllTargetsAsmParsers
+  AllTargetsMCAs          # CustomBehaviour and InstrPostProcess
   AllTargetsDescs
   AllTargetsDisassemblers
   AllTargetsInfos
@@ -32,9 +31,3 @@ add_llvm_tool(llvm-mca
   )
 
 set(LLVM_MCA_SOURCE_DIR ${CURRENT_SOURCE_DIR})
-
-target_link_libraries(llvm-mca PRIVATE
-  ${LLVM_MCA_CUSTOMBEHAVIOUR_TARGETS}
-  )
-
-add_definitions(${LLVM_MCA_MACROS_TO_DEFINE})

llvm/tools/llvm-mca/CodeRegionGenerator.h

@@ -37,7 +37,7 @@ protected:
   CodeRegionGenerator &operator=(const CodeRegionGenerator &) = delete;
 
 public:
-  CodeRegionGenerator(SourceMgr &SM) : Regions(SM) {}
+  CodeRegionGenerator(llvm::SourceMgr &SM) : Regions(SM) {}
   virtual ~CodeRegionGenerator();
   virtual Expected<const CodeRegions &>
   parseCodeRegions(const std::unique_ptr<MCInstPrinter> &IP) = 0;
@@ -54,7 +54,7 @@ class AsmCodeRegionGenerator final : public CodeRegionGenerator {
   unsigned AssemblerDialect; // This is set during parsing.
 
 public:
-  AsmCodeRegionGenerator(const Target &T, SourceMgr &SM, MCContext &C,
+  AsmCodeRegionGenerator(const Target &T, llvm::SourceMgr &SM, MCContext &C,
                          const MCAsmInfo &A, const MCSubtargetInfo &S,
                          const MCInstrInfo &I)
       : CodeRegionGenerator(SM), TheTarget(T), Ctx(C), MAI(A), STI(S), MCII(I),

llvm/tools/llvm-mca/lib/AMDGPU/AMDGPUCustomBehaviour.cpp

@@ -1,33 +0,0 @@
-//===------------------ AMDGPUCustomBehaviour.cpp ---------------*-C++ -* -===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-/// \file
-///
-/// This file implements methods from the AMDGPUCustomBehaviour class.
-///
-//===----------------------------------------------------------------------===//
-
-#include "AMDGPUCustomBehaviour.h"
-#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
-#include "SIInstrInfo.h"
-#include "llvm/Support/WithColor.h"
-
-namespace llvm {
-namespace mca {
-
-AMDGPUCustomBehaviour::AMDGPUCustomBehaviour(const MCSubtargetInfo &STI,
-                                             const SourceMgr &SrcMgr,
-                                             const MCInstrInfo &MCII)
-    : CustomBehaviour(STI, SrcMgr, MCII) {}
-
-unsigned AMDGPUCustomBehaviour::checkCustomHazard(ArrayRef<InstRef> IssuedInst,
-                                                  const InstRef &IR) {
-  return 0;
-}
-
-} // namespace mca
-} // namespace llvm

llvm/tools/llvm-mca/lib/AMDGPU/CMakeLists.txt

@@ -1,18 +0,0 @@
-include_directories(
-  ${LLVM_MAIN_SRC_DIR}/lib/Target/AMDGPU
-  ${LLVM_BINARY_DIR}/lib/Target/AMDGPU
-  )
-
-set(LLVM_LINK_COMPONENTS
-  AMDGPU
-  Core
-  MCA
-  Support
-  )
-
-add_llvm_library(LLVMMCACustomBehaviourAMDGPU
-  AMDGPUCustomBehaviour.cpp
-
-  DEPENDS
-  AMDGPUCommonTableGen
-  )

llvm/tools/llvm-mca/lib/CMakeLists.txt

@@ -1,11 +0,0 @@
-set(TARGETS_TO_APPEND "")
-set(MACROS_TO_APPEND "")
-
-if (LLVM_TARGETS_TO_BUILD MATCHES "AMDGPU")
-  add_subdirectory(AMDGPU)
-  list(APPEND TARGETS_TO_APPEND LLVMMCACustomBehaviourAMDGPU)
-  list(APPEND MACROS_TO_APPEND -DHAS_AMDGPU)
-endif()
-
-set(LLVM_MCA_CUSTOMBEHAVIOUR_TARGETS ${TARGETS_TO_APPEND} PARENT_SCOPE)
-set(LLVM_MCA_MACROS_TO_DEFINE ${MACROS_TO_APPEND} PARENT_SCOPE)

llvm/tools/llvm-mca/llvm-mca.cpp

@@ -32,9 +32,6 @@
 #include "Views/SchedulerStatistics.h"
 #include "Views/SummaryView.h"
 #include "Views/TimelineView.h"
-#ifdef HAS_AMDGPU
-#include "lib/AMDGPU/AMDGPUCustomBehaviour.h"
-#endif
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCCodeEmitter.h"
@@ -293,39 +290,6 @@ static void processViewOptions(bool IsOutOfOrder) {
     processOptionImpl(PrintRetireStats, Default);
 }
 
-std::unique_ptr<mca::InstrPostProcess>
-createInstrPostProcess(const Triple &TheTriple, const MCSubtargetInfo &STI,
-                       const MCInstrInfo &MCII) {
-  // Might be a good idea to have a separate flag so that InstrPostProcess
-  // can be used with or without CustomBehaviour
-  if (DisableCustomBehaviour)
-    return std::make_unique<mca::InstrPostProcess>(STI, MCII);
-#ifdef HAS_AMDGPU
-  if (TheTriple.isAMDGPU())
-    return std::make_unique<mca::AMDGPUInstrPostProcess>(STI, MCII);
-#endif
-  return std::make_unique<mca::InstrPostProcess>(STI, MCII);
-}
-
-std::unique_ptr<mca::CustomBehaviour>
-createCustomBehaviour(const Triple &TheTriple, const MCSubtargetInfo &STI,
-                      const mca::SourceMgr &SrcMgr, const MCInstrInfo &MCII) {
-  // Build the appropriate CustomBehaviour object for the current target.
-  // The CustomBehaviour class should never depend on the source code,
-  // but it can depend on the list of mca::Instruction and any classes
-  // that can be built using just the target info. If you need extra
-  // information from the source code or the list of MCInst, consider
-  // adding that information to the mca::Instruction class and setting
-  // it during InstrBuilder::createInstruction().
-  if (DisableCustomBehaviour)
-    return std::make_unique<mca::CustomBehaviour>(STI, SrcMgr, MCII);
-#ifdef HAS_AMDGPU
-  if (TheTriple.isAMDGPU())
-    return std::make_unique<mca::AMDGPUCustomBehaviour>(STI, SrcMgr, MCII);
-#endif
-  return std::make_unique<mca::CustomBehaviour>(STI, SrcMgr, MCII);
-}
-
 // Returns true on success.
 static bool runPipeline(mca::Pipeline &P) {
   // Handle pipeline errors here.
@@ -344,6 +308,7 @@ int main(int argc, char **argv) {
   InitializeAllTargetInfos();
   InitializeAllTargetMCs();
   InitializeAllAsmParsers();
+  InitializeAllTargetMCAs();
 
   // Enable printing of available targets when flag --version is specified.
   cl::AddExtraVersionPrinter(TargetRegistry::printRegisteredTargetsForVersion);
@@ -532,8 +497,18 @@ int main(int argc, char **argv) {
     // Lower the MCInst sequence into an mca::Instruction sequence.
     ArrayRef<MCInst> Insts = Region->getInstructions();
     mca::CodeEmitter CE(*STI, *MAB, *MCE, Insts);
-    std::unique_ptr<mca::InstrPostProcess> IPP =
-        createInstrPostProcess(TheTriple, *STI, *MCII);
+
+    std::unique_ptr<mca::InstrPostProcess> IPP;
+    if (!DisableCustomBehaviour) {
+      IPP = std::unique_ptr<mca::InstrPostProcess>(
+          TheTarget->createInstrPostProcess(*STI, *MCII));
+    }
+    if (!IPP)
+      // If the target doesn't have its own IPP implemented (or the
+      // -disable-cb flag is set) then we use the base class
+      // (which does nothing).
+      IPP = std::make_unique<mca::InstrPostProcess>(*STI, *MCII);
+
     std::vector<std::unique_ptr<mca::Instruction>> LoweredSequence;
     for (const MCInst &MCI : Insts) {
       Expected<std::unique_ptr<mca::Instruction>> Inst =
@@ -602,8 +577,14 @@ int main(int argc, char **argv) {
     // the source code (but it can depend on the list of
     // mca::Instruction or any objects that can be reconstructed
     // from the target information).
-    std::unique_ptr<mca::CustomBehaviour> CB =
-        createCustomBehaviour(TheTriple, *STI, S, *MCII);
+    std::unique_ptr<mca::CustomBehaviour> CB;
+    if (!DisableCustomBehaviour)
+      CB = std::unique_ptr<mca::CustomBehaviour>(
+          TheTarget->createCustomBehaviour(*STI, S, *MCII));
+    if (!CB)
+      // If the target doesn't have its own CB implemented (or the -disable-cb
+      // flag is set) then we use the base class (which does nothing).
+      CB = std::make_unique<mca::CustomBehaviour>(*STI, S, *MCII);
 
     // Create a basic pipeline simulating an out-of-order backend.
     auto P = MCA.createDefaultPipeline(PO, S, *CB);