chore: import upstream snapshot with attribution

tensorflow/c/experimental/stream_executor/BUILD

@@ -0,0 +1,138 @@
+# Description:
+# StreamExecutor C API.
+
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+)
+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 = [
+        "stream_executor.h",
+    ],
+    visibility = ["//tensorflow:__subpackages__"],
+)
+
+cc_library(
+    name = "stream_executor_hdrs",
+    hdrs = ["stream_executor.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/c:c_api_macros_hdrs",
+        "//tensorflow/c:tf_status_headers",
+    ],
+)
+
+cc_library(
+    name = "stream_executor",
+    srcs = ["stream_executor.cc"],
+    hdrs = ["stream_executor.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":stream_executor_internal",
+        "//tensorflow/c:c_api_macros",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/common_runtime/device:device_utils",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:optional",
+        "@xla//xla/stream_executor:allocator_stats",
+        "@xla//xla/stream_executor:device_description",
+        "@xla//xla/stream_executor:device_memory",
+        "@xla//xla/stream_executor:event",
+        "@xla//xla/stream_executor:executor_cache",
+        "@xla//xla/stream_executor:generic_memory_allocation",
+        "@xla//xla/stream_executor:generic_memory_allocator",
+        "@xla//xla/stream_executor:memory_allocation",
+        "@xla//xla/stream_executor:memory_allocator",
+        "@xla//xla/stream_executor:memory_space",
+        "@xla//xla/stream_executor:platform",
+        "@xla//xla/stream_executor:platform_manager",
+        "@xla//xla/stream_executor:stream",
+        "@xla//xla/stream_executor:stream_executor_common",
+        "@xla//xla/stream_executor:stream_executor_h",
+        "@xla//xla/tsl/framework:allocator",
+        "@xla//xla/tsl/platform:status",
+    ],
+)
+
+cc_library(
+    name = "stream_executor_internal",
+    hdrs = [
+        "stream_executor.h",
+        "stream_executor_internal.h",
+    ],
+    visibility = [
+        "//tensorflow/c:__subpackages__",
+        "//tensorflow/core/common_runtime/pluggable_device:__subpackages__",
+        "//tensorflow/core/kernels:__subpackages__",
+    ],
+    deps = [
+        "//tensorflow/c:c_api_macros",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_status_helper",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@xla//xla/stream_executor:device_memory",
+        "@xla//xla/stream_executor:event",
+        "@xla//xla/stream_executor:executor_cache",
+        "@xla//xla/stream_executor:platform",
+        "@xla//xla/stream_executor:stream",
+        "@xla//xla/stream_executor:stream_common",
+        "@xla//xla/stream_executor:stream_executor_h",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "stream_executor_test",
+    srcs = ["stream_executor_test.cc"],
+    deps = [
+        ":stream_executor",
+        ":stream_executor_internal",
+        ":stream_executor_test_util",
+        "//tensorflow/c:tf_status_headers",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/stream_executor:event",
+        "@xla//xla/stream_executor:platform_manager",
+        "@xla//xla/stream_executor:stream",
+        "@xla//xla/stream_executor:stream_executor_h",
+        "@xla//xla/tsl/platform:statusor",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+cc_library(
+    name = "stream_executor_test_util",
+    srcs = ["stream_executor_test_util.cc"],
+    hdrs = ["stream_executor_test_util.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":stream_executor_hdrs",
+        "//tensorflow/c:tf_status",
+    ],
+)

tensorflow/c/experimental/stream_executor/stream_executor.cc

@@ -0,0 +1,553 @@
+/* Copyright 2020 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.
+==============================================================================*/
+// This file extends/implements core stream executor base classes in terms of
+// the C API defined in stream_executor.h. A class "CSomething" represents a
+// "Something" that can be manipulated via calls in the C interface and a C
+// struct called "SP_Something".
+//
+// This file also contains stream_executor::Platform registration for pluggable
+// device.
+#include "tensorflow/c/experimental/stream_executor/stream_executor.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <variant>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/c_api_macros_internal.h"
+#include "tensorflow/c/experimental/stream_executor/stream_executor_internal.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "xla/stream_executor/allocator_stats.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/event.h"
+#include "xla/stream_executor/executor_cache.h"
+#include "xla/stream_executor/generic_memory_allocation.h"
+#include "xla/stream_executor/generic_memory_allocator.h"
+#include "xla/stream_executor/memory_allocation.h"
+#include "xla/stream_executor/memory_allocator.h"
+#include "xla/stream_executor/memory_space.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/stream_executor_common.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/status.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/common_runtime/device/device_utils.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/stringpiece.h"
+
+using tensorflow::StatusFromTF_Status;
+
+namespace stream_executor {
+using tensorflow::StringPiece;
+
+// TODO(penporn): Remove OwnedTFStatus.
+using OwnedTFStatus = tensorflow::TF_StatusPtr;
+
+namespace {
+absl::Status ValidateSPPlatform(const SP_Platform& platform) {
+  TF_VALIDATE_STRUCT_SIZE(SP_Platform, platform, SP_PLATFORM_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(SP_Platform, platform, name);
+  TF_VALIDATE_NOT_NULL(SP_Platform, platform, type);
+  TF_RETURN_IF_ERROR(
+      tensorflow::device_utils::ValidateDeviceType(platform.name));
+  TF_RETURN_IF_ERROR(
+      tensorflow::device_utils::ValidateDeviceType(platform.type));
+  // `visible_device_count` could be 0 at initialization time.
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSPPlatformFns(const SP_PlatformFns& platform_fns) {
+  TF_VALIDATE_STRUCT_SIZE(SP_PlatformFns, platform_fns,
+                          SP_PLATFORM_FNS_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(SP_PlatformFns, platform_fns, create_device);
+  TF_VALIDATE_NOT_NULL(SP_PlatformFns, platform_fns, destroy_device);
+  TF_VALIDATE_NOT_NULL(SP_PlatformFns, platform_fns, create_stream_executor);
+  TF_VALIDATE_NOT_NULL(SP_PlatformFns, platform_fns, destroy_stream_executor);
+  TF_VALIDATE_NOT_NULL(SP_PlatformFns, platform_fns, create_device_fns);
+  TF_VALIDATE_NOT_NULL(SP_PlatformFns, platform_fns, destroy_device_fns);
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSPAllocatorStats(const SP_AllocatorStats& stats) {
+  TF_VALIDATE_STRUCT_SIZE(SP_AllocatorStats, stats,
+                          SP_ALLOCATORSTATS_STRUCT_SIZE);
+  // All other fields could theoretically be zero/null.
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSPDeviceMemoryBase(const SP_DeviceMemoryBase& mem) {
+  TF_VALIDATE_STRUCT_SIZE(SP_DeviceMemoryBase, mem,
+                          SP_DEVICE_MEMORY_BASE_STRUCT_SIZE);
+  // All other fields could theoretically be zero/null.
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSPDevice(const SP_Device& device) {
+  TF_VALIDATE_STRUCT_SIZE(SP_Device, device, SP_DEVICE_STRUCT_SIZE);
+  // All other fields could theoretically be zero/null.
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSPDeviceFns(const SP_DeviceFns& device_fns) {
+  TF_VALIDATE_STRUCT_SIZE(SP_DeviceFns, device_fns, SP_DEVICE_FNS_STRUCT_SIZE);
+  // All other fields could theoretically be zero/null.
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSPStreamExecutor(const SP_StreamExecutor& se,
+                                      const SP_Platform& platform) {
+  TF_VALIDATE_STRUCT_SIZE(SP_StreamExecutor, se,
+                          SP_STREAM_EXECUTOR_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, allocate);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, deallocate);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, get_allocator_stats);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, host_memory_allocate);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, host_memory_deallocate);
+  if (platform.supports_unified_memory) {
+    TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, unified_memory_allocate);
+    TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, unified_memory_deallocate);
+  }
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, device_memory_usage);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, create_stream);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, destroy_stream);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, create_stream_dependency);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, get_stream_status);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, create_event);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, destroy_event);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, get_event_status);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, record_event);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, wait_for_event);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, memcpy_dtoh);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, memcpy_htod);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, sync_memcpy_dtoh);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, sync_memcpy_htod);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, block_host_for_event);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, synchronize_all_activity);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, host_callback);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, mem_zero);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, memset);
+  TF_VALIDATE_NOT_NULL(SP_StreamExecutor, se, memset32);
+  return absl::OkStatus();
+}
+
+absl::Status ValidateSEPlatformRegistrationParams(
+    const SE_PlatformRegistrationParams& params) {
+  TF_VALIDATE_STRUCT_SIZE(SE_PlatformRegistrationParams, params,
+                          SE_PLATFORM_REGISTRATION_PARAMS_STRUCT_SIZE);
+  TF_VALIDATE_NOT_NULL(SE_PlatformRegistrationParams, params, destroy_platform);
+  TF_VALIDATE_NOT_NULL(SE_PlatformRegistrationParams, params,
+                       destroy_platform_fns);
+  return absl::OkStatus();
+}
+#undef TF_VALIDATE_NOT_NULL
+
+DeviceAddressBase DeviceMemoryBaseFromC(const SP_DeviceMemoryBase& mem) {
+  DeviceAddressBase base(mem.opaque, mem.size);
+  base.SetPayload(mem.payload);
+  return base;
+}
+
+// Wrapper that allows passing std::function across C API.
+struct HostCallbackContext {
+  absl::AnyInvocable<absl::Status() &&> callback;
+};
+
+// This wrapper allows calling `HostCallbackContext::callback` across C API.
+// This function matches `SE_StatusCallbackFn` signature and will be passed as
+// `callback_fn` to `host_callback` in `SP_StreamExecutor`.
+void HostCallbackTrampoline(void* ctx, TF_Status* status) {
+  HostCallbackContext* host_ctx = static_cast<HostCallbackContext*>(ctx);
+  absl::Status s = std::move(host_ctx->callback)();
+  tsl::Set_TF_Status_from_Status(status, s);
+  delete host_ctx;
+}
+
+namespace {
+
+// Creates a MemoryAllocation that wraps a heap-allocated buffer.
+absl::StatusOr<std::unique_ptr<MemoryAllocation>> AllocateHostMemory(
+    SP_StreamExecutor* stream_executor, SP_Device* device, uint64_t size) {
+  void* ptr = stream_executor->host_memory_allocate(device, size);
+  if (ptr == nullptr && size > 0) {
+    return absl::InternalError("Failed to allocate host memory");
+  }
+  return std::make_unique<GenericMemoryAllocation>(
+      ptr, size, [stream_executor, device](void* ptr, uint64_t size) {
+        stream_executor->host_memory_deallocate(device, ptr);
+      });
+}
+}  // namespace
+
+class CStreamExecutor : public StreamExecutorCommon {
+ public:
+  explicit CStreamExecutor(Platform* se_platform, SP_Device device,
+                           SP_DeviceFns* device_fns,
+                           SP_StreamExecutor* stream_executor,
+                           SP_Platform* platform, SP_PlatformFns* platform_fns,
+                           const std::string& name)
+      : StreamExecutorCommon(se_platform),
+        device_(std::move(device)),
+        device_fns_(device_fns),
+        stream_executor_(stream_executor),
+        platform_(platform),
+        platform_fns_(platform_fns),
+        platform_name_(name) {}
+
+  ~CStreamExecutor() override {
+    platform_fns_->destroy_device(platform_, &device_);
+  }
+
+  absl::Status Init() override { return absl::OkStatus(); }
+
+  DeviceAddressBase Allocate(uint64_t size, int64_t memory_space) override {
+    SP_DeviceMemoryBase mem = {SP_DEVICE_MEMORY_BASE_STRUCT_SIZE};
+    stream_executor_->allocate(&device_, size, memory_space, &mem);
+    absl::Status status = ValidateSPDeviceMemoryBase(mem);
+    if (!status.ok()) {
+      LOG(ERROR) << status.message();
+    }
+    return DeviceMemoryBaseFromC(mem);
+  }
+  DeviceAddressBase Allocate(uint64_t size) {
+    return Allocate(size, /*memory_space=*/0);
+  }
+
+  void Deallocate(DeviceAddressBase* mem) override {
+    SP_DeviceMemoryBase device_memory_base = DeviceMemoryBaseToC(mem);
+    stream_executor_->deallocate(&device_, &device_memory_base);
+  }
+
+  absl::StatusOr<std::unique_ptr<MemoryAllocation>> HostMemoryAllocate(
+      uint64_t size) override {
+    return AllocateHostMemory(stream_executor_, &device_, size);
+  }
+
+  absl::optional<AllocatorStats> GetAllocatorStats() override {
+    SP_AllocatorStats c_stats{SP_ALLOCATORSTATS_STRUCT_SIZE};
+    TF_Bool has_stats =
+        stream_executor_->get_allocator_stats(&device_, &c_stats);
+    if (!has_stats) {
+      return absl::nullopt;
+    }
+    absl::Status status = ValidateSPAllocatorStats(c_stats);
+    if (!status.ok()) {
+      LOG(ERROR) << status.message();
+      return absl::nullopt;
+    }
+    ::stream_executor::AllocatorStats stats;
+    stats.num_allocs = c_stats.num_allocs;
+    stats.bytes_in_use = c_stats.bytes_in_use;
+    stats.peak_bytes_in_use = c_stats.peak_bytes_in_use;
+    stats.largest_alloc_size = c_stats.largest_alloc_size;
+    if (c_stats.has_bytes_limit) {
+      stats.bytes_limit = c_stats.bytes_limit;
+    }
+    stats.bytes_reserved = c_stats.bytes_reserved;
+    stats.peak_bytes_reserved = c_stats.peak_bytes_reserved;
+    if (c_stats.has_bytes_reservable_limit) {
+      stats.bytes_reservable_limit = c_stats.bytes_reservable_limit;
+    }
+    stats.largest_free_block_bytes = c_stats.largest_free_block_bytes;
+    return stats;
+  }
+  bool SynchronizeAllActivity() override {
+    OwnedTFStatus c_status(TF_NewStatus());
+    stream_executor_->synchronize_all_activity(&device_, c_status.get());
+    if (TF_GetCode(c_status.get()) != TF_OK) {
+      LOG(ERROR) << TF_Message(c_status.get());
+      return false;
+    }
+    return true;
+  }
+  absl::Status SynchronousMemcpy(DeviceAddressBase* gpu_dst,
+                                 const void* host_src, uint64_t size) override {
+    OwnedTFStatus c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_memory_base = DeviceMemoryBaseToC(gpu_dst);
+    stream_executor_->sync_memcpy_htod(&device_, &device_memory_base, host_src,
+                                       size, c_status.get());
+    return StatusFromTF_Status(c_status.get());
+  }
+  absl::Status SynchronousMemcpy(void* host_dst,
+                                 const DeviceAddressBase& gpu_src,
+                                 uint64_t size) override {
+    OwnedTFStatus c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_memory_base = DeviceMemoryBaseToC(&gpu_src);
+    stream_executor_->sync_memcpy_dtoh(&device_, host_dst, &device_memory_base,
+                                       size, c_status.get());
+    return StatusFromTF_Status(c_status.get());
+  }
+  void DeallocateStream(Stream* stream) override {
+    static_cast<CStream*>(stream)->Destroy();
+  }
+  absl::Status BlockHostForEvent(Stream* stream, Event* event) {
+    OwnedTFStatus c_status(TF_NewStatus());
+    SP_Event event_handle = static_cast<CEvent*>(event)->Handle();
+    stream_executor_->block_host_for_event(&device_, event_handle,
+                                           c_status.get());
+    return StatusFromTF_Status(c_status.get());
+  }
+
+  absl::Status EnablePeerAccessTo(StreamExecutor* other) override {
+    return absl::UnimplementedError(
+        "EnablePeerAccessTo is not supported by pluggable device.");
+  }
+  bool CanEnablePeerAccessTo(StreamExecutor* other) override { return false; }
+
+  bool DeviceMemoryUsage(int64_t* free, int64_t* total) const override {
+    return stream_executor_->device_memory_usage(
+        &device_, reinterpret_cast<int64_t*>(free),
+        reinterpret_cast<int64_t*>(total));
+  }
+
+  // Creates a new DeviceDescription object.
+  // Ownership is transferred to the caller.
+  absl::StatusOr<std::unique_ptr<DeviceDescription>> CreateDeviceDescription()
+      const override {
+    OwnedTFStatus c_status(TF_NewStatus());
+
+    DeviceDescription desc;
+    if (device_.hardware_name != nullptr) {
+      desc.set_name(device_.hardware_name);
+    }
+    if (device_.device_vendor != nullptr) {
+      desc.set_device_vendor(device_.device_vendor);
+    }
+    if (device_.pci_bus_id != nullptr) {
+      desc.set_pci_bus_id(device_.pci_bus_id);
+    }
+
+    if (device_fns_->get_numa_node != nullptr) {
+      int32_t numa_node = device_fns_->get_numa_node(&device_);
+      if (numa_node >= 0) {
+        desc.set_numa_node(numa_node);
+      }
+    }
+
+    if (device_fns_->get_memory_bandwidth != nullptr) {
+      int64_t memory_bandwidth = device_fns_->get_memory_bandwidth(&device_);
+      if (memory_bandwidth >= 0) {
+        desc.set_memory_bandwidth(memory_bandwidth);
+      }
+    }
+    // TODO(annarev): Add gflops field in DeviceDescription and set it here.
+    // TODO(annarev): Perhaps add `supports_unified_memory` in
+    // DeviceDescription.
+    return std::make_unique<DeviceDescription>(std::move(desc));
+  }
+
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override {
+    auto c_event = std::make_unique<CEvent>(&device_, stream_executor_);
+    TF_RETURN_IF_ERROR(c_event->Create());
+    return std::move(c_event);
+  }
+
+  absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
+      std::optional<std::variant<StreamPriority, int>> priority) override {
+    SP_StreamOptions options{SP_STREAM_OPTIONS_STRUCT_SIZE};
+    options.has_priority = priority.has_value();
+    if (priority.has_value()) {
+      options.priority =
+          std::visit([](auto p) { return static_cast<int>(p); }, *priority);
+    }
+    auto stream = std::make_unique<CStream>(&device_, stream_executor_, this);
+    TF_RETURN_IF_ERROR(stream->Create(&options));
+    return std::move(stream);
+  }
+
+  absl::StatusOr<std::unique_ptr<MemoryAllocator>> CreateMemoryAllocator(
+      MemorySpace type) override {
+    if (type == MemorySpace::kUnified) {
+      return std::make_unique<GenericMemoryAllocator>(
+          [this](uint64_t size)
+              -> absl::StatusOr<std::unique_ptr<MemoryAllocation>> {
+            void* ptr =
+                stream_executor_->unified_memory_allocate(&device_, size);
+            if (ptr == nullptr) {
+              return absl::InternalError("Failed to allocate unified memory");
+            }
+            return std::make_unique<GenericMemoryAllocation>(
+                ptr, size, [this](void* ptr, uint64_t size) {
+                  stream_executor_->unified_memory_deallocate(&device_, ptr);
+                });
+          });
+    } else if (type == MemorySpace::kHost) {
+      return std::make_unique<GenericMemoryAllocator>(
+          [this](uint64_t size)
+              -> absl::StatusOr<std::unique_ptr<MemoryAllocation>> {
+            return AllocateHostMemory(stream_executor_, &device_, size);
+          });
+    }
+    return absl::UnimplementedError(
+        absl::StrFormat("Unsupported memory type %d", type));
+  }
+
+ private:
+  SP_Device device_;
+  SP_DeviceFns* device_fns_;
+  SP_StreamExecutor* stream_executor_;
+  SP_Platform* platform_;
+  SP_PlatformFns* platform_fns_;
+  std::string platform_name_;
+};
+}  // namespace
+
+CPlatform::CPlatform(SP_Platform platform,
+                     void (*destroy_platform)(SP_Platform*),
+                     SP_PlatformFns platform_fns,
+                     void (*destroy_platform_fns)(SP_PlatformFns*),
+                     SP_DeviceFns device_fns, SP_StreamExecutor stream_executor,
+                     SP_TimerFns timer_fns)
+    : platform_(std::move(platform)),
+      destroy_platform_(destroy_platform),
+      platform_fns_(std::move(platform_fns)),
+      destroy_platform_fns_(destroy_platform_fns),
+      device_fns_(std::move(device_fns)),
+      stream_executor_(std::move(stream_executor)),
+      timer_fns_(std::move(timer_fns)),
+      name_(platform_.name),
+      platform_id_info_(platform_.name) {}
+
+CPlatform::~CPlatform() {
+  platform_fns_.destroy_device_fns(&platform_, &device_fns_);
+  platform_fns_.destroy_stream_executor(&platform_, &stream_executor_);
+  platform_fns_.destroy_timer_fns(&platform_, &timer_fns_);
+  destroy_platform_(&platform_);
+  destroy_platform_fns_(&platform_fns_);
+}
+
+absl::StatusOr<std::unique_ptr<DeviceDescription>>
+CPlatform::DescriptionForDevice(int ordinal) const {
+  TF_ASSIGN_OR_RETURN(StreamExecutor * executor, executor_cache_.Get(ordinal));
+  DeviceDescription desc = executor->GetDeviceDescription();
+  return std::make_unique<DeviceDescription>(std::move(desc));
+}
+absl::StatusOr<StreamExecutor*> CPlatform::FindExisting(int ordinal) {
+  return executor_cache_.Get(ordinal);
+}
+absl::StatusOr<StreamExecutor*> CPlatform::ExecutorForDevice(int ordinal) {
+  return executor_cache_.GetOrCreate(
+      ordinal, [this, ordinal]() { return GetUncachedExecutor(ordinal); });
+}
+absl::StatusOr<std::unique_ptr<StreamExecutor>> CPlatform::GetUncachedExecutor(
+    int ordinal) {
+  // Fill device creation params
+  SE_CreateDeviceParams device_params{SE_CREATE_DEVICE_PARAMS_STRUCT_SIZE};
+  SP_Device device{SP_DEVICE_STRUCT_SIZE};
+  device_params.device = &device;
+  device_params.ext = nullptr;
+  device_params.ordinal = ordinal;
+  OwnedTFStatus c_status(TF_NewStatus());
+
+  // Create Device
+  platform_fns_.create_device(&platform_, &device_params, c_status.get());
+  TF_RETURN_IF_ERROR(StatusFromTF_Status(c_status.get()));
+  TF_RETURN_IF_ERROR(ValidateSPDevice(device));
+
+  return std::make_unique<CStreamExecutor>(this, std::move(device),
+                                           &device_fns_, &stream_executor_,
+                                           &platform_, &platform_fns_, Name());
+}
+
+absl::Status InitStreamExecutorPlugin(void* dso_handle,
+                                      std::string* device_type,
+                                      std::string* platform_name) {
+  tensorflow::Env* env = tensorflow::Env::Default();
+
+  // Step 1: Load symbol for `TF_InitPlugin`
+  void* dso_symbol;
+  TF_RETURN_IF_ERROR(
+      env->GetSymbolFromLibrary(dso_handle, "SE_InitPlugin", &dso_symbol));
+
+  // Step 2: Call `TF_InitPlugin`
+  auto init_fn = reinterpret_cast<SEInitPluginFn>(dso_symbol);
+  return InitStreamExecutorPlugin(init_fn, device_type, platform_name);
+}
+
+absl::Status InitStreamExecutorPlugin(SEInitPluginFn init_fn,
+                                      std::string* device_type,
+                                      std::string* platform_name) {
+  SE_PlatformRegistrationParams params{
+      SE_PLATFORM_REGISTRATION_PARAMS_STRUCT_SIZE};
+  SP_Platform platform{SP_PLATFORM_STRUCT_SIZE};
+  SP_PlatformFns platform_fns{SP_PLATFORM_FNS_STRUCT_SIZE};
+  params.major_version = SE_MAJOR;
+  params.minor_version = SE_MINOR;
+  params.patch_version = SE_PATCH;
+  params.platform = &platform;
+  params.platform_fns = &platform_fns;
+
+  OwnedTFStatus c_status(TF_NewStatus());
+  init_fn(&params, c_status.get());
+  TF_RETURN_IF_ERROR(tensorflow::StatusFromTF_Status(c_status.get()));
+  TF_RETURN_IF_ERROR(ValidateSEPlatformRegistrationParams(params));
+  TF_RETURN_IF_ERROR(ValidateSPPlatform(platform));
+  TF_RETURN_IF_ERROR(ValidateSPPlatformFns(platform_fns));
+
+  // Fill SP_DeviceFns creation params
+  SE_CreateDeviceFnsParams device_fns_params{
+      SE_CREATE_DEVICE_FNS_PARAMS_STRUCT_SIZE};
+  SP_DeviceFns device_fns{SP_DEVICE_FNS_STRUCT_SIZE};
+  device_fns_params.device_fns = &device_fns;
+
+  // Create StreamExecutor
+  platform_fns.create_device_fns(&platform, &device_fns_params, c_status.get());
+  TF_RETURN_IF_ERROR(tensorflow::StatusFromTF_Status(c_status.get()));
+  TF_RETURN_IF_ERROR(ValidateSPDeviceFns(device_fns));
+
+  // Fill stream executor creation params
+  SE_CreateStreamExecutorParams se_params{
+      SE_CREATE_STREAM_EXECUTOR_PARAMS_STRUCT_SIZE};
+  SP_StreamExecutor se{SP_STREAMEXECUTOR_STRUCT_SIZE};
+  se_params.stream_executor = &se;
+
+  // Create StreamExecutor
+  platform_fns.create_stream_executor(&platform, &se_params, c_status.get());
+  TF_RETURN_IF_ERROR(tensorflow::StatusFromTF_Status(c_status.get()));
+  TF_RETURN_IF_ERROR(ValidateSPStreamExecutor(se, platform));
+
+  SP_TimerFns timer_fns{SP_TIMER_FNS_STRUCT_SIZE};
+  TF_RETURN_IF_ERROR(tensorflow::StatusFromTF_Status(c_status.get()));
+
+  // Register new platform
+  *device_type = platform.type;
+  *platform_name = platform.name;
+  std::unique_ptr<stream_executor::CPlatform> cplatform(
+      new stream_executor::CPlatform(
+          std::move(platform), params.destroy_platform, std::move(platform_fns),
+          params.destroy_platform_fns, std::move(device_fns), std::move(se),
+          std::move(timer_fns)));
+  TF_CHECK_OK(
+      stream_executor::PlatformManager::RegisterPlatform(std::move(cplatform)));
+  // TODO(annarev): Return `use_bfc_allocator` value in some way so that it is
+  // available in `PluggableDeviceProcessState` once the latter is checked in.
+  return absl::OkStatus();
+}
+}  // namespace stream_executor

