Repository: tensorflow/tensorflow (language: cpp)

Issue:
Added One TC for the custom ops
Added one error condition check for the allow_custom_ops set to False


Discussion / hints:
@gargn can you pls review the same
@gargn  can you pls review the PR and let me know your comments
@gargn , thanks alot for your comments, i have updated the code as per your comments. Also the TC file is very well written and i just found a missing condition which was not captured in the TC, so added the same. Kindly approve.

Regards
Amit
@gargn & @rthadur , I have updated the code as per your comments, kindly check, if everything is ok, kindly approve the PR.
@rthadur , can you pls update the label to review requested.
@gargn , thanks for approving the PR, can you please help to get this merged.

Regards
Amit
@rthadur , can you please help to get this merged.

Regards
Amit
@gargn thanks for approving the PR.
@rthadur , can you please help to get this merged.

Regards
Amit
@rthadur , can you please help to get this merged.

Regards
Amit
Can one of the admins verify this patch?
I'm going to go ahead and close this PR, because it seems to have stalled. If you're still interested in pursing this (and responding to my comments), please feel free to reopen!

Fix patch:
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 06a20a621cf838..35efa180941ea3 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -616,6 +616,7 @@ xla_test(
         "//xla/stream_executor:typed_kernel_factory",
         "//xla/stream_executor/rocm:rocm_platform_id",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -623,8 +624,6 @@ xla_test(
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
-        "@local_tsl//tsl/platform:statusor",
-        "@local_tsl//tsl/platform:test",
     ],
 )
 


Test patch:
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc
index ae9a2b602555f4..e56e7572b7e409 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <array>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -39,6 +40,7 @@ limitations under the License.
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/typed_kernel_factory.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/statusor.h"
 
 namespace stream_executor::gpu {
@@ -164,7 +166,12 @@ TEST_F(GpuKernelTest, ArrayArgByValue) {
     )";
 
   MultiKernelLoaderSpec spec(/*arity=*/2);
-  spec.AddCudaPtxInMemory(copy_kernel, "copy_kernel");
+  if (executor_->GetPlatform()->id() ==
+      stream_executor::rocm::kROCmPlatformId) {
+    spec.AddInProcessSymbol(internal::GetCopyKernel(), "copy_kernel");
+  } else {
+    spec.AddCudaPtxInMemory(copy_kernel, "copy_kernel");
+  }
 
   TF_ASSERT_OK_AND_ASSIGN(auto stream, executor_->CreateStream());
   TF_ASSERT_OK_AND_ASSIGN(auto kernel, executor_->LoadKernel(spec));
@@ -174,17 +181,14 @@ TEST_F(GpuKernelTest, ArrayArgByValue) {
   DeviceMemory<char> dst = executor_->AllocateArray<char>(kLength, 0);
   TF_ASSERT_OK(stream->MemZero(&dst, kLength));
 
-  struct ByValArg {
-    std::byte storage[16];
-  };
-  ByValArg arg;
+  std::array<std::byte, 16> storage;
   int i = 0;
-  for (auto& element : arg.storage) {
+  for (auto& element : storage) {
     element = static_cast<std::byte>(i++);
   }
 
   // Launch kernel.
-  auto args = stream_executor::PackKernelArgs(/*shmem_bytes=*/0, dst, arg);
+  auto args = stream_executor::PackKernelArgs(/*shmem_bytes=*/0, dst, storage);
   TF_ASSERT_OK(kernel->Launch(ThreadDim(), BlockDim(), stream.get(), *args));
 
   // Copy data back to host.
@@ -192,7 +196,7 @@ TEST_F(GpuKernelTest, ArrayArgByValue) {
   TF_ASSERT_OK(stream->Memcpy(dst_host, dst, kLength));
   TF_ASSERT_OK(stream->BlockHostUntilDone());
 
-  EXPECT_THAT(dst_host, ::testing::ElementsAreArray(arg.storage));
+  EXPECT_THAT(dst_host, ::testing::ElementsAreArray(storage));
 }
 }  // namespace
 }  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.cu.cc b/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.cu.cc
index 94a03f569feea6..f41f713ae15272 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.cu.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.cu.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/stream_executor/gpu/gpu_test_kernels.h"
 
+#include <array>
 #include <cstdint>
 
 #include "xla/stream_executor/kernel_spec.h"
@@ -46,6 +47,14 @@ __global__ void AddI32Ptrs3(Ptrs3<int32_t> ptrs) {
   int index = threadIdx.x + blockIdx.x * blockDim.x;
   ptrs.c[index] = ptrs.a[index] + ptrs.b[index];
 }
+
+__global__ void CopyKernel(std::byte* dst, std::array<std::byte, 16> byval) {
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < byval.size(); i++) {
+      dst[i] = byval[i];
+    }
+  }
+}
 }
 
 void* GetAddI32Kernel() { return reinterpret_cast<void*>(&AddI32); }
@@ -56,6 +65,8 @@ void* GetIncAndCmpKernel() { return reinterpret_cast<void*>(&IncAndCmp); }
 
 void* GetAddI32Ptrs3Kernel() { return reinterpret_cast<void*>(&AddI32Ptrs3); }
 
+void* GetCopyKernel() { return reinterpret_cast<void*>(&CopyKernel); }
+
 }  // namespace internal
 
 MultiKernelLoaderSpec GetAddI32KernelSpec() {
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.h b/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.h
index f64d1015f2ad4e..ff2fd2f9d28c55 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_test_kernels.h
@@ -98,6 +98,8 @@ void* GetIncAndCmpKernel();
 // StreamExecutor arguments packing for custom C++ types.
 void* GetAddI32Ptrs3Kernel();
 
+void* GetCopyKernel();
+
 }  // namespace internal
 
 // Returns an in-process kernel loader spec for the `AddI32` kernel above.