[BOLT] Refactor heatmap to be standalone tool

Separate heatmap from bolt and build it as standalone tool.

Reviewed By: maksfb

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

5 files changed, 115 insertions, 42 deletions

diff --git a/bolt/tools/CMakeLists.txt b/bolt/tools/CMakeLists.txt
index 2153a9b720b0..bd5a1d17af4c 100644
--- a/bolt/tools/CMakeLists.txt
+++ b/bolt/tools/CMakeLists.txt
@@ -1,2 +1,3 @@
 add_subdirectory(driver)
 add_subdirectory(merge-fdata)
+add_subdirectory(heatmap)
diff --git a/bolt/tools/driver/CMakeLists.txt b/bolt/tools/driver/CMakeLists.txt
index c8027197b056..d17083e7f41e 100644
--- a/bolt/tools/driver/CMakeLists.txt
+++ b/bolt/tools/driver/CMakeLists.txt
@@ -23,13 +23,11 @@ add_llvm_tool(llvm-bolt
 
 add_llvm_tool_symlink(perf2bolt llvm-bolt)
 add_llvm_tool_symlink(llvm-boltdiff llvm-bolt)
-add_llvm_tool_symlink(llvm-bolt-heatmap llvm-bolt)
 
 set(BOLT_DEPENDS
   llvm-bolt
   perf2bolt
   llvm-boltdiff
-  llvm-bolt-heatmap
   )
 
 add_custom_target(bolt DEPENDS ${BOLT_DEPENDS})
@@ -37,7 +35,6 @@ install(PROGRAMS
   ${CMAKE_BINARY_DIR}/bin/llvm-bolt
   ${CMAKE_BINARY_DIR}/bin/perf2bolt
   ${CMAKE_BINARY_DIR}/bin/llvm-boltdiff
-  ${CMAKE_BINARY_DIR}/bin/llvm-bolt-heatmap
   TYPE BIN
   COMPONENT bolt
   )
diff --git a/bolt/tools/driver/llvm-bolt.cpp b/bolt/tools/driver/llvm-bolt.cpp
index 717f2762809a..f21dd46d0911 100644
--- a/bolt/tools/driver/llvm-bolt.cpp
+++ b/bolt/tools/driver/llvm-bolt.cpp
@@ -125,34 +125,6 @@ void perf2boltMode(int argc, char **argv) {
   opts::AggregateOnly = true;
 }
 
-void heatmapMode(int argc, char **argv) {
-  // Insert a fake subcommand if invoked via a command alias.
-  std::unique_ptr<char *[]> FakeArgv;
-  if (argc == 1 || strcmp(argv[1], "heatmap")) {
-    ++argc;
-    FakeArgv.reset(new char *[argc + 1]);
-    FakeArgv[0] = argv[0];
-    FakeArgv[1] = const_cast<char *>("heatmap");
-    for (int I = 2; I < argc; ++I)
-      FakeArgv[I] = argv[I - 1];
-    FakeArgv[argc] = nullptr;
-    argv = FakeArgv.get();
-  }
-
-  cl::ParseCommandLineOptions(argc, argv, "");
-
-  if (!sys::fs::exists(opts::InputFilename))
-    report_error(opts::InputFilename, errc::no_such_file_or_directory);
-
-  if (opts::PerfData.empty()) {
-    errs() << ToolName << ": expected -perfdata=<filename> option.\n";
-    exit(1);
-  }
-
-  opts::HeatmapMode = true;
-  opts::AggregateOnly = true;
-}
-
 void boltDiffMode(int argc, char **argv) {
   cl::HideUnrelatedOptions(makeArrayRef(opts::BoltDiffCategories));
   cl::AddExtraVersionPrinter(printBoltRevision);
@@ -194,8 +166,8 @@ void boltMode(int argc, char **argv) {
   }
 }
 
-std::string GetExecutablePath(const char *Argv0) {
-  SmallString<128> ExecutablePath(Argv0);
+static std::string GetExecutablePath(const char *Argv0) {
+  SmallString<256> ExecutablePath(Argv0);
   // Do a PATH lookup if Argv0 isn't a valid path.
   if (!llvm::sys::fs::exists(ExecutablePath))
     if (llvm::ErrorOr<std::string> P =
@@ -224,19 +196,10 @@ int main(int argc, char **argv) {
 
   ToolName = argv[0];
 
-  // Pre-process subcommands.
-  if (argc > 1 && *argv[1] != '-') {
-    if (!strcmp(argv[1], "heatmap"))
-      opts::HeatmapMode = true;
-  }
-
   if (llvm::sys::path::filename(ToolName) == "perf2bolt")
     perf2boltMode(argc, argv);
   else if (llvm::sys::path::filename(ToolName) == "llvm-boltdiff")
     boltDiffMode(argc, argv);
-  else if (llvm::sys::path::filename(ToolName) == "llvm-bolt-heatmap" ||
-           opts::HeatmapMode)
-    heatmapMode(argc, argv);
   else
     boltMode(argc, argv);
 
diff --git a/bolt/tools/heatmap/CMakeLists.txt b/bolt/tools/heatmap/CMakeLists.txt
new file mode 100644
index 000000000000..2a04704683f6
--- /dev/null
+++ b/bolt/tools/heatmap/CMakeLists.txt
@@ -0,0 +1,15 @@
+set(LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  BOLTProfile
+  BOLTRewrite
+  BOLTUtils
+  MC
+  Object
+  Support
+  )
+
+add_llvm_tool(llvm-bolt-heatmap
+  heatmap.cpp
+  )
+
+set_target_properties(llvm-bolt-heatmap PROPERTIES FOLDER "BOLT")
diff --git a/bolt/tools/heatmap/heatmap.cpp b/bolt/tools/heatmap/heatmap.cpp
new file mode 100644
index 000000000000..93f03808b5c3
--- /dev/null
+++ b/bolt/tools/heatmap/heatmap.cpp
@@ -0,0 +1,97 @@
+#include "bolt/Profile/DataAggregator.h"
+#include "bolt/Rewrite/RewriteInstance.h"
+#include "bolt/Utils/CommandLineOpts.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Object/Binary.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/TargetSelect.h"
+
+using namespace llvm;
+using namespace bolt;
+
+namespace opts {
+
+static cl::OptionCategory *HeatmapCategories[] = {&HeatmapCategory,
+                                                  &BoltOutputCategory};
+
+static cl::opt<std::string> InputFilename(cl::Positional,
+                                          cl::desc("<executable>"),
+                                          cl::Required,
+                                          cl::cat(HeatmapCategory));
+
+} // namespace opts
+
+static StringRef ToolName;
+
+static void report_error(StringRef Message, std::error_code EC) {
+  assert(EC);
+  errs() << ToolName << ": '" << Message << "': " << EC.message() << ".\n";
+  exit(1);
+}
+
+static void report_error(StringRef Message, Error E) {
+  assert(E);
+  errs() << ToolName << ": '" << Message << "': " << toString(std::move(E))
+         << ".\n";
+  exit(1);
+}
+
+static std::string GetExecutablePath(const char *Argv0) {
+  SmallString<256> ExecutablePath(Argv0);
+  // Do a PATH lookup if Argv0 isn't a valid path.
+  if (!llvm::sys::fs::exists(ExecutablePath))
+    if (llvm::ErrorOr<std::string> P =
+            llvm::sys::findProgramByName(ExecutablePath))
+      ExecutablePath = *P;
+  return std::string(ExecutablePath.str());
+}
+
+int main(int argc, char **argv) {
+  cl::HideUnrelatedOptions(makeArrayRef(opts::HeatmapCategories));
+  cl::ParseCommandLineOptions(argc, argv, "");
+
+  if (opts::PerfData.empty()) {
+    errs() << ToolName << ": expected -perfdata=<filename> option.\n";
+    exit(1);
+  }
+
+  opts::HeatmapMode = true;
+  opts::AggregateOnly = true;
+  if (!sys::fs::exists(opts::InputFilename))
+    report_error(opts::InputFilename, errc::no_such_file_or_directory);
+
+  // Output to stdout by default
+  if (opts::OutputFilename.empty())
+    opts::OutputFilename = "-";
+
+  // Initialize targets and assembly printers/parsers.
+  llvm::InitializeAllTargetInfos();
+  llvm::InitializeAllTargetMCs();
+  llvm::InitializeAllAsmParsers();
+  llvm::InitializeAllDisassemblers();
+
+  llvm::InitializeAllTargets();
+  llvm::InitializeAllAsmPrinters();
+
+  ToolName = argv[0];
+  std::string ToolPath = GetExecutablePath(argv[0]);
+  Expected<OwningBinary<Binary>> BinaryOrErr =
+      createBinary(opts::InputFilename);
+  if (Error E = BinaryOrErr.takeError())
+    report_error(opts::InputFilename, std::move(E));
+  Binary &Binary = *BinaryOrErr.get().getBinary();
+
+  if (auto *e = dyn_cast<ELFObjectFileBase>(&Binary)) {
+    RewriteInstance RI(e, argc, argv, ToolPath);
+    if (Error E = RI.setProfile(opts::PerfData))
+      report_error(opts::PerfData, std::move(E));
+
+    RI.run();
+  } else {
+    report_error(opts::InputFilename, object_error::invalid_file_type);
+  }
+
+  return EXIT_SUCCESS;
+}