[stream_executor] NFC: Remove internal API uses outside of StreamExec…

…utor package PiperOrigin-RevId: 573401361
Intel-tensorflow · Oct 14, 2023 · 9ff124f · 9ff124f
1 parent 0acfec1
commit 9ff124f
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.cc b/tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.cc
@@ -151,7 +151,7 @@ extern "C" void _mlir_ciface_tf_launch_kernel(void *ctx, void *module_blob,
   // Get the GPU module.
   stream_executor::Stream *se_stream =
       op_kernel_ctx->op_device_context()->stream();
-  void *stream = se_stream->implementation()->GpuStreamHack();
+  void *stream = se_stream->platform_specific_handle().stream;
   GPURuntimeCache::GPUModule module = cache->LookupOrLoadModule(module_blob);
   GPURuntimeCache::GPUFunction function =
       cache->LookupOrGetFunction(module, kernel_name);

diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -711,7 +711,7 @@ Event::Status GpuExecutor::PollForEventStatus(Event* event) {
 bool GpuExecutor::AllocateStream(Stream* stream) {
   absl::MutexLock l(&alive_gpu_streams_mu_);
   bool out = AsGpuStream(stream)->Init();
-  alive_gpu_streams_[stream->implementation()->GpuStreamHack()] = stream;
+  alive_gpu_streams_[stream->platform_specific_handle().stream] = stream;
   return out;
 }
 

diff --git a/third_party/xla/xla/stream_executor/cuda/stream_search_test.cc b/third_party/xla/xla/stream_executor/cuda/stream_search_test.cc
@@ -51,8 +51,8 @@ TEST_F(StreamSearchTest, FoundPrevExecutor) {
   Stream s2(*executor);
   s2.Init();
 
-  void* gpu_ptr = s.implementation()->GpuStreamHack();
-  void* gpu_ptr_2 = s2.implementation()->GpuStreamHack();
+  void* gpu_ptr = s.platform_specific_handle().stream;
+  void* gpu_ptr_2 = s2.platform_specific_handle().stream;
 
   StreamExecutorConfig c;
   c.gpu_stream = gpu_ptr;

diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_gpu_executor.cc b/third_party/xla/xla/stream_executor/rocm/rocm_gpu_executor.cc
@@ -594,7 +594,7 @@ Event::Status GpuExecutor::PollForEventStatus(Event* event) {
 bool GpuExecutor::AllocateStream(Stream* stream) {
   absl::MutexLock l(&alive_gpu_streams_mu_);
   bool out = AsGpuStream(stream)->Init();
-  alive_gpu_streams_[stream->implementation()->GpuStreamHack()] = stream;
+  alive_gpu_streams_[stream->platform_specific_handle().stream] = stream;
   return out;
 }