[offload][OMPT] Add device-specific tracing control (llvm#1564)

Author: dhruvachak

offload/include/OmptCommonDefs.h

@@ -92,7 +92,7 @@
   } while (0)
 
 typedef ompt_set_result_t (*libomptarget_ompt_set_trace_ompt_t)(
-    ompt_device_t *Device, unsigned int Enable, unsigned int EventTy);
+    int Device, unsigned int Enable, unsigned int EventTy);
 typedef int (*libomptarget_ompt_start_trace_t)(int,
                                                ompt_callback_buffer_request_t,
                                                ompt_callback_buffer_complete_t);

offload/include/OmptTracing.h

@@ -76,56 +76,66 @@ void setOmptHostToDeviceRate(double Slope, double Offset);
 /// Set / store the number of granted teams
 void setOmptGrantedNumTeams(uint64_t NumTeams);
 
-/// Activate / deactivate tracing
-void setTracingState(bool Enabled);
+/// Check if (1) tracing is globally active (2) the given device is actively
+/// traced and (3) the given event type is traced on the device
+bool isTracingEnabled(int DeviceId, unsigned int EventTy);
 
-/// Check if the given tracing type is monitored
-bool isTracingTypeEnabled(unsigned int EventTy);
+/// Check if the given device is actively traced
+bool isTracedDevice(int DeviceId);
 
-/// Set whether the given tracing type should be monitored (or not)
-void setTracingTypeEnabled(unsigned int EventTy, bool Enable);
+/// Check if the given device is monitoring the provided tracing type
+bool isTracingTypeEnabled(int DeviceId, unsigned int EventTy);
+
+/// Check if the given device is monitoring the provided tracing type 'group'
+/// Where group means we will check for both: EMI and non-EMI event types
+bool isTracingTypeGroupEnabled(int DeviceId, unsigned int EventTy);
+
+/// Set whether the given tracing type should be monitored (or not) on the
+/// device
+void setTracingTypeEnabled(uint64_t &TracedEventTy, bool Enable,
+                           unsigned int EventTy);
 
 /// Set / reset the given tracing types (EventTy = 0 corresponds to 'all')
-ompt_set_result_t setTraceEventTy(ompt_device_t *Device, unsigned int Enable,
+ompt_set_result_t setTraceEventTy(int DeviceId, unsigned int Enable,
                                   unsigned int EventTy);
 
 /// Return thread id
 uint64_t getThreadId();
 
-// Mutexes to serialize invocation of device-independent entry points
+/// Mutexes to serialize invocation of device registration and checks
+extern std::mutex DeviceAccessMutex;
+
+/// Mutexes to serialize invocation of device-independent entry points
 extern std::mutex TraceAccessMutex;
 extern std::mutex TraceControlMutex;
 
-// Ensure serialization of calls to std::hash
+/// Ensure serialization of calls to std::hash
 extern std::mutex TraceHashThreadMutex;
 
-// Protect map from device-id to the corresponding buffer-request and
-// buffer-completion callback functions.
+/// Protect map from device-id to the corresponding buffer-request and
+/// buffer-completion callback functions.
 extern std::mutex BufferManagementFnMutex;
 
-// Map from device-id to the corresponding buffer-request and buffer-completion
-// callback functions.
+/// Map from device-id to the corresponding buffer-request and buffer-completion
+/// callback functions.
 extern std::unordered_map<int, std::pair<ompt_callback_buffer_request_t,
                                          ompt_callback_buffer_complete_t>>
     BufferManagementFns;
 
-// Thread local variables used by the plugin to communicate OMPT information
-// that are then used to populate trace records. This method assumes a
-// synchronous implementation, otherwise it won't work.
+/// Thread local variables used by the plugin to communicate OMPT information
+/// that are then used to populate trace records. This method assumes a
+/// synchronous implementation, otherwise it won't work.
 extern thread_local uint32_t TraceRecordNumGrantedTeams;
 extern thread_local uint64_t TraceRecordStartTime;
 extern thread_local uint64_t TraceRecordStopTime;
 
-// Thread local thread-id.
+/// Thread local thread-id.
 extern thread_local uint64_t ThreadId;
 
-// Manage all tracing records in one place
+/// Manage all tracing records in one place.
 extern OmptTracingBufferMgr TraceRecordManager;
 
-// Keep track of enabled tracing event types
-extern std::atomic<uint64_t> TracingTypesEnabled;
-
-/// OMPT tracing status; (Re-)Set via 'setTracingState'
+/// OMPT global tracing status. Indicates if at least one device is traced.
 extern bool TracingActive;
 
 } // namespace ompt

