chore: import upstream snapshot with attribution

tensorflow/c/experimental/pluggable_profiler/BUILD

@@ -0,0 +1,68 @@
+# Description:
+# Profiler C API
+
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+filegroup(
+    name = "headers",
+    srcs = [
+        "pluggable_profiler.h",
+    ],
+    visibility = ["//visibility:private"],
+)
+
+cc_library(
+    name = "pluggable_profiler_hdrs",
+    hdrs = ["pluggable_profiler.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/c:c_api_macros",
+        "//tensorflow/c:tf_status_headers",
+    ],
+)
+
+cc_library(
+    name = "pluggable_profiler",
+    srcs = ["pluggable_profiler.cc"],
+    hdrs = ["pluggable_profiler.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":pluggable_profiler_internal",
+        "//tensorflow/c:c_api_macros",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/common_runtime/device:device_utils",
+        "//tensorflow/core/profiler/lib:profiler_factory",
+        "//tensorflow/core/profiler/lib:profiler_interface",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+    ],
+)
+
+cc_library(
+    name = "pluggable_profiler_internal",
+    hdrs = [
+        "pluggable_profiler.h",
+        "pluggable_profiler_internal.h",
+    ],
+    visibility = ["//tensorflow/core/common_runtime/pluggable_device:__subpackages__"],
+    deps = [
+        "//tensorflow/c:c_api_macros",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/profiler/lib:profiler_interface",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+    ],
+)

tensorflow/c/experimental/pluggable_profiler/pluggable_profiler.cc

@@ -0,0 +1,191 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/c_api_macros_internal.h"
+#include "tensorflow/c/experimental/pluggable_profiler/pluggable_profiler_internal.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/common_runtime/device/device_utils.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/profiler/lib/profiler_factory.h"
+#include "tensorflow/core/profiler/lib/profiler_interface.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+
+namespace {
+
+Status ValidateTPProfilerRegistrationParams(
+    const TF_ProfilerRegistrationParams& params) {
+  TF_VALIDATE_STRUCT_SIZE(TF_ProfilerRegistrationParams, params,
+                          TF_PROFILER_REGISTRATION_PARAMS_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(TF_ProfilerRegistrationParams, params, destroy_profiler);
+  TF_VALIDATE_NOT_NULL(TF_ProfilerRegistrationParams, params,
+                       destroy_profiler_fns);
+  return absl::OkStatus();
+}
+
+Status ValidateTPProfiler(const TP_Profiler& profiler) {
+  TF_VALIDATE_STRUCT_SIZE(TP_Profiler, profiler, TP_PROFILER_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(TP_Profiler, profiler, device_type);
+  TF_RETURN_IF_ERROR(
+      tensorflow::device_utils::ValidateDeviceType(profiler.device_type));
+  return absl::OkStatus();
+}
+
+Status ValidateTPProfilerFns(const TP_ProfilerFns& profiler_fns) {
+  TF_VALIDATE_STRUCT_SIZE(TP_ProfilerFns, profiler_fns,
+                          TF_PROFILER_FNS_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(TP_ProfilerFns, profiler_fns, start);
+  TF_VALIDATE_NOT_NULL(TP_ProfilerFns, profiler_fns, stop);
+  TF_VALIDATE_NOT_NULL(TP_ProfilerFns, profiler_fns, collect_data_xspace);
+  return absl::OkStatus();
+}
+
+class PluggableProfiler : public tensorflow::profiler::ProfilerInterface {
+ public:
+  // The caller must have validated profiler_fns and profiler.
+  static std::unique_ptr<tensorflow::profiler::ProfilerInterface>
+  CreatePluggableProfiler(const ProfileOptions& options, TP_Profiler profiler,
+                          TP_ProfilerFns profiler_fns) {
+    if (options.device_tracer_level() == 0) {
+      return nullptr;
+    }
+    if (options.device_type() != ProfileOptions::PLUGGABLE_DEVICE &&
+        options.device_type() != ProfileOptions::UNSPECIFIED) {
+      return nullptr;
+    }
+    return absl::WrapUnique(new PluggableProfiler(profiler_fns, profiler));
+  }
+
+  Status Start() override {
+    tensorflow::TF_StatusPtr status(TF_NewStatus());
+    profiler_fns_.start(&profiler_, status.get());
+    return tensorflow::StatusFromTF_Status(status.get());
+  }
+
+  Status Stop() override {
+    tensorflow::TF_StatusPtr status(TF_NewStatus());
+    profiler_fns_.stop(&profiler_, status.get());
+    return tensorflow::StatusFromTF_Status(status.get());
+  }
+
+  Status CollectData(XSpace* space) override {
+    tensorflow::TF_StatusPtr status(TF_NewStatus());
+    // Get size of buffer required for Plugin to serialize XSpace into it.
+    size_t size_in_bytes;
+    profiler_fns_.collect_data_xspace(&profiler_, /*buffer=*/nullptr,
+                                      &size_in_bytes, status.get());
+
+    if (size_in_bytes <= 0)
+      return tensorflow::StatusFromTF_Status(status.get());
+
+    // Prepare an appropriately sized buffer.
+    std::vector<uint8_t> buffer(size_in_bytes);
+    profiler_fns_.collect_data_xspace(&profiler_, buffer.data(), &size_in_bytes,
+                                      status.get());
+    // Deserialize XSpace from the buffer and return it.
+    XSpace plugin_space;
+    plugin_space.ParseFromArray(buffer.data(), buffer.size());
+    for (XPlane& plugin_plane : *plugin_space.mutable_planes()) {
+      XPlane* plane = space->add_planes();
+      plane->Swap(&plugin_plane);
+    }
+    return tensorflow::StatusFromTF_Status(status.get());
+  }
+
+ private:
+  PluggableProfiler(TP_ProfilerFns profiler_fns, TP_Profiler profiler)
+      : profiler_fns_(profiler_fns), profiler_(profiler) {}
+  TP_ProfilerFns profiler_fns_;
+  TP_Profiler profiler_;
+};
+
+class PluggableProfilerFactory {
+ public:
+  PluggableProfilerFactory(TP_Profiler profiler,
+                           void (*destroy_profiler)(TP_Profiler*),
+                           TP_ProfilerFns profiler_fns,
+                           void (*destroy_profiler_fns)(TP_ProfilerFns*))
+      : profiler_(std::move(profiler)),
+        destroy_profiler_(destroy_profiler),
+        profiler_fns_(std::move(profiler_fns)),
+        destroy_profiler_fns_(destroy_profiler_fns) {}
+
+  ~PluggableProfilerFactory() {
+    destroy_profiler_(&profiler_);
+    destroy_profiler_fns_(&profiler_fns_);
+  }
+
+  std::unique_ptr<tensorflow::profiler::ProfilerInterface>
+  CreatePluggableProfiler(const ProfileOptions& options) {
+    return PluggableProfiler::CreatePluggableProfiler(options, profiler_,
+                                                      profiler_fns_);
+  }
+
+ private:
+  TP_Profiler profiler_{TP_PROFILER_STRUCT_SIZE};
+  void (*destroy_profiler_)(TP_Profiler*);
+  TP_ProfilerFns profiler_fns_{TP_PROFILER_FNS_STRUCT_SIZE};
+  void (*destroy_profiler_fns_)(TP_ProfilerFns*);
+};
+
+}  // namespace
+
+Status InitPluginProfiler(TFInitProfilerFn init_fn) {
+  TF_ProfilerRegistrationParams params{
+      TF_PROFILER_REGISTRATION_PARAMS_STRUCT_SIZE};
+  TP_Profiler profiler{TP_PROFILER_STRUCT_SIZE};
+  TP_ProfilerFns profiler_fns{TP_PROFILER_FNS_STRUCT_SIZE};
+  params.major_version = TP_MAJOR;
+  params.minor_version = TP_MINOR;
+  params.patch_version = TP_PATCH;
+  params.profiler = &profiler;
+  params.profiler_fns = &profiler_fns;
+  tensorflow::TF_StatusPtr status(TF_NewStatus());
+  init_fn(&params, status.get());
+  TF_RETURN_IF_ERROR(tensorflow::StatusFromTF_Status(status.get()));
+  TF_RETURN_IF_ERROR(ValidateTPProfilerRegistrationParams(params));
+  TF_RETURN_IF_ERROR(ValidateTPProfiler(profiler));
+  TF_RETURN_IF_ERROR(ValidateTPProfilerFns(profiler_fns));
+
+  PluggableProfilerFactory factory(std::move(profiler), params.destroy_profiler,
+                                   std::move(profiler_fns),
+                                   params.destroy_profiler_fns);
+  std::function<std::unique_ptr<ProfilerInterface>(const ProfileOptions&)>
+      create_func = [factory = std::move(factory)](
+                        const ProfileOptions& options) mutable {
+        return factory.CreatePluggableProfiler(options);
+      };
+
+  tsl::profiler::RegisterProfilerFactory(std::move(create_func));
+  return absl::OkStatus();
+}
+
+}  // namespace profiler
+}  // namespace tensorflow

tensorflow/c/experimental/pluggable_profiler/pluggable_profiler.h

@@ -0,0 +1,178 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EXPERIMENTAL_PLUGGABLE_PROFILER_PLUGGABLE_PROFILER_H_
+#define TENSORFLOW_C_EXPERIMENTAL_PLUGGABLE_PROFILER_PLUGGABLE_PROFILER_H_
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/tf_status.h"
+
+// C API for Pluggable Profiler. The API is under active development and
+// eventually should allow registering a profiler with TensorFlow.
+//
+// Conventions:
+//   * Struct prefix indicates whether struct fields should be filled by the
+//   plug-in or core TensorFlow implementation:
+//     * TF_: Set/filled by core, unless marked otherwise.
+//     * TP_: Set/filled by plug-in, unless marked otherwise.
+//     * This prefix rule only applies to structures. Enumerations and methods
+//     are all prefixed with TP_.
+//   * Structs begin with two fields:
+//     * size_t struct_size: Stores the unpadded size of the struct.
+//     * void* ext: A reserved field that may be populated by a plugin in TP_*
+//     structs or potential future extension points in TF_ structs. Must be set
+//     to zero by default if it unused.
+//  * We use struct_size for version checking by both core and plug-in.
+//    * It is exempt from the TF/TP rule above and must be set both by core and
+//    plug-in.
+//    * It can be checked programmatically to determine which struct fields are
+//    available in the structure.
+//  * When a member is added to a struct, the struct size definition must be
+//  updated to use the new last member of the struct.
+//
+// Example usage:
+//   /* Sample TensorFlow code below, exact implementation might differ. */
+//   // Version checking uses `struct_size`. It is exempt from the `TF/TP` rule
+//   // above and should be set both by core and the plugin."
+//
+//   /* Plugin code below */
+//   void profiler_start(const TP_Profiler* profiler, TF_Status* status) {
+//     /* Enable profiler */
+//     ...
+//   }
+//
+//  void profiler_stop(const TP_Profiler* profiler, TF_Status* status) {
+//    /* Disable Profiler */
+//    ...
+//  }
+//
+//  void profiler_collect_data_xspace(const TP_Profiler* profiler, uint8_t*
+//  buffer, size_t* size_in_bytes, TF_Status* status) {
+//    /* Plugin generates Xspace based on collected profiler data. */
+//    Xspace xspace = get_my_xspace();
+//    size_t buffer_size_in_bytes = *size_in_bytes;
+//    *size_in_bytes = xspace.ByteSizeLong(); /* get the size of Xspace */
+//    if (buffer == nullptr) {
+//      /* TensorFlow will first get the size of Xspace, then allocate the big
+//         enough buffer and pass it to the plugin for retrieving Xspace. */
+//      return;
+//    }
+//    bool success = xspace.SerializeToArray(buffer, buffer_size_in_bytes);
+//  }
+//
+// void TF_InitProfiler(TF_ProfilerRegistrationParams* params, TF_Status*
+// status) {
+//   *params = { TF_PROFILER_REGISTRATION_PARAMS_STRUCT_SIZE };
+//   params->profiler->struct_size = TP_PROFILER_STRUCT_SIZE;
+//   params->profiler_fns->struct_size = TP_PROFILER_FNS_STRUCT_SIZE;
+//
+//   params->profiler->type = "MyDeviceType";
+//
+//   params->profiler_fns->start =  profiler_start;
+//   params->profiler_fns->stop = profiler_stop;
+//   params->profiler_fns->collect_data_xspace = profiler_collect_data_xspace;
+//   params->destroy_profiler = profiler_destroy_profiler;
+//   params->destroy_profiler_fns = profiler_destroy_profiler_fns;
+// }
+
+#define TP_MAJOR 0
+#define TP_MINOR 0
+#define TP_PATCH 1
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// --------------------------------------------------------------------------
+// TP_Profiler holds a pointer to device type filed by the plug-in.
+typedef struct TP_Profiler {
+  size_t struct_size;
+  void* ext;  // free-form data set by plugin.
+  const char* device_type;
+
+  // The struct size must be updated when adding new members.
+#define TP_PROFILER_STRUCT_SIZE TF_OFFSET_OF_END(TP_Profiler, device_type)
+} TP_Profiler;
+
+// --------------------------------------------------------------------------
+// TP_ProfilerFns holds the profiler interface function pointers filled by the
+// plug-in.
+typedef struct TP_ProfilerFns {
+  size_t struct_size;
+
+  void* ext;  // reserved for future use.
+  // Starts profiling.
+  void (*start)(const TP_Profiler* profiler, TF_Status* status);
+  // Stops profiling.
+  void (*stop)(const TP_Profiler* profiler, TF_Status* status);
+
+  // Saves collected profile data into XSpace and serializes it to the buffer.
+  // - If `buffer` is null, returns the required buffer size in `size_in_bytes`.
+  // - If `buffer` is not null and `size_in_bytes` is the required buffer size,
+  //   `buffer` is populated with profile data in serialized XSpace format.
+  //
+  // Only the first call with a non-null `buffer` following successful calls to
+  // start and stop might return data. Subsequent calls might return empty data
+  // unless start and stop are successfully called again.
+  void (*collect_data_xspace)(const TP_Profiler* profiler, uint8_t* buffer,
+                              size_t* size_in_bytes, TF_Status* status);
+
+  // The struct size must be updated when adding new members.
+#define TP_PROFILER_FNS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(TP_ProfilerFns, collect_data_xspace)
+} TP_ProfilerFns;
+
+// TF_ProfilerRegistrationParams holds the pointers to TP_Profiler and
+// TP_ProfilerFns, the memory of TP_Profiler and TP_ProfilerFns is owned by Core
+// TensorFlow and populated by the plug-in.
+typedef struct TF_ProfilerRegistrationParams {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+
+  // TensorFlow Profiler C API version.
+  int32_t major_version;
+  int32_t minor_version;
+  int32_t patch_version;
+
+  // [in/out] Memory owned by core but attributes within are populated by the
+  // plugin.
+  TP_Profiler* profiler;
+  // [in/out] Memory owned by core but attributes within are populated by the
+  // plugin.
+  TP_ProfilerFns* profiler_fns;
+  // [out] Pointer to plugin's `TP_Profiler` clean up function.
+  // Cleans up fields inside `TP_Profiler` that were allocated
+  // by the plugin. `profiler` itself must not be deleted by the plugin.
+  void (*destroy_profiler)(TP_Profiler* profiler);
+  // [out] Pointer to plugin's `TP_ProfilerFns` clean up function.
+  // Cleans up fields inside `TP_ProfilerFns` that were allocated
+  // by the plugin. `profiler_fns` itself must not be deleted by the plugin.
+  void (*destroy_profiler_fns)(TP_ProfilerFns* profiler_fns);
+
+  // The struct size must be updated when adding new members.
+#define TF_PROFILER_REGISTRATION_PARAMS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(TF_ProfilerRegistrationParams, destroy_profiler_fns)
+} TF_ProfilerRegistrationParams;
+
+// TF_InitProfiler to do profiler registration.
+// Plug-in should implement TF_InitProfiler to register the profiler.
+void TF_InitProfiler(TF_ProfilerRegistrationParams* params, TF_Status* status);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_PLUGGABLE_PROFILER_PLUGGABLE_PROFILER_H_

tensorflow/c/experimental/pluggable_profiler/pluggable_profiler_internal.h

@@ -0,0 +1,38 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EXPERIMENTAL_PLUGGABLE_PROFILER_PLUGGABLE_PROFILER_INTERNAL_H_
+#define TENSORFLOW_C_EXPERIMENTAL_PLUGGABLE_PROFILER_PLUGGABLE_PROFILER_INTERNAL_H_
+#include "tensorflow/c/experimental/pluggable_profiler/pluggable_profiler.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/profiler/lib/profiler_interface.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+
+// Plugin initialization function that a device plugin must define. Returns
+// a TF_Status output specifying whether the initialization is successful.
+using TFInitProfilerFn = void (*)(TF_ProfilerRegistrationParams* const,
+                                  TF_Status* const);
+
+// Registers plugin's profiler to TensorFlow's profiler registry.
+absl::Status InitPluginProfiler(TFInitProfilerFn init_fn);
+
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_PLUGGABLE_PROFILER_PLUGGABLE_PROFILER_INTERNAL_H_