Basic VTune support for JITLink

[pchintalapudi]

Summary

This patch introduces VTuneSupportPlugin as a JITLink equivalent to IntelJITEventListener to report JIT compiled methods. Integration of debug info and actual registration will be added in following revisions.

Depends on D146391

[lhames]

llvm/include/llvm/ExecutionEngine/Orc/VTuneSupportPlugin.h

13–14

Include guards should match the file name.

llvm/lib/ExecutionEngine/Orc/VTuneSupportPlugin.cpp

86

It looks like you're deregistering in notifyRemovingResources below, but you sould add the deregistration call as a deallocation action here (replacing the no-op (}) if possible. Deallocation actions reduce IPC traffic, and they can be run during executor teardown even if the connection to the JIT has dropped out. It also saves you the hassle of tracking resources in your plugin (it becomes the memory manager's problem).

149–161

You can use the lookupAndRecordAddrs utility in llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h to simplify this to:

auto &ES = EPC.getExecutionSession();
ExecutorAddr RegisterImplAddr;
ExecutorAddr UnregisterImplAddr;
if (auto Err = lookupAndRecordAddrs(
      ES, LookupKind::Static, makeJITDylibSearchOrder(&PJD),
      {
        {ES.intern(RegisterVTuneImplName), &RegisterImplAddr},
        {ES.intern(UnregisterVTuneImplName), &UnregisterImplAddr},
      }))
  return Err;
return std::make_unique<VTuneSupportPlugin>(
    EPC, RegisterImplAddr, UnregisterImplAddr, EmitDebugInfo);

llvm/include/llvm/ExecutionEngine/Orc/VTuneSupportPlugin.h

+//===--- VTuneSupportPlugin.h -- Support for VTune profiler ---*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Handles support for registering code with VIntel Tune's Amplifier JIT API.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_AMPLIFIERSUPPORTPLUGIN_H
+#define LLVM_EXECUTIONENGINE_ORC_AMPLIFIERSUPPORTPLUGIN_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+
+#include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
+
+namespace llvm {
+
+namespace orc {
+
+class VTuneSupportPlugin : public ObjectLinkingLayer::Plugin {
+public:
+  VTuneSupportPlugin(ExecutorProcessControl &EPC, ExecutorAddr RegisterImplAddr,
+                     ExecutorAddr UnregisterImplAddr, bool EmitDebugInfo)
+      : EPC(EPC), RegisterVTuneImplAddr(RegisterImplAddr),
+        UnregisterVTuneImplAddr(UnregisterImplAddr),
+        EmitDebugInfo(EmitDebugInfo) {}
+
+  void modifyPassConfig(MaterializationResponsibility &MR,
+                        jitlink::LinkGraph &G,
+                        jitlink::PassConfiguration &Config) override;
+
+  Error notifyEmitted(MaterializationResponsibility &MR) override;
+  Error notifyFailed(MaterializationResponsibility &MR) override;
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override;
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                   ResourceKey SrcKey) override;
+
+  static Expected<std::unique_ptr<VTuneSupportPlugin>>
+  Create(ExecutorProcessControl &EPC, JITDylib &JD, bool EmitDebugInfo);
+
+private:
+  ExecutorProcessControl &EPC;
+  ExecutorAddr RegisterVTuneImplAddr;
+  ExecutorAddr UnregisterVTuneImplAddr;
+  std::mutex PluginMutex;
+  uint64_t NextMethodID{0};
+  DenseMap<MaterializationResponsibility *, std::pair<uint64_t, uint64_t>>
+      PendingMethodIDs;
+  DenseMap<ResourceKey, std::vector<std::pair<uint64_t, uint64_t>>>
+      LoadedMethodIDs;
+  bool EmitDebugInfo;
+};
+
+typedef std::vector<std::pair<unsigned, unsigned>> VTuneLineTable;
+
+// SI = String Index, 1-indexed into the VTuneMethodBatch::Strings table.
+// SI == 0 means replace with nullptr.
+
+// MI = Method Index, 1-indexed into the VTuneMethodBatch::Methods table.
+// MI == 0 means this is a parent method and was not inlined.
+
+struct VTuneMethodInfo {
+  VTuneLineTable LineTable;
+  ExecutorAddr LoadAddr;
+  uint64_t LoadSize;
+  uint64_t MethodID;
+  uint32_t NameSI;
+  uint32_t ClassFileSI;
+  uint32_t SourceFileSI;
+  uint32_t ParentMI;
+};
+
+typedef std::vector<VTuneMethodInfo> VTuneMethodTable;
+typedef std::vector<std::string> VTuneStringTable;
+
+struct VTuneMethodBatch {
+  VTuneMethodTable Methods;
+  VTuneStringTable Strings;
+};
+
+typedef std::vector<std::pair<uint64_t, uint64_t>> VTuneUnloadedMethodIDs;
+
+namespace shared {
+
+using SPSVTuneLineTable = SPSSequence<SPSTuple<uint32_t, uint32_t>>;
+using SPSVTuneMethodInfo =
+    SPSTuple<SPSVTuneLineTable, SPSExecutorAddr, uint64_t, uint64_t, uint32_t,
+             uint32_t, uint32_t, uint32_t>;
+using SPSVTuneMethodTable = SPSSequence<SPSVTuneMethodInfo>;
+using SPSVTuneStringTable = SPSSequence<SPSString>;
+using SPSVTuneMethodBatch = SPSTuple<SPSVTuneMethodTable, SPSVTuneStringTable>;
+using SPSVTuneUnloadedMethodIDs = SPSSequence<SPSTuple<uint64_t, uint64_t>>;
+
+template <> class SPSSerializationTraits<SPSVTuneMethodInfo, VTuneMethodInfo> {
+public:
+  static size_t size(const VTuneMethodInfo &MI) {
+    return SPSVTuneMethodInfo::AsArgList::size(
+        MI.LineTable, MI.LoadAddr, MI.LoadSize, MI.MethodID, MI.NameSI,
+        MI.ClassFileSI, MI.SourceFileSI, MI.ParentMI);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, VTuneMethodInfo &MI) {
+    return SPSVTuneMethodInfo::AsArgList::deserialize(
+        IB, MI.LineTable, MI.LoadAddr, MI.LoadSize, MI.MethodID, MI.NameSI,
+        MI.ClassFileSI, MI.SourceFileSI, MI.ParentMI);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const VTuneMethodInfo &MI) {
+    return SPSVTuneMethodInfo::AsArgList::serialize(
+        OB, MI.LineTable, MI.LoadAddr, MI.LoadSize, MI.MethodID, MI.NameSI,
+        MI.ClassFileSI, MI.SourceFileSI, MI.ParentMI);
+  }
+};
+
+template <>
+class SPSSerializationTraits<SPSVTuneMethodBatch, VTuneMethodBatch> {
+public:
+  static size_t size(const VTuneMethodBatch &MB) {
+    return SPSVTuneMethodBatch::AsArgList::size(MB.Methods, MB.Strings);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, VTuneMethodBatch &MB) {
+    return SPSVTuneMethodBatch::AsArgList::deserialize(IB, MB.Methods,
+                                                       MB.Strings);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const VTuneMethodBatch &MB) {
+    return SPSVTuneMethodBatch::AsArgList::serialize(OB, MB.Methods,
+                                                     MB.Strings);
+  }
+};
+
+} // end namespace shared
+
+} // end namespace orc
+
+} // end namespace llvm
+
+#endif

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