offload/include/OpenMP/OMPT/Interface.h

@@ -15,13 +15,9 @@
 
 // Only provide functionality if target OMPT support is enabled
 #ifdef OMPT_SUPPORT
-#include <functional>
-#include <tuple>
-
-#include "Shared/Debug.h"
-
 #include "Callback.h"
 #include "Shared/APITypes.h"
+#include "Shared/Debug.h"
 #include "omp-tools.h"
 
 #include "llvm/Support/ErrorHandling.h"
@@ -52,8 +48,12 @@ extern ompt_get_task_data_t ompt_get_task_data_fn;
 extern ompt_get_target_task_data_t ompt_get_target_task_data_fn;
 extern ompt_set_frame_enter_t ompt_set_frame_enter_fn;
 
+/// OMPT global tracing status. Indicates if at least one device is traced.
 extern bool TracingActive;
 
+/// Check if this device traces the given event type
+extern bool isTracingEnabled(int DeviceId, unsigned int EventTy);
+
 /// Used to maintain execution state for this thread
 class Interface {
 public:
@@ -191,7 +191,6 @@ class Interface {
                                              unsigned int NumTeams = 1);
 
   ompt_record_ompt_t *stopTargetSubmitTraceAsync(ompt_record_ompt_t *DataPtr,
-                                                 int64_t DeviceId,
                                                  unsigned int NumTeams,
                                                  uint64_t NanosStart,
                                                  uint64_t NanosStop);
@@ -485,17 +484,17 @@ struct OmptEventInfoTy {
 /// extends the original with async capabilities. That is: It takes an
 /// additional AsyncInfo reference as argument to populate the relevant fields.
 /// The AsyncInfoTy propagates the info into the RTL / plugins.
+/// TracedDeviceId represents the trace record's device affinity. EventType is
+/// the callback type that needs to be enabled via ompt_set_trace_ompt.
 template <typename FunctionPairTy, typename AsyncInfoTy, typename... ArgsTy>
 class TracerInterfaceRAII {
 public:
   TracerInterfaceRAII(FunctionPairTy Callbacks, AsyncInfoTy &AsyncInfo,
+                      int TracedDeviceId, ompt_callbacks_t EventType,
                       ArgsTy... Args)
       : Arguments(Args...), beginFunction(std::get<0>(Callbacks)) {
     __tgt_async_info *AI = AsyncInfo;
-    if (!llvm::omp::target::ompt::TracingActive) {
-      assert(AI->OmptEventInfo == nullptr &&
-             "The OmptEventInfo was not nullptr");
-    } else {
+    if (isTracingEnabled(TracedDeviceId, EventType)) {
       auto Record = begin();
       // Gets freed in interface.cpp, functions
       // targetKernel and targetData once launching target operations returns.
@@ -505,6 +504,9 @@ class TracerInterfaceRAII {
       AI->OmptEventInfo->NumTeams = 0;
       AI->OmptEventInfo->RegionInterface = &RegionInterface;
       AI->OmptEventInfo->RIFunction = std::get<1>(Callbacks);
+    } else {
+      // Actively prevent further tracing of this event
+      AI->OmptEventInfo = nullptr;
     }
   }
 

offload/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -1913,7 +1913,7 @@ struct AMDGPUStreamTy {
            "Invalid RegionInterface pointer in OMPT profiling");
     assert(OmptEventInfo.TraceRecord && "Invalid TraceRecord");
     std::invoke(RIFunc, OmptEventInfo.RegionInterface,
-                OmptEventInfo.TraceRecord, 0, OmptEventInfo.NumTeams, StartTime,
+                OmptEventInfo.TraceRecord, OmptEventInfo.NumTeams, StartTime,
                 EndTime);
 
     return Plugin::success();

offload/plugins-nextgen/common/OMPT/OmptDeviceTracing.h

@@ -48,8 +48,11 @@ FOREACH_OMPT_DEVICE_TRACING_FN_IMPLEMENTAIONS(declareOmptTracingFnMutex)
 
 extern std::mutex DeviceIdWritingMutex;
 
-/// Activate / deactivate tracing
-void setTracingState(bool Enabled);
+/// Activate tracing on the given device
+void enableDeviceTracing(int DeviceId);
+
+/// Deactivate tracing on the given device
+void disableDeviceTracing(int DeviceId);
 
 /// Set 'start' and 'stop' in trace records
 void setOmptTimestamp(uint64_t StartTime, uint64_t EndTime);
@@ -69,6 +72,9 @@ void setDeviceId(ompt_device_t *Device, int32_t DeviceId);
 /// Rempve the given device pointer from the current mapping
 void removeDeviceId(ompt_device_t *Device);
 
+/// Check whether the provided device is currently traced.
+bool isTracedDevice(int32_t DeviceId);
+
 /// Provide name based lookup for the device tracing functions
 extern ompt_interface_fn_t
 lookupDeviceTracingFn(const char *InterfaceFunctionName);
@@ -82,10 +88,11 @@ extern double HostToDeviceOffset;
 /// Mapping of device pointers to their corresponding RTL device ID
 extern std::map<ompt_device_t *, int32_t> Devices;
 
-// Keep track of enabled tracing event types
-extern std::atomic<uint64_t> TracingTypesEnabled;
+/// Mapping of RTL device IDs to their currently enabled tracing event types.
+/// Note: Event type '0' (bit position) indicates if this device is traced.
+extern std::map<int32_t, uint64_t> TracedDevices;
 
-/// OMPT tracing status; (Re-)Set via 'setTracingState'
+/// OMPT global tracing status. Indicates if at least one device is traced.
 extern bool TracingActive;
 
 /// Parent library pointer

offload/plugins-nextgen/common/OMPT/OmptTracing.cpp

@@ -78,43 +78,61 @@ int llvm::omp::target::ompt::getDeviceId(ompt_device_t *Device) {
   std::unique_lock<std::mutex> Lock(DeviceIdWritingMutex);
   auto DeviceIterator = Devices.find(Device);
   if (Device == nullptr || DeviceIterator == Devices.end()) {
-    REPORT("Failed to get ID for device=%p\n", Device);
+    REPORT("Failed to get ID for Device=%p\n", Device);
     return -1;
   }
   return DeviceIterator->second;
 }
 
 void llvm::omp::target::ompt::setDeviceId(ompt_device_t *Device,
                                           int32_t DeviceId) {
-  assert(Device && "Mapping device id to nullptr is not allowed");
-  if (Device == nullptr) {
-    REPORT("Failed to set ID for nullptr device\n");
+  assert(Device && "Mapping device ID to nullptr is not allowed");
+  if (Device == nullptr || DeviceId < 0) {
+    REPORT("Failed to set ID=%d for Device=%p\n", DeviceId, Device);
     return;
   }
   std::unique_lock<std::mutex> Lock(DeviceIdWritingMutex);
+  auto DeviceIterator = Devices.find(Device);
+  if (DeviceIterator != Devices.end()) {
+    auto CurrentDeviceId = DeviceIterator->second;
+    if (DeviceId == CurrentDeviceId)
+      REPORT("Tried to duplicate OMPT Device=%p (ID=%d)\n", Device, DeviceId);
+    else
+      REPORT("Tried to overwrite OMPT Device=%p (ID=%d with new ID=%d)\n",
+             Device, CurrentDeviceId, DeviceId);
+    return;
+  }
   Devices.emplace(Device, DeviceId);
 }
 
 void llvm::omp::target::ompt::removeDeviceId(ompt_device_t *Device) {
-  if (Device == nullptr) {
-    REPORT("Failed to remove ID for nullptr device\n");
+  int DeviceId = getDeviceId(Device);
+  if (DeviceId < 0) {
+    REPORT("Failed to remove Device=%p (ID=%d)\n", Device, DeviceId);
     return;
   }
   std::unique_lock<std::mutex> Lock(DeviceIdWritingMutex);
   Devices.erase(Device);
+  TracedDevices.erase(DeviceId);
 }
 
 OMPT_API_ROUTINE ompt_set_result_t ompt_set_trace_ompt(ompt_device_t *Device,
                                                        unsigned int Enable,
                                                        unsigned int EventTy) {
   DP("Executing ompt_set_trace_ompt\n");
 
-  // TODO handle device
+  int DeviceId = getDeviceId(Device);
+  if (DeviceId < 0) {
+    REPORT("Failed to set trace events for Device=%p (Unknown device)\n",
+           Device);
+    return ompt_set_never;
+  }
+
   std::unique_lock<std::mutex> Lock(ompt_set_trace_ompt_mutex);
   ensureFuncPtrLoaded<libomptarget_ompt_set_trace_ompt_t>(
       "libomptarget_ompt_set_trace_ompt", &ompt_set_trace_ompt_fn);
   assert(ompt_set_trace_ompt_fn && "libomptarget_ompt_set_trace_ompt loaded");
-  return ompt_set_trace_ompt_fn(Device, Enable, EventTy);
+  return ompt_set_trace_ompt_fn(DeviceId, Enable, EventTy);
 }
 
 OMPT_API_ROUTINE int
@@ -123,12 +141,18 @@ ompt_start_trace(ompt_device_t *Device, ompt_callback_buffer_request_t Request,
   DP("Executing ompt_start_trace\n");
 
   int DeviceId = getDeviceId(Device);
+  if (DeviceId < 0) {
+    REPORT("Failed to start trace for Device=%p (Unknown device)\n", Device);
+    // Indicate failure
+    return 0;
+  }
+
   {
     // Protect the function pointer
     std::unique_lock<std::mutex> Lock(ompt_start_trace_mutex);
 
     if (Request && Complete) {
-      llvm::omp::target::ompt::setTracingState(/*Enabled=*/true);
+      llvm::omp::target::ompt::enableDeviceTracing(DeviceId);
       // Enable asynchronous memory copy profiling
       setOmptAsyncCopyProfile(/*Enable=*/true);
       // Enable queue dispatch profiling
@@ -150,7 +174,6 @@ ompt_start_trace(ompt_device_t *Device, ompt_callback_buffer_request_t Request,
 OMPT_API_ROUTINE int ompt_flush_trace(ompt_device_t *Device) {
   DP("Executing ompt_flush_trace\n");
 
-  // TODO handle device
   std::unique_lock<std::mutex> Lock(ompt_flush_trace_mutex);
   ensureFuncPtrLoaded<libomptarget_ompt_flush_trace_t>(
       "libomptarget_ompt_flush_trace", &ompt_flush_trace_fn);
@@ -161,15 +184,20 @@ OMPT_API_ROUTINE int ompt_flush_trace(ompt_device_t *Device) {
 OMPT_API_ROUTINE int ompt_stop_trace(ompt_device_t *Device) {
   DP("Executing ompt_stop_trace\n");
 
-  // TODO handle device
+  int DeviceId = getDeviceId(Device);
+  if (DeviceId < 0) {
+    REPORT("Failed to stop trace for Device=%p (Unknown device)\n", Device);
+    // Indicate failure
+    return 0;
+  }
+
   {
     // Protect the function pointer
     std::unique_lock<std::mutex> Lock(ompt_stop_trace_mutex);
-    llvm::omp::target::ompt::setTracingState(/*Enabled=*/false);
+    llvm::omp::target::ompt::disableDeviceTracing(DeviceId);
     // Disable asynchronous memory copy profiling
     setOmptAsyncCopyProfile(/*Enable=*/false);
     // Disable queue dispatch profiling
-    int DeviceId = getDeviceId(Device);
     if (DeviceId >= 0)
       setGlobalOmptKernelProfile(Device, /*Enable=*/0);
     else
@@ -179,7 +207,7 @@ OMPT_API_ROUTINE int ompt_stop_trace(ompt_device_t *Device) {
         "libomptarget_ompt_stop_trace", &ompt_stop_trace_fn);
     assert(ompt_stop_trace_fn && "libomptarget_ompt_stop_trace loaded");
   }
-  return ompt_stop_trace_fn(getDeviceId(Device));
+  return ompt_stop_trace_fn(DeviceId);
 }
 
 OMPT_API_ROUTINE ompt_record_ompt_t *

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -675,10 +675,15 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
           printLaunchInfo(GenericDevice, KernelArgs, NumThreads, NumBlocks))
     return Err;
 
-  OMPT_IF_TRACING_ENABLED(setOmptGrantedNumTeams(NumBlocks);
-                          // Set number of granted teams for OMPT
-                          __tgt_async_info *AI = AsyncInfoWrapper;
-                          AI->OmptEventInfo->NumTeams = NumBlocks;);
+  OMPT_IF_TRACING_ENABLED(if (llvm::omp::target::ompt::isTracedDevice(
+                                  getDeviceId(&GenericDevice))) {
+    __tgt_async_info *AI = AsyncInfoWrapper;
+    if (AI->OmptEventInfo != nullptr) {
+      // Set number of granted teams for OMPT
+      setOmptGrantedNumTeams(NumBlocks);
+      AI->OmptEventInfo->NumTeams = NumBlocks;
+    }
+  });
 
   return launchImpl(GenericDevice, NumThreads, NumBlocks, KernelArgs,
                     LaunchParams, AsyncInfoWrapper);
@@ -871,7 +876,7 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
 
 #ifdef OMPT_SUPPORT
   auto DevicePtr = reinterpret_cast<ompt_device_t *>(this);
-  ompt::setDeviceId(DevicePtr, DeviceId);
+  ompt::setDeviceId(DevicePtr, Plugin.getUserId(DeviceId));
   if (ompt::CallbacksInitialized) {
     bool ExpectedStatus = false;
     if (OmptInitialized.compare_exchange_strong(ExpectedStatus, true))