tensorflow/c/experimental/stream_executor/stream_executor.h

@@ -0,0 +1,552 @@
+/* Copyright 2020 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_STREAM_EXECUTOR_STREAM_EXECUTOR_H_
+#define TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_H_
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/tf_status.h"
+
+// --------------------------------------------------------------------------
+// C API for StreamExecutor. The API is under active development and eventually
+// should allow registering a pluggable device with TensorFlow.
+//
+// Conventions:
+//   * Struct prefix indicates whether struct fields should be filled by the
+//     plugin or core implementation:
+//     * SE_ : set/filled by core unless explicitly marked otherwise.
+//     * SP_ : set/filled by plugin unless explicitly marked otherwise.
+//   * We use `struct_size` for version checking. It is exempt from the `SE/SP`
+//     rule above and should be set both by core and the plugin.
+//     * For example, `create_device` function receives `SP_Device*` as input
+//       with `struct_size` populated by core. The plugin is responsible for
+//       setting `struct_size` as well, along with all other fields.
+//     * Refer to "TensorFlow Versioning Strategy" section at
+//       https://github.com/tensorflow/community/pull/257/files.
+//     * Note that the API is still under active development and doesn't have
+//       versioning guarantees yet.
+//   * `void* ext` is a free-form field that can be populated by
+//     a plugin in `SP_*` structs or potential future extension points in `SE_`
+//     structs.
+//
+// Example usage:
+//
+//   /* Sample TensorFlow code below, exact implementation might differ. */
+//   // Version checking uses `struct_size`. It is exempt from the `SE/SP` rule
+//   // above and should be set both by core and the plugin."
+//   SP_Device device { SP_DEVICE_STRUCT_SIZE };
+//   SE_CreateDeviceParams params { SE_CREATE_DEVICE_PARAMS_STRUCT_SIZE } ;
+//   params.device = &device;
+//
+//   /* Plugin code below */
+//   constexpr char DEVICE_NAME[] = "MY_DEVICE";
+//   constexpr char DEVICE_TYPE[] = "GPU";
+//
+//   void create_device(const SP_Platform* platform,
+//                      SE_CreateDeviceParams* params, TF_Status* status) {
+//     // Custom actions based on TensorFlow's view of SP_Device.
+//     OnTFDeviceView(params->device->struct_size);
+//     params->device = { SP_DEVICE_STRUCT_SIZE };
+//     params->device->device_handle = get_my_device_handle(device->ordinal);
+//     params->device->ordinal = params->ordinal;
+//     ...
+//   }
+//
+//   void destroy_device(const SP_Platform* platform, SP_Device* device) {
+//     delete_my_device_handle(device->device_handle);
+//   }
+//
+//   void SE_InitPlugin(
+//       SE_PlatformRegistrationParams* params,
+//       TF_Status* status) {
+//     params->platform = { SP_PLATFORM_STRUCT_SIZE };
+//     // Values such as `name` and `type` must outlive SE_InitPlugin call.
+//     params->platform->name = DEVICE_NAME;
+//     params->platform->type = DEVICE_TYPE;
+//     params->platform_fns->get_device_count = get_device_count;
+//     params->platform_fns->create_device = create_device;
+//     params->platform_fns->destroy_device = destroy_device;
+//     ...
+//   }
+
+#define SE_MAJOR 0
+#define SE_MINOR 0
+#define SE_PATCH 1
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct SP_Stream_st* SP_Stream;
+typedef struct SP_Event_st* SP_Event;
+typedef struct SP_Timer_st* SP_Timer;
+
+typedef struct SP_StreamOptions {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+  TF_Bool has_priority;
+  int priority;
+} SP_StreamOptions;
+
+#define SP_STREAM_OPTIONS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SP_StreamOptions, priority)
+
+// Takes `callback_arg` passed to `host_callback` as the first argument.
+typedef void (*SE_StatusCallbackFn)(void* const, TF_Status* const);
+
+typedef struct SP_TimerFns {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+  uint64_t (*nanoseconds)(SP_Timer timer);
+} SP_TimerFns;
+
+#define SP_TIMER_FNS_STRUCT_SIZE TF_OFFSET_OF_END(SP_TimerFns, nanoseconds)
+
+typedef struct SP_AllocatorStats {
+  size_t struct_size;
+  int64_t num_allocs;
+  int64_t bytes_in_use;
+  int64_t peak_bytes_in_use;
+  int64_t largest_alloc_size;
+
+  int8_t has_bytes_limit;
+  int64_t bytes_limit;
+
+  int64_t bytes_reserved;
+  int64_t peak_bytes_reserved;
+
+  int8_t has_bytes_reservable_limit;
+  int64_t bytes_reservable_limit;
+
+  int64_t largest_free_block_bytes;
+} SP_AllocatorStats;
+
+#define SP_ALLOCATORSTATS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SP_AllocatorStats, largest_free_block_bytes)
+
+// Potential states for an SP_Event. If `poll_for_status` returns anything aside
+// from kPending or kComplete, an error has occurred; kUnknown is a bad state.
+typedef enum SE_EventStatus {
+  SE_EVENT_UNKNOWN,
+  SE_EVENT_ERROR,
+  SE_EVENT_PENDING,
+  SE_EVENT_COMPLETE,
+} SE_EventStatus;
+
+// Memory allocation information.
+// This matches DeviceMemoryBase defined here:
+// https://cs.opensource.google/tensorflow/tensorflow/+/refs/tags/v2.3.0:tensorflow/compiler/xla/stream_executor/device_memory.h;l=57
+typedef struct SP_DeviceMemoryBase {
+  size_t struct_size;
+  void* ext;  // Reserved for future use
+  // Platform-dependent value representing allocated memory.
+  // Note that the pointer does not have to be to the virtual address itself.
+  void* opaque;
+  uint64_t size;     // Size in bytes of this allocation.
+  uint64_t payload;  // Value for plugin's use
+} SP_DeviceMemoryBase;
+
+#define SP_DEVICE_MEMORY_BASE_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SP_DeviceMemoryBase, payload)
+
+typedef struct SP_Device {
+  size_t struct_size;
+  void* ext;        // free-form data set by plugin
+  int32_t ordinal;  // device index
+
+  // Device vendor can store handle to their device representation
+  // here.
+  void* device_handle;
+
+  // [Optional]
+  // Device hardware name. Used for printing.
+  // Must be null-terminated.
+  const char* hardware_name;
+
+  // [Optional]
+  // Device vendor name. Used for printing.
+  // Must be null-terminated.
+  const char* device_vendor;
+
+  // [Optional]
+  // Returns the PCI bus identifier for this device, of the form
+  // [domain]:[bus]:[device].[function]
+  // where domain number is usually 0000.
+  // Example: 0000:00:02.1
+  // For more information see:
+  // https://en.wikipedia.org/wiki/PCI_configuration_space
+  // https://www.oreilly.com/library/view/linux-device-drivers/0596005903/ch12.html
+  // Used for printing. Must be null-terminated.
+  const char* pci_bus_id;
+} SP_Device;
+
+#define SP_DEVICE_STRUCT_SIZE TF_OFFSET_OF_END(SP_Device, pci_bus_id)
+
+typedef struct SE_CreateDeviceParams {
+  size_t struct_size;
+  void* ext;        // reserved for future use
+  int32_t ordinal;  // device index
+
+  SP_Device* device;  // Input/output, struct_size set by TF for plugin to read.
+                      // Subsequently plugin fills the entire struct.
+} SE_CreateDeviceParams;
+
+#define SE_CREATE_DEVICE_PARAMS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SE_CreateDeviceParams, device)
+
+typedef struct SP_DeviceFns {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+
+  // [Optional]
+  // Returns the NUMA node associated with this device, for use in
+  // determining socket locality. If the NUMA node could not be determined, -1
+  // is returned.
+  // Negative values are treated as "unset".
+  int32_t (*get_numa_node)(const SP_Device* device);
+
+  // [Optional]
+  // Device's memory bandwidth in bytes/sec.  (This is for reads/writes to/from
+  // the device's own memory, not for transfers between the host and device.)
+  // Negative values are treated as "unset".
+  int64_t (*get_memory_bandwidth)(const SP_Device* device);
+
+  // [Optional]
+  // Estimate of average number of floating point operations per second for
+  // this device * 10e-9.
+  // Negative values are treated as "unset".
+  double (*get_gflops)(const SP_Device* device);
+} SP_DeviceFns;
+
+#define SP_DEVICE_FNS_STRUCT_SIZE TF_OFFSET_OF_END(SP_DeviceFns, get_gflops)
+
+typedef struct SE_CreateDeviceFnsParams {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+
+  SP_DeviceFns* device_fns;  // output, to be filled by plugin
+} SE_CreateDeviceFnsParams;
+
+#define SE_CREATE_DEVICE_FNS_PARAMS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SE_CreateDeviceFnsParams, device_fns)
+
+typedef struct SP_StreamExecutor {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+
+  /*** ALLOCATION CALLBACKS ***/
+  // Synchronously allocates `size` bytes on the underlying platform and returns
+  // `SP_DeviceMemoryBase` representing that allocation. In the case of failure,
+  // nullptr is returned.
+  // `memory_space` is reserved for a potential future usage and should be set
+  // to 0.
+  void (*allocate)(const SP_Device* device, uint64_t size, int64_t memory_space,
+                   SP_DeviceMemoryBase* mem);
+
+  // Deallocate the device memory previously allocated via this interface.
+  // Deallocation of a nullptr-representative value is permitted.
+  void (*deallocate)(const SP_Device* device, SP_DeviceMemoryBase* memory);
+
+  // Allocates a region of host memory and registers it with the platform API.
+  // Memory allocated in this manner is required for use in asynchronous memcpy
+  // operations, such as `memcpy_dtoh`.
+  void* (*host_memory_allocate)(const SP_Device* device, uint64_t size);
+
+  // Deallocates a region of host memory allocated by `host_memory_allocate`.
+  void (*host_memory_deallocate)(const SP_Device* device, void* mem);
+
+  // Allocates unified memory space of the given size, if supported. Unified
+  // memory support should be added by setting `supports_unified_memory` field
+  // in `SP_Platform`.
+  void* (*unified_memory_allocate)(const SP_Device* device, uint64_t bytes);
+
+  // Deallocates unified memory space previously allocated with
+  // `unified_memory_allocate`. Unified
+  // memory support should be added by setting `supports_unified_memory` field
+  // in `SP_Platform`.
+  void (*unified_memory_deallocate)(const SP_Device* device, void* location);
+
+  // Fills SP_AllocatorStats with allocator statistics, if it is available.
+  // If it is not available, return false.
+  TF_Bool (*get_allocator_stats)(const SP_Device* device,
+                                 SP_AllocatorStats* stats);
+  // Fills the underlying device memory usage information, if it is
+  // available. If it is not available (false is returned), free/total need not
+  // be initialized.
+  TF_Bool (*device_memory_usage)(const SP_Device* device, int64_t* free,
+                                 int64_t* total);
+
+  /*** STREAM CALLBACKS ***/
+  // Creates SP_Stream. This call should also allocate stream
+  // resources on the underlying platform and initializes its
+  // internals.
+  void (*create_stream)(const SP_Device* device, SP_Stream* stream,
+                        TF_Status* status);
+
+  // Destroys SP_Stream and deallocates any underlying resources.
+  void (*destroy_stream)(const SP_Device* device, SP_Stream stream);
+
+  // Causes `dependent` to not begin execution until `other` has finished its
+  // last-enqueued work.
+  void (*create_stream_dependency)(const SP_Device* device, SP_Stream dependent,
+                                   SP_Stream other, TF_Status* status);
+
+  // Without blocking the device, retrieve the current stream status.
+  void (*get_stream_status)(const SP_Device* device, SP_Stream stream,
+                            TF_Status* status);
+
+  /*** EVENT CALLBACKS ***/
+  // Create SP_Event. Performs platform-specific allocation and initialization
+  // of an event.
+  void (*create_event)(const SP_Device* device, SP_Event* event,
+                       TF_Status* status);
+
+  // Destroy SE_Event and perform any platform-specific deallocation and
+  // cleanup of an event.
+  void (*destroy_event)(const SP_Device* device, SP_Event event);
+
+  // Requests the current status of the event from the underlying platform.
+  SE_EventStatus (*get_event_status)(const SP_Device* device, SP_Event event);
+  // Inserts the specified event at the end of the specified stream.
+  void (*record_event)(const SP_Device* device, SP_Stream stream,
+                       SP_Event event, TF_Status* status);
+
+  // Wait for the specified event at the end of the specified stream.
+  void (*wait_for_event)(const SP_Device* const device, SP_Stream stream,
+                         SP_Event event, TF_Status* const status);
+
+  /*** TIMER CALLBACKS ***/
+  // Creates SP_Timer. Allocates timer resources on the underlying platform
+  // and initializes its internals, setting `timer` output variable. Sets
+  // values in `timer_fns` struct.
+  void (*create_timer)(const SP_Device* device, SP_Timer* timer,
+                       TF_Status* status);
+
+  // Destroy timer and deallocates timer resources on the underlying platform.
+  void (*destroy_timer)(const SP_Device* device, SP_Timer timer);
+
+  // Records a start event for an interval timer.
+  void (*start_timer)(const SP_Device* device, SP_Stream stream, SP_Timer timer,
+                      TF_Status* status);
+
+  // Records a stop event for an interval timer.
+  void (*stop_timer)(const SP_Device* device, SP_Stream stream, SP_Timer timer,
+                     TF_Status* status);
+
+  /*** MEMCPY CALLBACKS ***/
+  // Enqueues a memcpy operation onto stream, with a host destination location
+  // `host_dst` and a device memory source, with target size `size`.
+  void (*memcpy_dtoh)(const SP_Device* device, SP_Stream stream, void* host_dst,
+                      const SP_DeviceMemoryBase* device_src, uint64_t size,
+                      TF_Status* status);
+
+  // Enqueues a memcpy operation onto stream, with a device destination
+  // location and a host memory source, with target size `size`.
+  void (*memcpy_htod)(const SP_Device* device, SP_Stream stream,
+                      SP_DeviceMemoryBase* device_dst, const void* host_src,
+                      uint64_t size, TF_Status* status);
+
+  // Enqueues a memcpy operation onto stream, with a device destination
+  // location and a device memory source, with target size `size`.
+  void (*memcpy_dtod)(const SP_Device* device, SP_Stream stream,
+                      SP_DeviceMemoryBase* device_dst,
+                      const SP_DeviceMemoryBase* device_src, uint64_t size,
+                      TF_Status* status);
+
+  // Blocks the caller while a data segment of the given size is
+  // copied from the device source to the host destination.
+  void (*sync_memcpy_dtoh)(const SP_Device* device, void* host_dst,
+                           const SP_DeviceMemoryBase* device_src, uint64_t size,
+                           TF_Status* status);
+
+  // Blocks the caller while a data segment of the given size is
+  // copied from the host source to the device destination.
+  void (*sync_memcpy_htod)(const SP_Device* device,
+                           SP_DeviceMemoryBase* device_dst,
+                           const void* host_src, uint64_t size,
+                           TF_Status* status);
+
+  // Blocks the caller while a data segment of the given size is copied from the
+  // device source to the device destination.
+  void (*sync_memcpy_dtod)(const SP_Device* device,
+                           SP_DeviceMemoryBase* device_dst,
+                           const SP_DeviceMemoryBase* device_src, uint64_t size,
+                           TF_Status* status);
+
+  // Causes the host code to synchronously wait for the event to complete.
+  void (*block_host_for_event)(const SP_Device* device, SP_Event event,
+                               TF_Status* status);
+
+  // [Optional]
+  // Causes the host code to synchronously wait for operations entrained onto
+  // stream to complete. Effectively a join on the asynchronous device
+  // operations enqueued on the stream before this program point.
+  // If not set, then corresponding functionality will be implemented
+  // by registering an event on the `stream` and waiting for it using
+  // `block_host_for_event`.
+  void (*block_host_until_done)(const SP_Device* device, SP_Stream stream,
+                                TF_Status* status);
+
+  // Synchronizes all activity occurring in the StreamExecutor's context (most
+  // likely a whole device).
+  void (*synchronize_all_activity)(const SP_Device* device, TF_Status* status);
+
+  // Zero out `size` bytes starting at the location.
+  void (*mem_zero)(const SP_Device* device, SP_Stream stream,
+                   SP_DeviceMemoryBase* location, uint64_t size,
+                   TF_Status* status);
+
+  // Set the 8-bit patterns starting at the location with `size` bytes.
+  void (*memset)(const SP_Device* device, SP_Stream stream,
+                 SP_DeviceMemoryBase* location, uint8_t pattern, uint64_t size,
+                 TF_Status* status);
+
+  // Set the 32-bit patterns starting at the location with `size` bytes.
+  void (*memset32)(const SP_Device* device, SP_Stream stream,
+                   SP_DeviceMemoryBase* location, uint32_t pattern,
+                   uint64_t size, TF_Status* status);
+
+  // Enqueues on a stream a user-specified function to be run on the host.
+  // `callback_arg` should be passed as the first argument to `callback_fn`.
+  TF_Bool (*host_callback)(const SP_Device* device, SP_Stream stream,
+                           SE_StatusCallbackFn callback_fn, void* callback_arg);
+
+  // Same as create_stream but takes additional options.
+  void (*create_stream_with_options)(const SP_Device* device,
+                                     const SP_StreamOptions* options,
+                                     SP_Stream* stream, TF_Status* status);
+} SP_StreamExecutor;
+
+#define SP_STREAMEXECUTOR_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SP_StreamExecutor, create_stream_with_options)
+
+typedef struct SE_CreateStreamExecutorParams {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+
+  SP_StreamExecutor* stream_executor;  // output, to be filled by plugin
+} SE_CreateStreamExecutorParams;
+
+#define SE_CREATE_STREAM_EXECUTOR_PARAMS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SE_CreateStreamExecutorParams, stream_executor)
+
+typedef struct SP_Platform {
+  size_t struct_size;
+
+  void* ext;  // free-form data set by plugin
+
+  // Platform name (also referred to as subtype), for example MY_DEVICE.
+  // The name must start with a capital letter and consist of
+  // capital letters and underscores.
+  // Must be null-terminated.
+  const char* name;
+
+  // Device type name, for example GPU. Must be null-terminated.
+  // The name must start with a capital letter and consist of
+  // capital letters and underscores.
+  const char* type;
+
+  // Whether this platform supports unified memory.
+  // Unified memory is a single memory address space accessible from any device.
+  TF_Bool supports_unified_memory;
+
+  // Whether to wrap allocator for this device with an allocator that uses BFC
+  // (best-fit with coalescing) strategy.
+  TF_Bool use_bfc_allocator;
+
+  // Whether to force the memory allocations to grow over time instead of
+  // allocating it all at once. When this is set to true, the value of
+  // allow_growth is ignored.
+  TF_Bool force_memory_growth;
+} SP_Platform;
+
+#define SP_PLATFORM_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SP_Platform, force_memory_growth)
+
+typedef struct SP_PlatformFns {
+  size_t struct_size;
+
+  void* ext;  // reserved for future use
+
+  // Callbacks for getting device count
+  void (*get_device_count)(const SP_Platform* platform, int* device_count,
+                           TF_Status* status);
+  // Callbacks for creating/destroying SP_Device.
+  void (*create_device)(const SP_Platform* platform,
+                        SE_CreateDeviceParams* params, TF_Status* status);
+
+  // Clean up fields inside SP_Device that were allocated
+  // by the plugin. `device` itself should not be deleted here.
+  void (*destroy_device)(const SP_Platform* platform, SP_Device* device);
+
+  // Callbacks for creating/destroying SP_DeviceFns.
+  void (*create_device_fns)(const SP_Platform* platform,
+                            SE_CreateDeviceFnsParams* params,
+                            TF_Status* status);
+
+  // Clean up fields inside SP_DeviceFns that were allocated
+  // by the plugin. `device_fns` itself should not be deleted here.
+  void (*destroy_device_fns)(const SP_Platform* platform,
+                             SP_DeviceFns* device_fns);
+
+  // Callbacks for creating/destroying SP_StreamExecutor.
+  void (*create_stream_executor)(const SP_Platform* platform,
+                                 SE_CreateStreamExecutorParams* params,
+                                 TF_Status* status);
+  // Clean up fields inside SP_StreamExecutor that were allocated
+  // by the plugin. `stream_executor` itself should not be deleted here.
+  void (*destroy_stream_executor)(const SP_Platform* platform,
+                                  SP_StreamExecutor* stream_executor);
+
+  // Callbacks for creating/destroying SP_TimerFns.
+  void (*create_timer_fns)(const SP_Platform* platform, SP_TimerFns* timer,
+                           TF_Status* status);
+
+  void (*destroy_timer_fns)(const SP_Platform* platform,
+                            SP_TimerFns* timer_fns);
+} SP_PlatformFns;
+
+#define SP_PLATFORM_FNS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SP_PlatformFns, destroy_timer_fns)
+
+typedef struct SE_PlatformRegistrationParams {
+  size_t struct_size;
+  void* ext;  // reserved for future use
+
+  // StreamExecutor C API version.
+  int32_t major_version;
+  int32_t minor_version;
+  int32_t patch_version;
+
+  SP_Platform* platform;         // output, set by plugin
+  SP_PlatformFns* platform_fns;  // output, set by plugin
+  // Clean up fields inside SP_Platform that were allocated
+  // by the plugin. `platform` itself should not be deleted here.
+  void (*destroy_platform)(SP_Platform* platform);  // out, set by plugin
+  void (*destroy_platform_fns)(
+      SP_PlatformFns* platform_fns);  // out, set by plugin
+} SE_PlatformRegistrationParams;
+
+#define SE_PLATFORM_REGISTRATION_PARAMS_STRUCT_SIZE \
+  TF_OFFSET_OF_END(SE_PlatformRegistrationParams, destroy_platform_fns)
+
+void SE_InitPlugin(SE_PlatformRegistrationParams* params, TF_Status* status);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_H_