add_llvm_component_library(LLVMOrcJIT
   PerfSupportPlugin.cpp
   RTDyldObjectLinkingLayer.cpp
   SimpleRemoteEPC.cpp
   Speculation.cpp
   SpeculateAnalyses.cpp
   ExecutorProcessControl.cpp
   TaskDispatch.cpp
   ThreadSafeModule.cpp
+  VTuneSupportPlugin.cpp
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
 
   DEPENDS
   intrinsics_gen
 
   LINK_LIBS
   ${LLVM_PTHREAD_LIB}

llvm/lib/ExecutionEngine/Orc/VTuneSupportPlugin.cpp

+//===--- VTuneSupportPlugin.cpp -- Support for VTune profiler --*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Handles support for registering code with VIntel Tune's Amplfiier JIT API.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/VTuneSupportPlugin.h"
+
+#include "llvm/ExecutionEngine/Orc/DebugInfoSupport.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+namespace {
+
+constexpr StringRef RegisterVTuneImplName = "llvm_orc_registerVTuneImpl";
+constexpr StringRef UnregisterVTuneImplName = "llvm_orc_unregisterVTuneImpl";
+
+static VTuneMethodBatch getMethodBatch(LinkGraph &G, bool EmitDebugInfo) {
+  VTuneMethodBatch Batch;
+  auto GetStringIdx = [Deduplicator = StringMap<uint32_t>(),
+                       &Batch](StringRef S) mutable {
+    auto I = Deduplicator.find(S);
+    if (I != Deduplicator.end()) {
+      return I->second;
+    }
+    Batch.Strings.push_back(S.str());
+    return Deduplicator[S] = Batch.Strings.size();
+  };
+  for (auto Sym : G.defined_symbols()) {
+    if (!Sym->isCallable()) {
+      continue;
+    }
+
+    Batch.Methods.push_back(VTuneMethodInfo());
+    auto &Method = Batch.Methods.back();
+    Method.MethodID = 0;
+    Method.ParentMI = 0;
+    Method.LoadAddr = Sym->getAddress();
+    Method.LoadSize = Sym->getSize();
+    Method.NameSI = GetStringIdx(Sym->getName());
+    Method.ClassFileSI = 0;
+    Method.SourceFileSI = 0;
+
+    if (!EmitDebugInfo) {
+      continue;
+    }
+
+    // TODO: Emit debug info
+  }
+  return Batch;
+}
+
+} // namespace
+
+void VTuneSupportPlugin::modifyPassConfig(MaterializationResponsibility &MR,
+                                          LinkGraph &G,
+                                          PassConfiguration &Config) {
+  Config.PostFixupPasses.push_back([this, MR = &MR](LinkGraph &G) {
+    auto Batch = getMethodBatch(G, EmitDebugInfo);
+    if (Batch.Methods.empty()) {
+      // Nothing to do.
+      return Error::success();
+    }
+    {
+      std::lock_guard<std::mutex> Lock(PluginMutex);
+      uint64_t Allocated = Batch.Methods.size();
+      uint64_t Start = NextMethodID;
+      NextMethodID += Allocated;
+      for (auto &M : Batch.Methods) {
+        M.MethodID = Start++;
+      }
+      this->PendingMethodIDs[MR] = {Start, Allocated};
+    }
+    G.allocActions().push_back(
+        {cantFail(shared::WrapperFunctionCall::Create<
+                  shared::SPSArgList<shared::SPSVTuneMethodBatch>>(
+             RegisterVTuneImplAddr, Batch)),
+         {}});
+    return Error::success();
+  });
+}
+
+Error VTuneSupportPlugin::notifyEmitted(MaterializationResponsibility &MR) {
+  if (auto Err = MR.withResourceKeyDo([this, MR = &MR](ResourceKey K) {
+        std::lock_guard<std::mutex> Lock(PluginMutex);
+        auto I = PendingMethodIDs.find(MR);
+        if (I == PendingMethodIDs.end()) {
+          return;
+        }
+        LoadedMethodIDs[K].push_back(I->second);
+        PendingMethodIDs.erase(I);
+      })) {
+    return Err;
+  }
+  return Error::success();
+}
+
+Error VTuneSupportPlugin::notifyFailed(MaterializationResponsibility &MR) {
+  std::lock_guard<std::mutex> Lock(PluginMutex);
+  PendingMethodIDs.erase(&MR);
+  return Error::success();
+}
+
+Error VTuneSupportPlugin::notifyRemovingResources(JITDylib &JD, ResourceKey K) {
+  // Unregistration not required if not provided
+  if (!UnregisterVTuneImplAddr) {
+    return Error::success();
+  }
+  VTuneUnloadedMethodIDs UnloadedIDs;
+  {
+    std::lock_guard<std::mutex> Lock(PluginMutex);
+    auto I = LoadedMethodIDs.find(K);
+    if (I == LoadedMethodIDs.end()) {
+      return Error::success();
+    }
+    UnloadedIDs = std::move(I->second);
+    LoadedMethodIDs.erase(I);
+  }
+  if (auto Err = EPC.callSPSWrapper<void(shared::SPSVTuneUnloadedMethodIDs)>(
+          UnregisterVTuneImplAddr, UnloadedIDs)) {
+    return Err;
+  }
+  return Error::success();
+}
+
+void VTuneSupportPlugin::notifyTransferringResources(JITDylib &JD,
+                                                     ResourceKey DstKey,
+                                                     ResourceKey SrcKey) {
+  std::lock_guard<std::mutex> Lock(PluginMutex);
+  auto I = LoadedMethodIDs.find(SrcKey);
+  if (I == LoadedMethodIDs.end()) {
+    return;
+  }
+  auto &Dest = LoadedMethodIDs[DstKey];
+  Dest.insert(Dest.end(), I->second.begin(), I->second.end());
+}
+
+Expected<std::unique_ptr<VTuneSupportPlugin>>
+VTuneSupportPlugin::Create(ExecutorProcessControl &EPC, JITDylib &JD,
+                           bool EmitDebugInfo) {
+  auto &ES = EPC.getExecutionSession();
+  auto RegisterImplName = ES.intern(RegisterVTuneImplName);
+  auto UnregisterImplName = ES.intern(UnregisterVTuneImplName);
+  SymbolLookupSet SLS{RegisterImplName, UnregisterImplName};
+  auto Res = ES.lookup(makeJITDylibSearchOrder({&JD}), std::move(SLS));
+  if (!Res)
+    return Res.takeError();
+  ExecutorAddr RegisterImplAddr(
+      Res->find(RegisterImplName)->second.getAddress());
+  ExecutorAddr UnregisterImplAddr(
+      Res->find(UnregisterImplName)->second.getAddress());
+  return std::make_unique<VTuneSupportPlugin>(
+      EPC, RegisterImplAddr, UnregisterImplAddr, EmitDebugInfo);
+}