[ROCm] Fix build after removal of cuda_kernels library

xla/stream_executor/gpu/BUILD

@@ -147,6 +147,9 @@ cc_library(
     name = "gpu_command_buffer",
     srcs = if_gpu_is_configured(["gpu_command_buffer.cc"]),
     hdrs = if_gpu_is_configured(["gpu_command_buffer.h"]),
+    local_defines = if_rocm_is_configured([
+        "TENSORFLOW_USE_ROCM=1",
+    ]),
     deps = [
         ":gpu_driver_header",
         ":gpu_executor_header",
@@ -174,6 +177,7 @@ cc_library(
         "//xla/stream_executor/cuda:cuda_conditional_kernels",
     ]) + if_rocm_is_configured([
         "//xla/stream_executor/rocm:hip_conditional_kernels",
+        "//xla/stream_executor/rocm:hip_noop_kernel",
     ]),
 )
 

xla/stream_executor/gpu/gpu_command_buffer.cc

@@ -197,7 +197,11 @@ absl::StatusOr<GpuCommandBuffer::NoOpKernel*> GpuCommandBuffer::GetNoOpKernel(
     auto noop_kernel = std::make_unique<NoOpKernel>(executor);
 
     MultiKernelLoaderSpec spec(/*arity=*/0);
+#if !defined(TENSORFLOW_USE_ROCM)
     spec.AddCudaPtxInMemory(gpu::kNoOpKernel, "noop");
+#else
+    spec.AddInProcessSymbol(gpu::GetNoOpKernel(), "noop");
+#endif // TENSORFLOW_USE_ROCM
     TF_RETURN_IF_ERROR(executor->GetKernel(spec, noop_kernel.get()));
 
     noop_kernel_ = std::move(noop_kernel);

xla/stream_executor/rocm/BUILD

@@ -176,6 +176,12 @@ rocm_library(
     deps = if_rocm_is_configured(["@local_config_rocm//rocm:rocm_headers"]),
 )
 
+rocm_library(
+    name = "hip_noop_kernel",
+    srcs = if_rocm_is_configured(["hip_noop_kernel.cu.cc"]),
+    deps = if_rocm_is_configured(["@local_config_rocm//rocm:rocm_headers"]),
+)
+
 cc_library(
     name = "rocm_platform",
     srcs = if_rocm_is_configured(["rocm_platform.cc"]),

xla/stream_executor/rocm/hip_noop_kernel.cu.cc

@@ -0,0 +1,28 @@
+/* Copyright 2024 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 <hip/hip_runtime.h>
+
+namespace stream_executor {
+namespace rocm {
+namespace {
+
+__global__ void NoOp() {}
+
+}  // namespace
+}  // namespace rocm
+
+namespace gpu {
+void* GetNoOpKernel() { return reinterpret_cast<void*>(&rocm::NoOp); }
+}  // namespace gpu
+
+}  // namespace stream_executor