[OpenMP] Add implementation to two OMPT API routines

This change adds implementation to ompt_finalize_tool() and
ompt_get_task_memory().

Patch by Hansang Bae

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


git-svn-id: https://llvm.org/svn/llvm-project/openmp/trunk@361309 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 182 insertions, 15 deletions

diff --git a/runtime/src/ompt-general.cpp b/runtime/src/ompt-general.cpp
index 55a96a4..00bf606 100644
--- a/runtime/src/ompt-general.cpp
+++ b/runtime/src/ompt-general.cpp
@@ -524,8 +524,7 @@ OMPT_API_ROUTINE int ompt_get_task_info(int ancestor_level, int *type,
 
 OMPT_API_ROUTINE int ompt_get_task_memory(void **addr, size_t *size,
                                           int block) {
-  // stub
-  return 0;
+  return __ompt_get_task_memory_internal(addr, size, block);
 }
 
 /*****************************************************************************
@@ -700,9 +699,7 @@ OMPT_API_ROUTINE uint64_t ompt_get_unique_id(void) {
   return __ompt_get_unique_id_internal();
 }
 
-OMPT_API_ROUTINE void ompt_finalize_tool(void) {
-  // stub
-}
+OMPT_API_ROUTINE void ompt_finalize_tool(void) { __kmp_internal_end_atexit(); }
 
 /*****************************************************************************
  * Target
diff --git a/runtime/src/ompt-specific.cpp b/runtime/src/ompt-specific.cpp
index c69cb46..99c4b61 100644
--- a/runtime/src/ompt-specific.cpp
+++ b/runtime/src/ompt-specific.cpp
@@ -427,6 +427,35 @@ int __ompt_get_task_info_internal(int ancestor_level, int *type,
   return 0;
 }
 
+int __ompt_get_task_memory_internal(void **addr, size_t *size, int blocknum) {
+  if (blocknum != 0)
+    return 0; // support only a single block
+
+  kmp_info_t *thr = ompt_get_thread();
+  kmp_taskdata_t *taskdata = thr->th.th_current_task;
+  kmp_task_t *task = KMP_TASKDATA_TO_TASK(taskdata);
+
+  if (taskdata->td_flags.tasktype != TASK_EXPLICIT)
+    return 0; // support only explicit task
+
+  void *ret_addr;
+  int64_t ret_size = taskdata->td_size_alloc - sizeof(kmp_taskdata_t);
+
+  // kmp_task_t->data1 is an optional member
+  if (taskdata->td_flags.destructors_thunk)
+    ret_addr = &task->data1 + 1;
+  else
+    ret_addr = &task->part_id + 1;
+
+  ret_size -= (char *)(ret_addr) - (char *)(task);
+  if (ret_size < 0)
+    return 0;
+
+  *addr = ret_addr;
+  *size = ret_size;
+  return 1;
+}
+
 //----------------------------------------------------------
 // team support
 //----------------------------------------------------------
diff --git a/runtime/test/ompt/callback.h b/runtime/test/ompt/callback.h
index c2e1530..dfb8a17 100755
--- a/runtime/test/ompt/callback.h
+++ b/runtime/test/ompt/callback.h
@@ -67,9 +67,11 @@ static ompt_set_callback_t ompt_set_callback;
 static ompt_get_callback_t ompt_get_callback;
 static ompt_get_state_t ompt_get_state;
 static ompt_get_task_info_t ompt_get_task_info;
+static ompt_get_task_memory_t ompt_get_task_memory;
 static ompt_get_thread_data_t ompt_get_thread_data;
 static ompt_get_parallel_info_t ompt_get_parallel_info;
 static ompt_get_unique_id_t ompt_get_unique_id;
+static ompt_finalize_tool_t ompt_finalize_tool;
 static ompt_get_num_procs_t ompt_get_num_procs;
 static ompt_get_num_places_t ompt_get_num_places;
 static ompt_get_place_proc_ids_t ompt_get_place_proc_ids;
@@ -196,6 +198,16 @@ ompt_label_##id:
          ((uint64_t)addr) / FUZZY_ADDRESS_DISCARD_BYTES + 1,                   \
          ((uint64_t)addr) / FUZZY_ADDRESS_DISCARD_BYTES + 2, addr)
 
+#define register_callback_t(name, type)                                        \
+  do {                                                                         \
+    type f_##name = &on_##name;                                                \
+    if (ompt_set_callback(name, (ompt_callback_t)f_##name) == ompt_set_never)  \
+      printf("0: Could not register callback '" #name "'\n");                  \
+  } while (0)
+
+#define register_callback(name) register_callback_t(name, name##_t)
+
+#ifndef USE_PRIVATE_TOOL
 static void
 on_ompt_callback_mutex_acquire(
   ompt_mutex_t kind,
@@ -711,16 +723,6 @@ on_ompt_callback_control_tool(
   return 0; //success
 }
 
-#define register_callback_t(name, type)                       \
-do{                                                           \
-  type f_##name = &on_##name;                                 \
-  if (ompt_set_callback(name, (ompt_callback_t)f_##name) ==   \
-      ompt_set_never)                                         \
-    printf("0: Could not register callback '" #name "'\n");   \
-}while(0)
-
-#define register_callback(name) register_callback_t(name, name##_t)
-
 int ompt_initialize(
   ompt_function_lookup_t lookup,
   int initial_device_num,
@@ -730,9 +732,11 @@ int ompt_initialize(
   ompt_get_callback = (ompt_get_callback_t) lookup("ompt_get_callback");
   ompt_get_state = (ompt_get_state_t) lookup("ompt_get_state");
   ompt_get_task_info = (ompt_get_task_info_t) lookup("ompt_get_task_info");
+  ompt_get_task_memory = (ompt_get_task_memory_t)lookup("ompt_get_task_memory");
   ompt_get_thread_data = (ompt_get_thread_data_t) lookup("ompt_get_thread_data");
   ompt_get_parallel_info = (ompt_get_parallel_info_t) lookup("ompt_get_parallel_info");
   ompt_get_unique_id = (ompt_get_unique_id_t) lookup("ompt_get_unique_id");
+  ompt_finalize_tool = (ompt_finalize_tool_t)lookup("ompt_finalize_tool");
 
   ompt_get_num_procs = (ompt_get_num_procs_t) lookup("ompt_get_num_procs");
   ompt_get_num_places = (ompt_get_num_places_t) lookup("ompt_get_num_places");
@@ -787,3 +791,4 @@ ompt_start_tool_result_t* ompt_start_tool(
 #ifdef __cplusplus
 }
 #endif
+#endif // ifndef USE_PRIVATE_TOOL
diff --git a/runtime/test/ompt/misc/finalize_tool.c b/runtime/test/ompt/misc/finalize_tool.c
new file mode 100644
index 0000000..0a406c8
--- /dev/null
+++ b/runtime/test/ompt/misc/finalize_tool.c
@@ -0,0 +1,28 @@
+// RUN: %libomp-compile-and-run | FileCheck %s
+// REQUIRES: ompt
+#include "callback.h"
+
+int main() {
+#pragma omp parallel num_threads(2)
+  {}
+
+  printf("Before ompt_finalize_tool\n");
+  ompt_finalize_tool();
+  printf("After ompt_finalize_tool\n");
+
+  return 0;
+}
+
+// CHECK: 0: NULL_POINTER=[[NULL:.*$]]
+// CHECK: {{^}}[[THREAD_ID:[0-9]+]]: ompt_event_thread_begin:
+// CHECK-SAME: thread_type=ompt_thread_initial=1
+
+// CHECK: {{^}}[[THREAD_ID]]: ompt_event_parallel_begin
+// CHECK: {{^}}[[THREAD_ID]]: ompt_event_parallel_end
+
+// CHECK: {{^}}Before ompt_finalize_tool
+
+// CHECK: {{^}}[[THREAD_ID]]: ompt_event_thread_end: thread_id=[[THREAD_ID]]
+// CHECK: 0: ompt_event_runtime_shutdown
+
+// CHECK: {{^}}After ompt_finalize_tool
diff --git a/runtime/test/ompt/tasks/task_memory.c b/runtime/test/ompt/tasks/task_memory.c
new file mode 100644
index 0000000..a48cef2
--- /dev/null
+++ b/runtime/test/ompt/tasks/task_memory.c
@@ -0,0 +1,108 @@
+// RUN: %libomp-compile-and-run | FileCheck %s
+// REQUIRES: ompt
+// UNSUPPORTED: gcc-4, gcc-5, gcc-6, gcc-7
+#define USE_PRIVATE_TOOL 1
+#include "callback.h"
+#include <omp.h>
+
+int main() {
+  int x;
+#pragma omp parallel num_threads(2)
+  {
+#pragma omp master
+    {
+#pragma omp task
+      { x++; }
+#pragma omp task firstprivate(x)
+      { x++; }
+    }
+  }
+
+  return 0;
+}
+
+static void on_ompt_callback_implicit_task(ompt_scope_endpoint_t endpoint,
+                                           ompt_data_t *parallel_data,
+                                           ompt_data_t *task_data,
+                                           unsigned int team_size,
+                                           unsigned int thread_num, int flag) {
+  void *addr = NULL;
+  size_t size = 0;
+  int result = ompt_get_task_memory(&addr, &size, 0);
+  switch (endpoint) {
+  case ompt_scope_begin:
+    task_data->value = ompt_get_unique_id();
+    printf("ompt_event_implicit_task_begin: task_id=%" PRIu64
+           ", memory_addr=%p, memory_size=%lu, result=%d \n",
+           task_data->value, addr, size, result);
+    break;
+  case ompt_scope_end:
+    printf("ompt_event_implicit_task_end: task_id=%" PRIu64
+           ", memory_addr=%p, memory_size=%lu, result=%d \n",
+           task_data->value, addr, size, result);
+    break;
+  }
+}
+
+static void
+on_ompt_callback_task_create(ompt_data_t *encountering_task_data,
+                             const ompt_frame_t *encountering_task_frame,
+                             ompt_data_t *new_task_data, int flags,
+                             int has_dependences, const void *codeptr_ra) {
+  if (flags & ompt_task_initial)
+    return; // not interested in the initial task
+  new_task_data->value = ompt_get_unique_id();
+  void *addr = NULL;
+  size_t size = 0;
+  printf("ompt_event_task_create: task_id=%" PRIu64 "\n", new_task_data->value);
+}
+
+static void on_ompt_callback_task_schedule(ompt_data_t *first_task_data,
+                                           ompt_task_status_t prior_task_status,
+                                           ompt_data_t *second_task_data) {
+  void *addr = NULL;
+  size_t size = 0;
+  int result = ompt_get_task_memory(&addr, &size, 0);
+  printf("ompt_event_task_schedule: task_id=%" PRIu64
+         ", memory_addr=%p, memory_size=%lu, result=%d\n",
+         first_task_data->value, addr, size, result);
+}
+
+int ompt_initialize(ompt_function_lookup_t lookup, int initial_device_num,
+                    ompt_data_t *tool_data) {
+  ompt_set_callback = (ompt_set_callback_t)lookup("ompt_set_callback");
+  ompt_get_unique_id = (ompt_get_unique_id_t)lookup("ompt_get_unique_id");
+  ompt_get_task_memory = (ompt_get_task_memory_t)lookup("ompt_get_task_memory");
+
+  register_callback(ompt_callback_implicit_task);
+  register_callback(ompt_callback_task_create);
+  register_callback(ompt_callback_task_schedule);
+  printf("0: NULL_POINTER=%p\n", (void *)NULL);
+  return 1; // success
+}
+
+void ompt_finalize(ompt_data_t *tool_data) {}
+
+ompt_start_tool_result_t *ompt_start_tool(unsigned int omp_version,
+                                          const char *runtime_version) {
+  static ompt_start_tool_result_t ompt_start_tool_result = {&ompt_initialize,
+                                                            &ompt_finalize, 0};
+  return &ompt_start_tool_result;
+}
+
+// CHECK: {{^}}0: NULL_POINTER=[[NULL:.*$]]
+
+// CHECK: ompt_event_implicit_task_begin: task_id=[[TASK_ID:[0-9]+]]
+// CHECK-SAME: memory_addr=[[NULL]], memory_size=0, result=0
+
+// CHECK: ompt_event_task_create: task_id=[[TASK_ID_0:[0-9]+]]
+// CHECK: ompt_event_task_create: task_id=[[TASK_ID_1:[0-9]+]]
+
+// Expects non-zero address, size, and result
+// CHECK-DAG: ompt_event_task_schedule: task_id=[[TASK_ID_0]],
+// memory_addr=0x{{[0-f]+}}, memory_size={{[1-9][0-9]*}}, result=1
+// CHECK-DAG: ompt_event_task_schedule: task_id=[[TASK_ID_1]],
+// memory_addr=0x{{[0-f]+}}, memory_size={{[1-9][0-9]*}}, result=1
+
+// CHECK: ompt_event_implicit_task_end: task_id=[[TASK_ID]]
+// CHECK-SAME: memory_addr=[[NULL]], memory_size=0, result=0