tensorflow/c/experimental/stream_executor/stream_executor_internal.h

@@ -0,0 +1,365 @@
+/* Copyright 2020 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.
+==============================================================================*/
+// Classes and utilities that work with StreamExecutor C API for internal use.
+// This includes functions used for device registration and interfaces needed
+// for testing.
+#ifndef TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_
+#define TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/experimental/stream_executor/stream_executor.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/event.h"
+#include "xla/stream_executor/executor_cache.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_common.h"
+#include "xla/stream_executor/stream_executor.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+
+namespace stream_executor {
+
+// Plugin initialization function that a device plugin
+// must define.
+typedef void (*SEInitPluginFn)(SE_PlatformRegistrationParams* const,
+                               TF_Status* const);
+
+// Registers StreamExecutor platform. `device_type` and `platform_name` are
+// output parameters.
+absl::Status InitStreamExecutorPlugin(void* dso_handle,
+                                      std::string* device_type,
+                                      std::string* platform_name);
+
+// Allow registering a StreamExecutor plugin using a function (used for
+// testing).
+absl::Status InitStreamExecutorPlugin(SEInitPluginFn init_fn,
+                                      std::string* device_type,
+                                      std::string* platform_name);
+
+// Converts DeviceMemoryBase to a C struct.
+inline SP_DeviceMemoryBase DeviceMemoryBaseToC(const DeviceAddressBase* mem) {
+  SP_DeviceMemoryBase device_memory_base{SP_DEVICE_MEMORY_BASE_STRUCT_SIZE};
+  // `opaque` field inside SP_DeviceMemoryBase is not const.
+  // Therefore, we need to cast away the constness before setting it.
+  device_memory_base.opaque = const_cast<void*>(mem->opaque());
+  device_memory_base.size = mem->size();
+  device_memory_base.payload = mem->payload();
+  return device_memory_base;
+}
+
+namespace internal {
+class InternalIdInfo : public Platform::IdInfo {
+ public:
+  explicit InternalIdInfo(absl::string_view name)
+      : Platform::IdInfo([](const Platform::IdInfo& self) {
+          auto real_self = static_cast<const InternalIdInfo*>(&self);
+          return real_self->name_;
+        }),
+        name_(name) {}
+
+ private:
+  absl::string_view name_;
+};
+}  // namespace internal
+
+// This file implements core stream executor base classes in terms of
+// the C API defined in stream_executor.h. A class "CSomething" represents a
+// "Something" that can be manipulated via calls in the C interface.
+class CPlatform : public Platform {
+ public:
+  explicit CPlatform(SP_Platform platform,
+                     void (*destroy_platform)(SP_Platform*),
+                     SP_PlatformFns platform_fns,
+                     void (*destroy_platform_fns)(SP_PlatformFns*),
+                     SP_DeviceFns device_fns, SP_StreamExecutor stream_executor,
+                     SP_TimerFns timer_fns);
+  ~CPlatform() override;
+
+  Id id() const override { return Platform::Id(&platform_id_info_); }
+  const std::string& Name() const override { return name_; }
+  int VisibleDeviceCount() const override {
+    int visible_device_count = 0;
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    platform_fns_.get_device_count(&platform_, &visible_device_count,
+                                   c_status.get());
+    if (TF_GetCode(c_status.get()) != TF_OK) {
+      LOG(ERROR) << TF_Message(c_status.get());
+      return 0;
+    }
+    return visible_device_count;
+  }
+  bool UseBfcAllocator() const { return platform_.use_bfc_allocator; }
+  bool ForceMemoryGrowth() const { return platform_.force_memory_growth; }
+  absl::StatusOr<std::unique_ptr<DeviceDescription>> DescriptionForDevice(
+      int ordinal) const override;
+  absl::StatusOr<StreamExecutor*> ExecutorForDevice(int ordinal) override;
+  absl::StatusOr<StreamExecutor*> FindExisting(int ordinal) override;
+
+ private:
+  // Returns a device constructed with the ordinal without
+  // looking in or storing to the Platform's executor cache.
+  // Ownership IS transferred to the caller.
+  absl::StatusOr<std::unique_ptr<StreamExecutor>> GetUncachedExecutor(
+      int ordinal);
+
+  SP_Platform platform_;
+  void (*destroy_platform_)(SP_Platform*);
+  SP_PlatformFns platform_fns_;
+  void (*destroy_platform_fns_)(SP_PlatformFns*);
+  SP_DeviceFns device_fns_;
+  SP_StreamExecutor stream_executor_;
+  SP_TimerFns timer_fns_;
+  std::string name_;
+  internal::InternalIdInfo platform_id_info_;
+  stream_executor::ExecutorCache executor_cache_;
+};
+
+class CEvent : public Event {
+ public:
+  CEvent(SP_Device* device, SP_StreamExecutor* stream_executor)
+      : device_(device),
+        stream_executor_(stream_executor),
+        event_handle_(nullptr) {}
+  ~CEvent() override { Destroy(); }
+
+  Event::Status PollForStatus() override {
+    SE_EventStatus event_status =
+        stream_executor_->get_event_status(device_, event_handle_);
+
+    switch (event_status) {
+      case SE_EVENT_ERROR:
+        return Event::Status::kError;
+      case SE_EVENT_PENDING:
+        return Event::Status::kPending;
+      case SE_EVENT_COMPLETE:
+        return Event::Status::kComplete;
+      default:
+        return Event::Status::kUnknown;
+    }
+  }
+
+  absl::Status Create() {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    stream_executor_->create_event(device_, &event_handle_, c_status.get());
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+
+  absl::Status Record(SP_Stream stream_handle) {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    stream_executor_->record_event(device_, stream_handle, event_handle_,
+                                   c_status.get());
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+
+  void Destroy() {
+    if (event_handle_ != nullptr) {
+      stream_executor_->destroy_event(device_, event_handle_);
+      event_handle_ = nullptr;
+    }
+  }
+
+  SP_Event Handle() { return event_handle_; }
+
+ private:
+  SP_Device* device_;
+  SP_StreamExecutor* stream_executor_;
+  SP_Event event_handle_;
+};
+
+class CStream : public StreamCommon {
+ public:
+  CStream(SP_Device* device, SP_StreamExecutor* stream_executor,
+          StreamExecutor* executor)
+      : StreamCommon(executor),
+        device_(device),
+        stream_executor_(stream_executor),
+        stream_handle_(nullptr) {}
+  ~CStream() override {
+    BlockHostUntilDone().IgnoreError();
+    parent()->DeallocateStream(this);
+    Destroy();
+  }
+
+  absl::Status Create(SP_StreamOptions* options) {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    if (stream_executor_->struct_size >=
+            TF_OFFSET_OF_END(SP_StreamExecutor, create_stream_with_options) &&
+        stream_executor_->create_stream_with_options != nullptr) {
+      stream_executor_->create_stream_with_options(
+          device_, options, &stream_handle_, c_status.get());
+    } else if (options != nullptr && options->has_priority) {
+      return absl::InvalidArgumentError(
+          "Stream executor does not implement `create_stream_with_options`, "
+          "priority is not supported.");
+    } else {
+      stream_executor_->create_stream(device_, &stream_handle_, c_status.get());
+    }
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+
+  void Destroy() {
+    if (stream_handle_ != nullptr) {
+      stream_executor_->destroy_stream(device_, stream_handle_);
+      stream_handle_ = nullptr;
+    }
+  }
+  absl::Status RefreshStatus() override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    stream_executor_->get_stream_status(device_, stream_handle_,
+                                        c_status.get());
+    absl::Status status = tensorflow::StatusFromTF_Status(c_status.get());
+    CheckStatus(status);
+    return status;
+  }
+
+  absl::Status RecordEvent(Event* event) override {
+    return static_cast<CEvent*>(event)->Record(stream_handle_);
+  }
+
+  absl::Status BlockHostUntilDone() override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_Stream stream_handle = Handle();
+
+    // If `block_host_until_done` is set, use it.
+    if (stream_executor_->block_host_until_done != nullptr) {
+      stream_executor_->block_host_until_done(device_, stream_handle,
+                                              c_status.get());
+      return tensorflow::StatusFromTF_Status(c_status.get());
+    }
+    // Create and record an event and then wait for it.
+    SP_Event event_handle;
+    stream_executor_->create_event(device_, &event_handle, c_status.get());
+    TF_RETURN_IF_ERROR(tensorflow::StatusFromTF_Status(c_status.get()));
+    stream_executor_->record_event(device_, stream_handle, event_handle,
+                                   c_status.get());
+    absl::Status s = tensorflow::StatusFromTF_Status(c_status.get());
+    if (!s.ok()) {
+      stream_executor_->destroy_event(device_, event_handle);
+      return s;
+    }
+    stream_executor_->block_host_for_event(device_, event_handle,
+                                           c_status.get());
+    stream_executor_->destroy_event(device_, event_handle);
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+
+  absl::Status WaitFor(Stream* other) override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_Stream other_handle = static_cast<CStream*>(other)->Handle();
+    stream_executor_->create_stream_dependency(device_, stream_handle_,
+                                               other_handle, c_status.get());
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  absl::Status WaitFor(Event* event) override {
+    SP_Event event_handle = static_cast<CEvent*>(event)->Handle();
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    stream_executor_->wait_for_event(device_, stream_handle_, event_handle,
+                                     c_status.get());
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  absl::Status MemZero(DeviceAddressBase* location, uint64_t size) override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_mem = DeviceMemoryBaseToC(location);
+    stream_executor_->mem_zero(device_, stream_handle_, &device_mem, size,
+                               c_status.get());
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  absl::Status Memset32(DeviceAddressBase* location, uint32_t pattern,
+                        uint64_t size) override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_mem = DeviceMemoryBaseToC(location);
+    stream_executor_->memset32(device_, stream_handle_, &device_mem, pattern,
+                               size, c_status.get());
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  absl::Status Memcpy(DeviceAddressBase* gpu_dst, const void* host_src,
+                      uint64_t size) override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_mem_dst = DeviceMemoryBaseToC(gpu_dst);
+    stream_executor_->memcpy_htod(device_, stream_handle_, &device_mem_dst,
+                                  host_src, size, c_status.get());
+    if (TF_GetCode(c_status.get()) != TF_OK) {
+      LOG(ERROR) << TF_Message(c_status.get());
+    }
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  absl::Status Memcpy(DeviceAddressBase* gpu_dst,
+                      const DeviceAddressBase& gpu_src,
+                      uint64_t size) override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_mem_dst = DeviceMemoryBaseToC(gpu_dst);
+    SP_DeviceMemoryBase device_mem_src = DeviceMemoryBaseToC(&gpu_src);
+    stream_executor_->memcpy_dtod(device_, stream_handle_, &device_mem_dst,
+                                  &device_mem_src, size, c_status.get());
+    if (TF_GetCode(c_status.get()) != TF_OK) {
+      LOG(ERROR) << TF_Message(c_status.get());
+    }
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  absl::Status Memcpy(void* host_dst, const DeviceAddressBase& gpu_src,
+                      uint64_t size) override {
+    tensorflow::TF_StatusPtr c_status(TF_NewStatus());
+    SP_DeviceMemoryBase device_mem_src = DeviceMemoryBaseToC(&gpu_src);
+    stream_executor_->memcpy_dtoh(device_, stream_handle_, host_dst,
+                                  &device_mem_src, size, c_status.get());
+    if (TF_GetCode(c_status.get()) != TF_OK) {
+      LOG(ERROR) << TF_Message(c_status.get());
+    }
+    return tensorflow::StatusFromTF_Status(c_status.get());
+  }
+  // Wrapper that allows passing std::function across C API.
+  struct HostCallbackContext {
+    absl::AnyInvocable<absl::Status() &&> callback;
+  };
+
+  // This wrapper allows calling `HostCallbackContext::callback` across C API.
+  // This function matches `SE_StatusCallbackFn` signature and will be passed as
+  // `callback_fn` to `host_callback` in `SP_StreamExecutor`.
+  static void HostCallbackTrampoline(void* ctx, TF_Status* status) {
+    HostCallbackContext* host_ctx = static_cast<HostCallbackContext*>(ctx);
+    absl::Status s = std::move(host_ctx->callback)();
+    tsl::Set_TF_Status_from_Status(status, s);
+    delete host_ctx;
+  }
+  absl::Status DoHostCallbackWithStatus(
+      absl::AnyInvocable<absl::Status() &&> callback) override {
+    HostCallbackContext* ctx = new HostCallbackContext{std::move(callback)};
+    if (stream_executor_->host_callback(device_, stream_handle_,
+                                        &HostCallbackTrampoline, ctx)) {
+      return absl::OkStatus();
+    }
+    return absl::InternalError("Failed to host callback.");
+  }
+  SP_Stream Handle() { return stream_handle_; }
+
+ private:
+  SP_Device* device_;
+  SP_StreamExecutor* stream_executor_;
+  SP_Stream stream_handle_;
+};
+
+}  // namespace stream_executor
+#endif  // TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_

tensorflow/c/experimental/stream_executor/stream_executor_test.cc

@@ -0,0 +1,760 @@
+/* Copyright 2020 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 "tensorflow/c/experimental/stream_executor/stream_executor.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/types/optional.h"
+#include "tensorflow/c/experimental/stream_executor/stream_executor_internal.h"
+#include "tensorflow/c/experimental/stream_executor/stream_executor_test_util.h"
+#include "tensorflow/c/tf_status.h"
+#include "xla/stream_executor/event.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+namespace stream_executor {
+namespace {
+
+/*** Registration tests ***/
+TEST(StreamExecutor, SuccessfulRegistration) {
+  auto plugin_init = [](SE_PlatformRegistrationParams* const params,
+                        TF_Status* const status) -> void {
+    TF_SetStatus(status, TF_OK, "");
+    test_util::PopulateDefaultPlatformRegistrationParams(params);
+  };
+  std::string device_type, platform_name;
+  absl::Status status =
+      InitStreamExecutorPlugin(plugin_init, &device_type, &platform_name);
+  TF_ASSERT_OK(status);
+  absl::StatusOr<Platform*> maybe_platform =
+      PlatformManager::PlatformWithName("MY_DEVICE");
+  TF_ASSERT_OK(maybe_platform.status());
+  Platform* platform = std::move(maybe_platform).value();
+  ASSERT_EQ(platform->Name(), test_util::kDeviceName);
+  ASSERT_EQ(platform->VisibleDeviceCount(), test_util::kDeviceCount);
+
+  absl::StatusOr<StreamExecutor*> maybe_executor =
+      platform->ExecutorForDevice(0);
+  TF_ASSERT_OK(maybe_executor.status());
+}
+
+TEST(StreamExecutor, NameNotSet) {
+  auto plugin_init = [](SE_PlatformRegistrationParams* const params,
+                        TF_Status* const status) -> void {
+    TF_SetStatus(status, TF_OK, "");
+    test_util::PopulateDefaultPlatformRegistrationParams(params);
+    params->platform->name = nullptr;
+  };
+
+  std::string device_type, platform_name;
+  absl::Status status =
+      InitStreamExecutorPlugin(plugin_init, &device_type, &platform_name);
+  ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
+  ASSERT_EQ(status.message(), "'name' field in SP_Platform must be set.");
+}
+
+TEST(StreamExecutor, InvalidNameWithSemicolon) {
+  auto plugin_init = [](SE_PlatformRegistrationParams* const params,
+                        TF_Status* const status) -> void {
+    TF_SetStatus(status, TF_OK, "");
+    test_util::PopulateDefaultPlatformRegistrationParams(params);
+    params->platform->name = "INVALID:NAME";
+  };
+
+  std::string device_type, platform_name;
+  absl::Status status =
+      InitStreamExecutorPlugin(plugin_init, &device_type, &platform_name);
+  ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
+  EXPECT_THAT(
+      status.message(),
+      testing::ContainsRegex("Device name/type 'INVALID:NAME' must match"));
+}
+
+TEST(StreamExecutor, InvalidNameWithSlash) {
+  auto plugin_init = [](SE_PlatformRegistrationParams* const params,
+                        TF_Status* const status) -> void {
+    TF_SetStatus(status, TF_OK, "");
+    test_util::PopulateDefaultPlatformRegistrationParams(params);
+    params->platform->name = "INVALID/";
+  };
+
+  std::string device_type, platform_name;
+  absl::Status status =
+      InitStreamExecutorPlugin(plugin_init, &device_type, &platform_name);
+  ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
+  EXPECT_THAT(status.message(),
+              testing::ContainsRegex("Device name/type 'INVALID/' must match"));
+}
+
+TEST(StreamExecutor, CreateDeviceNotSet) {
+  auto plugin_init = [](SE_PlatformRegistrationParams* const params,
+                        TF_Status* const status) -> void {
+    TF_SetStatus(status, TF_OK, "");
+    test_util::PopulateDefaultPlatformRegistrationParams(params);
+    params->platform_fns->create_device = nullptr;
+  };
+
+  std::string device_type, platform_name;
+  absl::Status status =
+      InitStreamExecutorPlugin(plugin_init, &device_type, &platform_name);
+  ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
+  ASSERT_EQ(status.message(),
+            "'create_device' field in SP_PlatformFns must be set.");
+}
+
+TEST(StreamExecutor, UnifiedMemoryAllocateNotSet) {
+  auto plugin_init = [](SE_PlatformRegistrationParams* const params,
+                        TF_Status* const status) -> void {
+    TF_SetStatus(status, TF_OK, "");
+    test_util::PopulateDefaultPlatformRegistrationParams(params);
+    params->platform->supports_unified_memory = true;
+  };
+
+  std::string device_type, platform_name;
+  absl::Status status =
+      InitStreamExecutorPlugin(plugin_init, &device_type, &platform_name);
+  ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
+  ASSERT_EQ(
+      status.message(),
+      "'unified_memory_allocate' field in SP_StreamExecutor must be set.");
+}
+
+/*** StreamExecutor behavior tests ***/
+class StreamExecutorTest : public ::testing::Test {
+ protected:
+  StreamExecutorTest() {}
+  void SetUp() override {
+    test_util::PopulateDefaultPlatform(&platform_, &platform_fns_);
+    test_util::PopulateDefaultDeviceFns(&device_fns_);
+    test_util::PopulateDefaultStreamExecutor(&se_);
+    test_util::PopulateDefaultTimerFns(&timer_fns_);
+  }
+  void TearDown() override {}
+
+  StreamExecutor* GetExecutor(int ordinal) {
+    if (!cplatform_) {
+      cplatform_ = absl::make_unique<CPlatform>(
+          platform_, test_util::DestroyPlatform, platform_fns_,
+          test_util::DestroyPlatformFns, device_fns_, se_, timer_fns_);
+    }
+    absl::StatusOr<StreamExecutor*> maybe_executor =
+        cplatform_->ExecutorForDevice(ordinal);
+    TF_CHECK_OK(maybe_executor.status());
+    return std::move(maybe_executor).value();
+  }
+  SP_Platform platform_;
+  SP_PlatformFns platform_fns_;
+  SP_DeviceFns device_fns_;
+  SP_StreamExecutor se_;
+  SP_TimerFns timer_fns_;
+  std::unique_ptr<CPlatform> cplatform_;
+};
+
+TEST_F(StreamExecutorTest, Allocate) {
+  se_.allocate = [](const SP_Device* const device, uint64_t size,
+                    int64_t memory_space, SP_DeviceMemoryBase* const mem) {
+    mem->struct_size = SP_DEVICE_MEMORY_BASE_STRUCT_SIZE;
+    mem->opaque = malloc(size);
+    mem->size = size;
+  };
+  se_.deallocate = [](const SP_Device* const device,
+                      SP_DeviceMemoryBase* const mem) {
+    EXPECT_EQ(mem->size, 2 * sizeof(int));
+    free(mem->opaque);
+    mem->opaque = nullptr;
+    mem->size = 0;
+  };
+  StreamExecutor* executor = GetExecutor(0);
+  DeviceAddress<int> mem = executor->AllocateArray<int>(2);
+  ASSERT_NE(mem.opaque(), nullptr);
+  ASSERT_EQ(mem.size(), 2 * sizeof(int));
+  executor->Deallocate(&mem);
+}
+
+TEST_F(StreamExecutorTest, HostMemoryAllocate) {
+  static bool allocate_called = false;
+  static bool deallocate_called = false;
+  se_.host_memory_allocate = [](const SP_Device* const device, uint64_t size) {
+    allocate_called = true;
+    return malloc(size);
+  };
+  se_.host_memory_deallocate = [](const SP_Device* const device, void* mem) {
+    free(mem);
+    deallocate_called = true;
+  };
+  StreamExecutor* executor = GetExecutor(0);
+  ASSERT_FALSE(allocate_called);
+  TF_ASSERT_OK_AND_ASSIGN(auto mem, executor->HostMemoryAllocate(8));
+  ASSERT_NE(mem->address().opaque(), nullptr);
+  ASSERT_TRUE(allocate_called);
+  ASSERT_FALSE(deallocate_called);
+  mem.reset();
+  ASSERT_TRUE(deallocate_called);
+}
+
+TEST_F(StreamExecutorTest, HostMemoryAllocator) {
+  static bool allocate_called = false;
+  static bool deallocate_called = false;
+  se_.host_memory_allocate = [](const SP_Device* const device, uint64_t size) {
+    allocate_called = true;
+    return malloc(size);
+  };
+  se_.host_memory_deallocate = [](const SP_Device* const device, void* mem) {
+    free(mem);
+    deallocate_called = true;
+  };
+  StreamExecutor* executor = GetExecutor(0);
+  ASSERT_FALSE(allocate_called);
+  TF_ASSERT_OK_AND_ASSIGN(auto allocator,
+                          executor->CreateMemoryAllocator(MemorySpace::kHost));
+  TF_ASSERT_OK_AND_ASSIGN(auto mem, allocator->Allocate(8));
+  ASSERT_NE(mem->address().opaque(), nullptr);
+  ASSERT_TRUE(allocate_called);
+  ASSERT_FALSE(deallocate_called);
+  mem.reset();
+  ASSERT_TRUE(deallocate_called);
+}
+
+TEST_F(StreamExecutorTest, UnifiedMemoryAllocate) {
+  static bool allocate_called = false;
+  static bool deallocate_called = false;
+  se_.unified_memory_allocate = [](const SP_Device* const device,
+                                   uint64_t size) {
+    allocate_called = true;
+    return malloc(size);
+  };
+  se_.unified_memory_deallocate = [](const SP_Device* const device, void* mem) {
+    free(mem);
+    deallocate_called = true;
+  };
+  StreamExecutor* executor = GetExecutor(0);
+  ASSERT_FALSE(allocate_called);
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto allocator, executor->CreateMemoryAllocator(MemorySpace::kUnified));
+  TF_ASSERT_OK_AND_ASSIGN(auto mem, allocator->Allocate(8));
+  ASSERT_NE(mem->address().opaque(), nullptr);
+  ASSERT_TRUE(allocate_called);
+  ASSERT_FALSE(deallocate_called);
+  mem.reset();
+  ASSERT_TRUE(deallocate_called);
+}
+
+TEST_F(StreamExecutorTest, GetAllocatorStats) {
+  se_.get_allocator_stats = [](const SP_Device* const device,
+                               SP_AllocatorStats* const stat) -> TF_Bool {
+    stat->struct_size = SP_ALLOCATORSTATS_STRUCT_SIZE;
+    stat->bytes_in_use = 123;
+    return true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  absl::optional<AllocatorStats> optional_stats = executor->GetAllocatorStats();
+  ASSERT_TRUE(optional_stats.has_value());
+  AllocatorStats stats = optional_stats.value();
+  ASSERT_EQ(stats.bytes_in_use, 123);
+}
+
+TEST_F(StreamExecutorTest, DeviceMemoryUsage) {
+  se_.device_memory_usage = [](const SP_Device* const device,
+                               int64_t* const free,
+                               int64_t* const total) -> TF_Bool {
+    *free = 45;
+    *total = 7;
+    return true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  int64_t free = 0;
+  int64_t total = 0;
+  executor->DeviceMemoryUsage(&free, &total);
+  ASSERT_EQ(free, 45);
+  ASSERT_EQ(total, 7);
+}
+
+TEST_F(StreamExecutorTest, CreateStream) {
+  static bool stream_created = false;
+  static bool stream_deleted = false;
+  se_.create_stream_with_options =
+      [](const SP_Device* const device, const SP_StreamOptions* options,
+         SP_Stream* stream, TF_Status* const status) -> void {
+    ASSERT_TRUE(options->has_priority);
+    *stream = new SP_Stream_st(14, options->priority);
+    stream_created = true;
+  };
+  se_.destroy_stream = [](const SP_Device* const device,
+                          SP_Stream stream) -> void {
+    auto custom_stream = static_cast<SP_Stream_st*>(stream);
+    ASSERT_EQ(custom_stream->stream_id, 14);
+    ASSERT_EQ(custom_stream->priority, 3);
+    delete custom_stream;
+    stream_deleted = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  ASSERT_FALSE(stream_created);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream(/*priority=*/3));
+  ASSERT_TRUE(stream_created);
+  ASSERT_FALSE(stream_deleted);
+  stream.reset();
+  ASSERT_TRUE(stream_deleted);
+}
+
+TEST_F(StreamExecutorTest, CreateStreamDependency) {
+  static bool create_stream_dependency_called = false;
+  se_.create_stream_dependency = [](const SP_Device* const device,
+                                    SP_Stream dependent, SP_Stream other,
+                                    TF_Status* const status) {
+    TF_SetStatus(status, TF_OK, "");
+    create_stream_dependency_called = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto dependent, executor->CreateStream());
+  TF_ASSERT_OK_AND_ASSIGN(auto other, executor->CreateStream());
+  ASSERT_FALSE(create_stream_dependency_called);
+  TF_ASSERT_OK(dependent->WaitFor(other.get()));
+  ASSERT_TRUE(create_stream_dependency_called);
+}
+
+TEST_F(StreamExecutorTest, StreamStatus) {
+  static bool status_ok = true;
+  se_.get_stream_status = [](const SP_Device* const device, SP_Stream stream,
+                             TF_Status* const status) -> void {
+    if (status_ok) {
+      TF_SetStatus(status, TF_OK, "");
+    } else {
+      TF_SetStatus(status, TF_INTERNAL, "Test error");
+    }
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  TF_ASSERT_OK(stream->RefreshStatus());
+  status_ok = false;
+  auto updated_status = stream->RefreshStatus();
+  ASSERT_FALSE(stream->ok());
+  ASSERT_EQ(updated_status.message(), "Test error");
+}
+
+TEST_F(StreamExecutorTest, CreateEvent) {
+  static bool event_created = false;
+  static bool event_deleted = false;
+  se_.create_event = [](const SP_Device* const device, SP_Event* event,
+                        TF_Status* const status) -> void {
+    *event = new SP_Event_st(123);
+    event_created = true;
+  };
+  se_.destroy_event = [](const SP_Device* const device,
+                         SP_Event event) -> void {
+    auto custom_event = static_cast<SP_Event_st*>(event);
+    ASSERT_EQ(custom_event->event_id, 123);
+    delete custom_event;
+    event_deleted = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  ASSERT_FALSE(event_created);
+  TF_ASSERT_OK_AND_ASSIGN(auto event, executor->CreateEvent());
+  ASSERT_TRUE(event_created);
+  ASSERT_FALSE(event_deleted);
+  event.reset();
+  ASSERT_TRUE(event_deleted);
+}
+
+TEST_F(StreamExecutorTest, PollForEventStatus) {
+  static SE_EventStatus event_status = SE_EVENT_COMPLETE;
+  se_.create_event = [](const SP_Device* const device, SP_Event* event,
+                        TF_Status* const status) -> void {
+    *event = new SP_Event_st(123);
+  };
+  se_.destroy_event = [](const SP_Device* const device,
+                         SP_Event event) -> void { delete event; };
+  se_.get_event_status = [](const SP_Device* const device,
+                            SP_Event event) -> SE_EventStatus {
+    EXPECT_EQ(event->event_id, 123);
+    return event_status;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto event, executor->CreateEvent());
+  ASSERT_EQ(event->PollForStatus(), Event::Status::kComplete);
+  event_status = SE_EVENT_ERROR;
+  ASSERT_EQ(event->PollForStatus(), Event::Status::kError);
+}
+
+TEST_F(StreamExecutorTest, RecordAndWaitForEvent) {
+  static bool record_called = false;
+  static bool wait_called = false;
+  se_.create_stream = [](const SP_Device* const device, SP_Stream* stream,
+                         TF_Status* const status) -> void {
+    *stream = new SP_Stream_st(1);
+  };
+  se_.destroy_stream = [](const SP_Device* const device,
+                          SP_Stream stream) -> void { delete stream; };
+  se_.create_event = [](const SP_Device* const device, SP_Event* event,
+                        TF_Status* const status) -> void {
+    *event = new SP_Event_st(2);
+  };
+  se_.destroy_event = [](const SP_Device* const device,
+                         SP_Event event) -> void { delete event; };
+  se_.record_event = [](const SP_Device* const device, SP_Stream stream,
+                        SP_Event event, TF_Status* const status) {
+    EXPECT_EQ(stream->stream_id, 1);
+    EXPECT_EQ(event->event_id, 2);
+    TF_SetStatus(status, TF_OK, "");
+    record_called = true;
+  };
+  se_.wait_for_event = [](const SP_Device* const device, SP_Stream stream,
+                          SP_Event event, TF_Status* const status) {
+    EXPECT_EQ(stream->stream_id, 1);
+    EXPECT_EQ(event->event_id, 2);
+    TF_SetStatus(status, TF_OK, "");
+    wait_called = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto event, executor->CreateEvent());
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  ASSERT_FALSE(record_called);
+  TF_ASSERT_OK(stream->RecordEvent(event.get()));
+  ASSERT_TRUE(record_called);
+  ASSERT_FALSE(wait_called);
+  TF_ASSERT_OK(stream->WaitFor(event.get()));
+  ASSERT_TRUE(wait_called);
+}
+
+TEST_F(StreamExecutorTest, MemcpyToHost) {
+  se_.create_stream = [](const SP_Device* const device, SP_Stream* stream,
+                         TF_Status* const status) -> void {
+    *stream = new SP_Stream_st(14);
+  };
+  se_.destroy_stream = [](const SP_Device* const device,
+                          SP_Stream stream) -> void { delete stream; };
+
+  se_.memcpy_dtoh = [](const SP_Device* const device, SP_Stream stream,
+                       void* host_dst,
+                       const SP_DeviceMemoryBase* const device_src,
+                       uint64_t size, TF_Status* const status) {
+    TF_SetStatus(status, TF_OK, "");
+    EXPECT_EQ(stream->stream_id, 14);
+    std::memcpy(host_dst, device_src->opaque, size);
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  size_t size = sizeof(int);
+  int src_data = 34;
+  int dst_data = 2;
+  DeviceAddressBase device_src(&src_data, size);
+  TF_ASSERT_OK(stream->Memcpy(&dst_data, device_src, size));
+  ASSERT_EQ(dst_data, 34);
+}
+
+TEST_F(StreamExecutorTest, MemcpyFromHost) {
+  se_.memcpy_htod = [](const SP_Device* const device, SP_Stream stream,
+                       SP_DeviceMemoryBase* const device_dst,
+                       const void* host_src, uint64_t size,
+                       TF_Status* const status) {
+    TF_SetStatus(status, TF_OK, "");
+    std::memcpy(device_dst->opaque, host_src, size);
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  size_t size = sizeof(int);
+  int src_data = 18;
+  int dst_data = 0;
+  DeviceAddressBase device_dst(&dst_data, size);
+  TF_ASSERT_OK(stream->Memcpy(&device_dst, &src_data, size));
+  ASSERT_EQ(dst_data, 18);
+}
+
+TEST_F(StreamExecutorTest, MemcpyDeviceToDevice) {
+  se_.memcpy_dtod = [](const SP_Device* const device, SP_Stream stream,
+                       SP_DeviceMemoryBase* const device_dst,
+                       const SP_DeviceMemoryBase* const device_src,
+                       uint64_t size, TF_Status* const status) {
+    TF_SetStatus(status, TF_OK, "");
+    std::memcpy(device_dst->opaque, device_src->opaque, size);
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  size_t size = sizeof(int);
+  int src_data = 18;
+  int dst_data = 0;
+  DeviceAddressBase device_dst(&dst_data, size);
+  DeviceAddressBase device_src(&src_data, size);
+  TF_ASSERT_OK(stream->Memcpy(&device_dst, device_src, size));
+  ASSERT_EQ(dst_data, 18);
+}
+
+TEST_F(StreamExecutorTest, SyncMemcpyToHost) {
+  se_.sync_memcpy_dtoh = [](const SP_Device* const device, void* host_dst,
+                            const SP_DeviceMemoryBase* const device_src,
+                            uint64_t size, TF_Status* const status) {
+    TF_SetStatus(status, TF_OK, "");
+    std::memcpy(host_dst, device_src->opaque, size);
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  size_t size = sizeof(int);
+  int src_data = 34;
+  int dst_data = 2;
+  DeviceAddressBase device_src(&src_data, size);
+  TF_ASSERT_OK(executor->SynchronousMemcpyD2H(device_src, size, &dst_data));
+  ASSERT_EQ(dst_data, 34);
+}
+
+TEST_F(StreamExecutorTest, SyncMemcpyFromHost) {
+  se_.sync_memcpy_htod =
+      [](const SP_Device* const device, SP_DeviceMemoryBase* const device_dst,
+         const void* host_src, uint64_t size, TF_Status* const status) {
+        TF_SetStatus(status, TF_OK, "");
+        std::memcpy(device_dst->opaque, host_src, size);
+      };
+
+  StreamExecutor* executor = GetExecutor(0);
+  size_t size = sizeof(int);
+  int src_data = 18;
+  int dst_data = 0;
+  DeviceAddressBase device_dst(&dst_data, size);
+  TF_ASSERT_OK(executor->SynchronousMemcpyH2D(&src_data, size, &device_dst));
+  ASSERT_EQ(dst_data, 18);
+}
+
+TEST_F(StreamExecutorTest, BlockHostForEvent) {
+  static bool block_host_for_event_called = false;
+  se_.create_event = [](const SP_Device* const device, SP_Event* event,
+                        TF_Status* const status) {
+    *event = new SP_Event_st(357);
+  };
+  se_.destroy_event = [](const SP_Device* const device, SP_Event event) {
+    delete event;
+  };
+  se_.block_host_for_event = [](const SP_Device* const device, SP_Event event,
+                                TF_Status* const status) -> void {
+    ASSERT_EQ(event->event_id, 357);
+    TF_SetStatus(status, TF_OK, "");
+    block_host_for_event_called = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  ASSERT_FALSE(block_host_for_event_called);
+  TF_ASSERT_OK(stream->BlockHostUntilDone());
+  ASSERT_TRUE(block_host_for_event_called);
+}
+
+TEST_F(StreamExecutorTest, BlockHostUntilDone) {
+  static bool block_host_until_done_called = false;
+  se_.create_stream = [](const SP_Device* const device, SP_Stream* stream,
+                         TF_Status* const status) {
+    *stream = new SP_Stream_st(58);
+  };
+  se_.destroy_stream = [](const SP_Device* const device, SP_Stream stream) {
+    delete stream;
+  };
+  se_.block_host_until_done = [](const SP_Device* const device,
+                                 SP_Stream stream,
+                                 TF_Status* const status) -> void {
+    ASSERT_EQ(stream->stream_id, 58);
+    TF_SetStatus(status, TF_OK, "");
+    block_host_until_done_called = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  ASSERT_FALSE(block_host_until_done_called);
+  TF_ASSERT_OK(stream->BlockHostUntilDone());
+  ASSERT_TRUE(block_host_until_done_called);
+}
+
+TEST_F(StreamExecutorTest, SynchronizeAllActivity) {
+  static bool synchronize_all_called = false;
+  se_.synchronize_all_activity = [](const SP_Device* const device,
+                                    TF_Status* const status) {
+    TF_SetStatus(status, TF_OK, "");
+    synchronize_all_called = true;
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  ASSERT_FALSE(synchronize_all_called);
+  ASSERT_TRUE(executor->SynchronizeAllActivity());
+  ASSERT_TRUE(synchronize_all_called);
+}
+
+TEST_F(StreamExecutorTest, HostCallbackOk) {
+  se_.host_callback = [](const SP_Device* const device, SP_Stream stream,
+                         SE_StatusCallbackFn const callback_fn,
+                         void* const callback_arg) -> TF_Bool {
+    TF_Status* status = TF_NewStatus();
+    callback_fn(callback_arg, status);
+    bool ok = TF_GetCode(status) == TF_OK;
+    TF_DeleteStatus(status);
+    return ok;
+  };
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  std::function<absl::Status()> callback = []() -> absl::Status {
+    return absl::OkStatus();
+  };
+  TF_ASSERT_OK(stream->DoHostCallbackWithStatus(callback));
+}
+
+TEST_F(StreamExecutorTest, HostCallbackError) {
+  se_.host_callback = [](const SP_Device* const device, SP_Stream stream,
+                         SE_StatusCallbackFn const callback_fn,
+                         void* const callback_arg) -> TF_Bool {
+    TF_Status* status = TF_NewStatus();
+    callback_fn(callback_arg, status);
+    bool ok = TF_GetCode(status) == TF_OK;
+    TF_DeleteStatus(status);
+    return ok;
+  };
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  std::function<absl::Status()> callback = []() -> absl::Status {
+    return absl::UnimplementedError("Unimplemented");
+  };
+  ASSERT_FALSE(stream->DoHostCallbackWithStatus(callback).ok());
+}
+
+TEST_F(StreamExecutorTest, DeviceDescription) {
+  static const char* hardware_name = "TestName";
+  static const char* vendor = "TestVendor";
+  static const char* pci_bus_id = "TestPCIBusId";
+  platform_fns_.create_device = [](const SP_Platform* platform,
+                                   SE_CreateDeviceParams* params,
+                                   TF_Status* status) {
+    params->device->hardware_name = hardware_name;
+    params->device->device_vendor = vendor;
+    params->device->pci_bus_id = pci_bus_id;
+  };
+
+  device_fns_.get_numa_node = [](const SP_Device* device) { return 123; };
+  device_fns_.get_memory_bandwidth = [](const SP_Device* device) -> int64_t {
+    return 54;
+  };
+  device_fns_.get_gflops = [](const SP_Device* device) -> double { return 32; };
+
+  StreamExecutor* executor = GetExecutor(0);
+  const DeviceDescription& description = executor->GetDeviceDescription();
+  ASSERT_EQ(description.name(), "TestName");
+  ASSERT_EQ(description.device_vendor(), "TestVendor");
+  ASSERT_EQ(description.pci_bus_id(), "TestPCIBusId");
+  ASSERT_EQ(description.numa_node(), 123);
+  ASSERT_EQ(description.memory_bandwidth(), 54);
+}
+
+TEST_F(StreamExecutorTest, DeviceDescriptionNumaNodeNotSet) {
+  static const char* hardware_name = "TestName";
+  static const char* vendor = "TestVendor";
+  static const char* pci_bus_id = "TestPCIBusId";
+  platform_fns_.create_device = [](const SP_Platform* platform,
+                                   SE_CreateDeviceParams* params,
+                                   TF_Status* status) {
+    params->device->hardware_name = hardware_name;
+    params->device->device_vendor = vendor;
+    params->device->pci_bus_id = pci_bus_id;
+  };
+
+  device_fns_.get_memory_bandwidth = [](const SP_Device* device) -> int64_t {
+    return 54;
+  };
+  device_fns_.get_gflops = [](const SP_Device* device) -> double { return 32; };
+
+  StreamExecutor* executor = GetExecutor(0);
+  const DeviceDescription& description = executor->GetDeviceDescription();
+  ASSERT_EQ(description.name(), "TestName");
+  ASSERT_EQ(description.device_vendor(), "TestVendor");
+  ASSERT_EQ(description.pci_bus_id(), "TestPCIBusId");
+  ASSERT_EQ(description.numa_node(), -1);
+  ASSERT_EQ(description.memory_bandwidth(), 54);
+}
+
+TEST_F(StreamExecutorTest, MemZero) {
+  se_.create_stream = [](const SP_Device* const device, SP_Stream* stream,
+                         TF_Status* const status) -> void {
+    *stream = new SP_Stream_st(14);
+  };
+  se_.destroy_stream = [](const SP_Device* const device,
+                          SP_Stream stream) -> void { delete stream; };
+
+  se_.mem_zero = [](const SP_Device* device, SP_Stream stream,
+                    SP_DeviceMemoryBase* location, uint64_t size,
+                    TF_Status* status) {
+    TF_SetStatus(status, TF_OK, "");
+    EXPECT_EQ(stream->stream_id, 14);
+    std::memset(location->opaque, 0, size);
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  size_t size = sizeof(int);
+  int data = 2;
+  DeviceAddressBase device_data(&data, size);
+  TF_ASSERT_OK(stream->MemZero(&device_data, size));
+  ASSERT_EQ(data, 0);
+}
+
+TEST_F(StreamExecutorTest, Memset32) {
+  se_.create_stream = [](const SP_Device* const device, SP_Stream* stream,
+                         TF_Status* const status) -> void {
+    *stream = new SP_Stream_st(14);
+  };
+  se_.destroy_stream = [](const SP_Device* const device,
+                          SP_Stream stream) -> void { delete stream; };
+
+  se_.memset32 = [](const SP_Device* device, SP_Stream stream,
+                    SP_DeviceMemoryBase* location, uint32_t pattern,
+                    uint64_t size, TF_Status* status) {
+    TF_SetStatus(status, TF_OK, "");
+    EXPECT_EQ(stream->stream_id, 14);
+    EXPECT_EQ(size % 4, 0);
+    auto ptr = static_cast<uint32_t*>(location->opaque);
+    for (int i = 0; i < size / 4; i++) {
+      *(ptr + i) = pattern;
+    }
+  };
+
+  StreamExecutor* executor = GetExecutor(0);
+  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  size_t size = sizeof(int);
+  int data = 2;
+  DeviceAddressBase device_data(&data, size);
+  TF_ASSERT_OK(stream->Memset32(&device_data, 18, size));
+  ASSERT_EQ(data, 18);
+}
+
+}  // namespace
+}  // namespace stream_executor

tensorflow/c/experimental/stream_executor/stream_executor_test_util.cc

@@ -0,0 +1,209 @@
+/* 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 "tensorflow/c/experimental/stream_executor/stream_executor_test_util.h"
+
+#include <cstdint>
+
+#include "tensorflow/c/experimental/stream_executor/stream_executor.h"
+
+namespace stream_executor {
+namespace test_util {
+
+/*** Functions for creating SP_StreamExecutor ***/
+void Allocate(const SP_Device* const device, uint64_t size,
+              int64_t memory_space, SP_DeviceMemoryBase* const mem) {}
+void Deallocate(const SP_Device* const device, SP_DeviceMemoryBase* const mem) {
+}
+void* HostMemoryAllocate(const SP_Device* const device, uint64_t size) {
+  return nullptr;
+}
+void HostMemoryDeallocate(const SP_Device* const device, void* mem) {}
+TF_Bool GetAllocatorStats(const SP_Device* const device,
+                          SP_AllocatorStats* const stats) {
+  return true;
+}
+TF_Bool DeviceMemoryUsage(const SP_Device* const device, int64_t* const free,
+                          int64_t* const total) {
+  return true;
+}
+void CreateStream(const SP_Device* const device, SP_Stream* stream,
+                  TF_Status* const status) {
+  *stream = nullptr;
+}
+void DestroyStream(const SP_Device* const device, SP_Stream stream) {}
+void CreateStreamDependency(const SP_Device* const device, SP_Stream dependent,
+                            SP_Stream other, TF_Status* const status) {}
+void GetStreamStatus(const SP_Device* const device, SP_Stream stream,
+                     TF_Status* const status) {}
+void CreateEvent(const SP_Device* const device, SP_Event* event,
+                 TF_Status* const status) {
+  *event = nullptr;
+}
+void DestroyEvent(const SP_Device* const device, SP_Event event) {}
+SE_EventStatus GetEventStatus(const SP_Device* const device, SP_Event event) {
+  return SE_EVENT_UNKNOWN;
+}
+void RecordEvent(const SP_Device* const device, SP_Stream stream,
+                 SP_Event event, TF_Status* const status) {}
+void WaitForEvent(const SP_Device* const device, SP_Stream stream,
+                  SP_Event event, TF_Status* const status) {}
+void CreateTimer(const SP_Device* const device, SP_Timer* timer,
+                 TF_Status* const status) {}
+void DestroyTimer(const SP_Device* const device, SP_Timer timer) {}
+void StartTimer(const SP_Device* const device, SP_Stream stream, SP_Timer timer,
+                TF_Status* const status) {}
+void StopTimer(const SP_Device* const device, SP_Stream stream, SP_Timer timer,
+               TF_Status* const status) {}
+void MemcpyDToH(const SP_Device* const device, SP_Stream stream, void* host_dst,
+                const SP_DeviceMemoryBase* const device_src, uint64_t size,
+                TF_Status* const status) {}
+void MemcpyHToD(const SP_Device* const device, SP_Stream stream,
+                SP_DeviceMemoryBase* const device_dst, const void* host_src,
+                uint64_t size, TF_Status* const status) {}
+void SyncMemcpyDToH(const SP_Device* const device, void* host_dst,
+                    const SP_DeviceMemoryBase* const device_src, uint64_t size,
+                    TF_Status* const status) {}
+void SyncMemcpyHToD(const SP_Device* const device,
+                    SP_DeviceMemoryBase* const device_dst, const void* host_src,
+                    uint64_t size, TF_Status* const status) {}
+void BlockHostForEvent(const SP_Device* const device, SP_Event event,
+                       TF_Status* const status) {}
+void SynchronizeAllActivity(const SP_Device* const device,
+                            TF_Status* const status) {}
+TF_Bool HostCallback(const SP_Device* const device, SP_Stream stream,
+                     SE_StatusCallbackFn const callback_fn,
+                     void* const callback_arg) {
+  return true;
+}
+
+void MemZero(const SP_Device* device, SP_Stream stream,
+             SP_DeviceMemoryBase* location, uint64_t size, TF_Status* status) {}
+
+void Memset(const SP_Device* device, SP_Stream stream,
+            SP_DeviceMemoryBase* location, uint8_t pattern, uint64_t size,
+            TF_Status* status) {}
+
+void Memset32(const SP_Device* device, SP_Stream stream,
+              SP_DeviceMemoryBase* location, uint32_t pattern, uint64_t size,
+              TF_Status* status) {}
+
+void PopulateDefaultStreamExecutor(SP_StreamExecutor* se) {
+  *se = {SP_STREAMEXECUTOR_STRUCT_SIZE};
+  se->allocate = Allocate;
+  se->deallocate = Deallocate;
+  se->host_memory_allocate = HostMemoryAllocate;
+  se->host_memory_deallocate = HostMemoryDeallocate;
+  se->get_allocator_stats = GetAllocatorStats;
+  se->device_memory_usage = DeviceMemoryUsage;
+  se->create_stream = CreateStream;
+  se->destroy_stream = DestroyStream;
+  se->create_stream_dependency = CreateStreamDependency;
+  se->get_stream_status = GetStreamStatus;
+  se->create_event = CreateEvent;
+  se->destroy_event = DestroyEvent;
+  se->get_event_status = GetEventStatus;
+  se->record_event = RecordEvent;
+  se->wait_for_event = WaitForEvent;
+  se->create_timer = CreateTimer;
+  se->destroy_timer = DestroyTimer;
+  se->start_timer = StartTimer;
+  se->stop_timer = StopTimer;
+  se->memcpy_dtoh = MemcpyDToH;
+  se->memcpy_htod = MemcpyHToD;
+  se->sync_memcpy_dtoh = SyncMemcpyDToH;
+  se->sync_memcpy_htod = SyncMemcpyHToD;
+  se->block_host_for_event = BlockHostForEvent;
+  se->synchronize_all_activity = SynchronizeAllActivity;
+  se->host_callback = HostCallback;
+  se->mem_zero = MemZero;
+  se->memset = Memset;
+  se->memset32 = Memset32;
+}
+
+void PopulateDefaultDeviceFns(SP_DeviceFns* device_fns) {
+  *device_fns = {SP_DEVICE_FNS_STRUCT_SIZE};
+}
+
+/*** Functions for creating SP_TimerFns ***/
+uint64_t Nanoseconds(SP_Timer timer) { return timer->timer_id; }
+
+void PopulateDefaultTimerFns(SP_TimerFns* timer_fns) {
+  timer_fns->nanoseconds = Nanoseconds;
+}
+
+/*** Functions for creating SP_Platform ***/
+void CreateTimerFns(const SP_Platform* platform, SP_TimerFns* timer_fns,
+                    TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  PopulateDefaultTimerFns(timer_fns);
+}
+void DestroyTimerFns(const SP_Platform* platform, SP_TimerFns* timer_fns) {}
+
+void CreateStreamExecutor(const SP_Platform* platform,
+                          SE_CreateStreamExecutorParams* params,
+                          TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  PopulateDefaultStreamExecutor(params->stream_executor);
+}
+void DestroyStreamExecutor(const SP_Platform* platform, SP_StreamExecutor* se) {
+}
+void GetDeviceCount(const SP_Platform* platform, int* device_count,
+                    TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  *device_count = kDeviceCount;
+}
+void CreateDevice(const SP_Platform* platform, SE_CreateDeviceParams* params,
+                  TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  params->device->struct_size = {SP_DEVICE_STRUCT_SIZE};
+}
+void DestroyDevice(const SP_Platform* platform, SP_Device* device) {}
+
+void CreateDeviceFns(const SP_Platform* platform,
+                     SE_CreateDeviceFnsParams* params, TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  params->device_fns->struct_size = {SP_DEVICE_FNS_STRUCT_SIZE};
+}
+void DestroyDeviceFns(const SP_Platform* platform, SP_DeviceFns* device_fns) {}
+
+void PopulateDefaultPlatform(SP_Platform* platform,
+                             SP_PlatformFns* platform_fns) {
+  *platform = {SP_PLATFORM_STRUCT_SIZE};
+  platform->name = kDeviceName;
+  platform->type = kDeviceType;
+  platform_fns->get_device_count = GetDeviceCount;
+  platform_fns->create_device = CreateDevice;
+  platform_fns->destroy_device = DestroyDevice;
+  platform_fns->create_device_fns = CreateDeviceFns;
+  platform_fns->destroy_device_fns = DestroyDeviceFns;
+  platform_fns->create_stream_executor = CreateStreamExecutor;
+  platform_fns->destroy_stream_executor = DestroyStreamExecutor;
+  platform_fns->create_timer_fns = CreateTimerFns;
+  platform_fns->destroy_timer_fns = DestroyTimerFns;
+}
+
+/*** Functions for creating SE_PlatformRegistrationParams ***/
+void DestroyPlatform(SP_Platform* platform) {}
+void DestroyPlatformFns(SP_PlatformFns* platform_fns) {}
+
+void PopulateDefaultPlatformRegistrationParams(
+    SE_PlatformRegistrationParams* const params) {
+  PopulateDefaultPlatform(params->platform, params->platform_fns);
+  params->destroy_platform = DestroyPlatform;
+  params->destroy_platform_fns = DestroyPlatformFns;
+}
+
+}  // namespace test_util
+}  // namespace stream_executor

tensorflow/c/experimental/stream_executor/stream_executor_test_util.h

@@ -0,0 +1,58 @@
+/* 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_STREAM_EXECUTOR_STREAM_EXECUTOR_TEST_UTIL_H_
+#define TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_TEST_UTIL_H_
+
+#include "tensorflow/c/experimental/stream_executor/stream_executor.h"
+
+struct SP_Stream_st {
+  explicit SP_Stream_st(int stream_id, int priority = 0)
+      : stream_id(stream_id), priority(priority) {}
+  int stream_id;
+  int priority;
+};
+
+struct SP_Event_st {
+  explicit SP_Event_st(int id) : event_id(id) {}
+  int event_id;
+};
+
+struct SP_Timer_st {
+  explicit SP_Timer_st(int id) : timer_id(id) {}
+  int timer_id;
+};
+
+namespace stream_executor {
+namespace test_util {
+
+constexpr int kDeviceCount = 2;
+constexpr char kDeviceName[] = "MY_DEVICE";
+constexpr char kDeviceType[] = "GPU";
+
+void PopulateDefaultStreamExecutor(SP_StreamExecutor* se);
+void PopulateDefaultDeviceFns(SP_DeviceFns* device_fns);
+void PopulateDefaultTimerFns(SP_TimerFns* timer_fns);
+void PopulateDefaultPlatform(SP_Platform* platform,
+                             SP_PlatformFns* platform_fns);
+void PopulateDefaultPlatformRegistrationParams(
+    SE_PlatformRegistrationParams* params);
+
+void DestroyPlatform(SP_Platform* platform);
+void DestroyPlatformFns(SP_PlatformFns* platform_fns);
+
+}  // namespace test_util
+}  // namespace stream_executor
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_TEST_UTIL_H_

tensorflow/c/experimental/stream_executor/test/BUILD

@@ -0,0 +1,33 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+# Description:
+# test for stream_executor
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_shared_object",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_cc_shared_object(
+    name = "test_pluggable_device.so",
+    srcs = ["test_pluggable_device.cc"],
+    visibility = ["//tensorflow/c:__subpackages__"],
+    deps = [
+        "//tensorflow/c/experimental/stream_executor:stream_executor_hdrs",
+        "//tensorflow/c/experimental/stream_executor:stream_executor_test_util",
+    ],
+)
+
+cc_library(
+    name = "test_pluggable_device",
+    srcs = ["test_pluggable_device.cc"],
+    visibility = ["//tensorflow/core:__subpackages__"],
+    deps = [
+        "//tensorflow/c/experimental/stream_executor:stream_executor_hdrs",
+        "//tensorflow/c/experimental/stream_executor:stream_executor_test_util",
+    ],
+)

tensorflow/c/experimental/stream_executor/test/test_pluggable_device.cc

@@ -0,0 +1,27 @@
+/* Copyright 2020 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 "tensorflow/c/experimental/stream_executor/stream_executor.h"
+#include "tensorflow/c/experimental/stream_executor/stream_executor_test_util.h"
+
+extern "C" {
+
+void SE_InitPlugin(SE_PlatformRegistrationParams* const params,
+                   TF_Status* const status) {
+  stream_executor::test_util::PopulateDefaultPlatformRegistrationParams(params);
+}
+
+void TF_InitKernel() {}
+}