Repository: ROCm/xla
Pull request #178: [ROCM] refactored redzone checker kernel to make it running similar to buffer comparator
Author: login: pemeliya
Created: 2025-05-05T10:28:22Z | Merged: 2025-05-08T13:01:16Z



Diff:
diff --git a/xla/stream_executor/gpu/BUILD b/xla/stream_executor/gpu/BUILD
index 00c3809ebe360..bd34af6f7cf45 100644
--- a/xla/stream_executor/gpu/BUILD
+++ b/xla/stream_executor/gpu/BUILD
@@ -23,6 +23,7 @@ load(
     "//xla/stream_executor:build_defs.bzl",
     "gpu_only_cc_library",
     "stream_executor_gpu_friends",
+    "if_gpu_is_configured",
 )
 load(
     "//xla/tests:build_defs.bzl",
@@ -353,62 +354,24 @@ gpu_only_cc_library(
     ],
 )
 
-cc_library(
-    name = "redzone_allocator_kernel_cuda",
-    srcs = [
-        "redzone_allocator_kernel.h",
-        "redzone_allocator_kernel_cuda.cc",
-    ],
-    tags = [
-        "cuda-only",
-        "gpu",
-        "manual",
-    ],
-    deps = [
-        ":gpu_asm_opts",
-        "//xla/stream_executor:device_memory",
-        "//xla/stream_executor:kernel",
-        "//xla/stream_executor:stream_executor_h",
-        "//xla/stream_executor:typed_kernel_factory",
-        "//xla/stream_executor/cuda:cuda_asm_compiler",
-        "@com_google_absl//absl/base",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:node_hash_map",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
-        "@com_google_absl//absl/types:span",
-        "@local_config_cuda//cuda:cuda_headers",
-        "@tsl//tsl/platform:statusor",
-    ],
-)
-
 gpu_kernel_library(
-    name = "redzone_allocator_kernel_rocm",
+    name = "redzone_allocator_kernel",
     srcs = [
-        "redzone_allocator_kernel.h",
-        "redzone_allocator_kernel_rocm.cu.cc",
+        "redzone_allocator_kernel.cu.cc",
     ],
-    tags = ["manual"],
-    deps = [
-        ":gpu_asm_opts",
-        "//xla/stream_executor:device_memory",
-        "//xla/stream_executor:kernel",
-        "//xla/stream_executor:stream_executor_h",
-        "//xla/stream_executor:typed_kernel_factory",
-        "@com_google_absl//absl/status:statusor",
+    tags = ["gpu"],
+    deps = if_cuda_is_configured([
+        "@local_config_cuda//cuda:cuda_headers",
+    ]) + if_rocm_is_configured([
         "@local_config_rocm//rocm:rocm_headers",
-        "@tsl//tsl/platform:statusor",
-    ],
+        "@local_config_rocm//rocm:rocm_config",
+    ]),
 )
 
 gpu_only_cc_library(
     name = "redzone_allocator",
     srcs = [
         "redzone_allocator.cc",
-        "redzone_allocator_kernel.h",
     ],
     hdrs = ["redzone_allocator.h"],
     visibility = internal_visibility([":friends"]),
@@ -435,10 +398,8 @@ gpu_only_cc_library(
         "@com_google_absl//absl/strings:string_view",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:statusor",
-    ] + if_rocm_is_configured([
-        ":redzone_allocator_kernel_rocm",
-    ]) + if_cuda_is_configured([
-        ":redzone_allocator_kernel_cuda",
+    ] + if_gpu_is_configured([
+        ":redzone_allocator_kernel",
     ]),
 )
 
diff --git a/xla/stream_executor/gpu/redzone_allocator.cc b/xla/stream_executor/gpu/redzone_allocator.cc
index 408e76996a04c..040d10c1dcc9f 100644
--- a/xla/stream_executor/gpu/redzone_allocator.cc
+++ b/xla/stream_executor/gpu/redzone_allocator.cc
@@ -33,10 +33,10 @@ limitations under the License.
 #include "xla/service/gpu/stream_executor_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/device_memory_handle.h"
 #include "xla/stream_executor/gpu/gpu_asm_opts.h"
-#include "xla/stream_executor/gpu/redzone_allocator_kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
@@ -60,6 +60,12 @@ constexpr int64_t kRhsRedzoneAlign = 4;
 
 using RedzoneCheckStatus = RedzoneAllocator::RedzoneCheckStatus;
 
+using ComparisonKernel = TypedKernel<DeviceMemory<uint8_t>, uint8_t, uint64_t,
+                                     DeviceMemory<uint64_t>>;
+namespace redzone_checker_kernel {
+  void* kernel();
+}
+
 RedzoneAllocator::RedzoneAllocator(Stream* stream,
                                    DeviceMemoryAllocator* memory_allocator,
                                    int64_t memory_limit, int64_t redzone_size,
@@ -269,10 +275,11 @@ absl::StatusOr<DeviceMemoryBase> RedzoneAllocator::CreateBuffer(
 absl::StatusOr<RedzoneCheckStatus> RedzoneAllocator::CheckRedzones() const {
   StreamExecutor* executor = stream_->parent();
 
-  TF_ASSIGN_OR_RETURN(ComparisonKernel * kernel,
-                      GetComparisonKernel(stream_->parent(), GpuAsmOpts()));
+  TF_ASSIGN_OR_RETURN(auto kernel, 
+      ComparisonKernel::FactoryType::Create(executor, "RedzoneCheckerKernel",
+                                      redzone_checker_kernel::kernel()));
 
-  stream_executor::DeviceMemoryHandle out_param(
+  DeviceMemoryHandle out_param(
       executor, executor->AllocateScalar<uint64_t>());
   TF_RETURN_IF_ERROR(
       stream_->MemZero(out_param.memory_ptr(), sizeof(uint64_t)));
@@ -282,7 +289,7 @@ absl::StatusOr<RedzoneCheckStatus> RedzoneAllocator::CheckRedzones() const {
         RedzoneCheckStatus redzone_status,
         CheckRedzonesForBuffer(stream_, *buf_and_size.first,
                                DeviceMemory<uint64_t>(out_param.memory()),
-                               *kernel, buf_and_size.second, redzone_size_,
+                               kernel, buf_and_size.second, redzone_size_,
                                redzone_pattern_));
     if (!redzone_status.ok()) {
       return redzone_status;
diff --git a/xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc b/xla/stream_executor/gpu/redzone_allocator_kernel.cu.cc
similarity index 53%
rename from xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc
rename to xla/stream_executor/gpu/redzone_allocator_kernel.cu.cc
index 2e701f2c0ddbb..173afa5d153fc 100644
--- a/xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc
+++ b/xla/stream_executor/gpu/redzone_allocator_kernel.cu.cc
@@ -15,16 +15,10 @@ limitations under the License.
 
 #include <cstdint>
 
-#include "absl/status/statusor.h"
-#include "xla/stream_executor/device_memory.h"
-#include "xla/stream_executor/gpu/redzone_allocator_kernel.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/stream_executor.h"
-#include "xla/stream_executor/typed_kernel_factory.h"
-#include "tsl/platform/statusor.h"
+namespace stream_executor::redzone_checker_kernel {
 
 namespace {
-__global__ void redzone_checker_kernel(uint8_t* input_buffer,
+__global__ void kernel_func(uint8_t* input_buffer,
                                        uint8_t redzone_pattern,
                                        uint64_t buffer_length,
                                        uint32_t* out_mismatched_ptr) {
@@ -34,18 +28,6 @@ __global__ void redzone_checker_kernel(uint8_t* input_buffer,
 }
 }  // namespace
 
-namespace stream_executor {
+void* kernel() { return reinterpret_cast<void*>(kernel_func); }
 
-absl::StatusOr<ComparisonKernel*> GetComparisonKernel(
-    StreamExecutor* executor, GpuAsmOpts /*gpu_asm_opts*/) {
-  static auto kernel = TypedKernelFactory<
-      DeviceMemory<uint8_t>, uint8_t, uint64_t,
-      DeviceMemory<uint64_t>>::Create(executor, "redzone_checker",
-                                      reinterpret_cast<void*>(
-                                          redzone_checker_kernel));
-
-  if (!kernel.ok()) return kernel.status();
-  return &kernel.value();
-}
-
-}  // namespace stream_executor
+}  // namespace stream_executor::redzone_checker_kernel
diff --git a/xla/stream_executor/gpu/redzone_allocator_kernel.h b/xla/stream_executor/gpu/redzone_allocator_kernel.h
deleted file mode 100644
index 7f1a3c3420ae0..0000000000000
--- a/xla/stream_executor/gpu/redzone_allocator_kernel.h
+++ /dev/null
@@ -1,42 +0,0 @@
-/* Copyright 2024 The OpenXLA Authors.
-
-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 XLA_STREAM_EXECUTOR_GPU_REDZONE_ALLOCATOR_KERNEL_H_
-#define XLA_STREAM_EXECUTOR_GPU_REDZONE_ALLOCATOR_KERNEL_H_
-
-#include <cstdint>
-
-#include "absl/status/statusor.h"
-#include "xla/stream_executor/device_memory.h"
-#include "xla/stream_executor/gpu/gpu_asm_opts.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/stream_executor.h"
-
-namespace stream_executor {
-using ComparisonKernel = TypedKernel<DeviceMemory<uint8_t>, uint8_t, uint64_t,
-                                     DeviceMemory<uint64_t>>;
-
-// Returns a GPU kernel that checks a memory location for redzone patterns.
-// Parameters are (buffer_address, redzone_pattern, buffer_length,
-// mismatch_count_ptr). For each byte in buffer `[buffer_address :
-// buffer_address
-// + buffer_length]` that is not equal to `redzone_pattern`,
-// `*mismatch_count_ptr` gets incremented by 1.
-absl::StatusOr<ComparisonKernel*> GetComparisonKernel(StreamExecutor* executor,
-                                                      GpuAsmOpts gpu_asm_opts);
-
-}  // namespace stream_executor
-
-#endif  // XLA_STREAM_EXECUTOR_GPU_REDZONE_ALLOCATOR_KERNEL_H_
diff --git a/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc b/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc
deleted file mode 100644
index 3c50e64850367..0000000000000
--- a/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc
+++ /dev/null
@@ -1,146 +0,0 @@
-/* Copyright 2024 The OpenXLA Authors.
-
-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 <cstdint>
-#include <tuple>
-#include <utility>
-
-#include "absl/base/call_once.h"
-#include "absl/base/const_init.h"
-#include "absl/base/thread_annotations.h"
-#include "absl/container/node_hash_map.h"
-#include "absl/log/check.h"
-#include "absl/log/log.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/string_view.h"
-#include "absl/synchronization/mutex.h"
-#include "absl/types/span.h"
-#include "xla/stream_executor/cuda/cuda_asm_compiler.h"
-#include "xla/stream_executor/device_memory.h"
-#include "xla/stream_executor/gpu/gpu_asm_opts.h"
-#include "xla/stream_executor/gpu/redzone_allocator_kernel.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/stream_executor.h"
-#include "xla/stream_executor/typed_kernel_factory.h"
-#include "tsl/platform/statusor.h"
-
-namespace stream_executor {
-// Maintains a cache of pointers to loaded kernels
-template <typename... Args>
-static absl::StatusOr<TypedKernel<Args...>*> LoadKernelOrGetPtr(
-    StreamExecutor* executor, absl::string_view kernel_name,
-    absl::string_view ptx, absl::Span<const uint8_t> cubin_data) {
-  using KernelPtrCacheKey =
-      std::tuple<StreamExecutor*, absl::string_view, absl::string_view>;
-
-  static absl::Mutex kernel_ptr_cache_mutex(absl::kConstInit);
-  static auto& kernel_ptr_cache ABSL_GUARDED_BY(kernel_ptr_cache_mutex) =
-      *new absl::node_hash_map<KernelPtrCacheKey, TypedKernel<Args...>>();
-  KernelPtrCacheKey kernel_ptr_cache_key{executor, kernel_name, ptx};
-  absl::MutexLock lock(&kernel_ptr_cache_mutex);
-
-  auto it = kernel_ptr_cache.find(kernel_ptr_cache_key);
-  if (it == kernel_ptr_cache.end()) {
-    TF_ASSIGN_OR_RETURN(TypedKernel<Args...> loaded,
-                        (TypedKernelFactory<Args...>::Create(
-                            executor, kernel_name, ptx, cubin_data)));
-    it =
-        kernel_ptr_cache.emplace(kernel_ptr_cache_key, std::move(loaded)).first;
-  }
-
-  CHECK(it != kernel_ptr_cache.end());
-  return &it->second;
-}
-
-// PTX blob for the function which checks that every byte in
-// input_buffer (length is buffer_length) is equal to redzone_pattern.
-//
-// On mismatch, increment the counter pointed to by out_mismatch_cnt_ptr.
-//
-// Generated from:
-// __global__ void redzone_checker(unsigned char* input_buffer,
-//                                 unsigned char redzone_pattern,
-//                                 unsigned long long buffer_length,
-//                                 int* out_mismatched_ptr) {
-//   unsigned long long idx = threadIdx.x + blockIdx.x * blockDim.x;
-//   if (idx >= buffer_length) return;
-//   if (input_buffer[idx] != redzone_pattern) atomicAdd(out_mismatched_ptr, 1);
-// }
-//
-// Code must compile for the oldest GPU XLA may be compiled for.
-static const char* redzone_checker_ptx = R"(
-.version 4.2
-.target sm_30
-.address_size 64
-
-.visible .entry redzone_checker(
-  .param .u64 input_buffer,
-  .param .u8 redzone_pattern,
-  .param .u64 buffer_length,
-  .param .u64 out_mismatch_cnt_ptr
-)
-{
-  .reg .pred   %p<3>;
-  .reg .b16   %rs<3>;
-  .reg .b32   %r<6>;
-  .reg .b64   %rd<8>;
-
-  ld.param.u64   %rd6, [buffer_length];
-  mov.u32   %r1, %tid.x;
-  mov.u32   %r2, %ctaid.x;
-  mov.u32   %r3, %ntid.x;
-  mad.lo.s32   %r4, %r3, %r2, %r1;
-  cvt.u64.u32   %rd3, %r4;
-  setp.ge.u64   %p1, %rd3, %rd6;
-  @%p1 bra   LBB6_3;
-  ld.param.u8   %rs1, [redzone_pattern];
-  ld.param.u64   %rd4, [input_buffer];
-  cvta.to.global.u64   %rd2, %rd4;
-  add.s64   %rd7, %rd2, %rd3;
-  ld.global.u8   %rs2, [%rd7];
-  setp.eq.s16   %p2, %rs2, %rs1;
-  @%p2 bra   LBB6_3;
-  ld.param.u64   %rd5, [out_mismatch_cnt_ptr];
-  cvta.to.global.u64   %rd1, %rd5;
-  atom.global.add.u32   %r5, [%rd1], 1;
-LBB6_3:
-  ret;
-}
-)";
-
-absl::StatusOr<ComparisonKernel*> GetComparisonKernel(StreamExecutor* executor,
-                                                      GpuAsmOpts gpu_asm_opts) {
-  absl::Span<const uint8_t> compiled_ptx = {};
-  absl::StatusOr<absl::Span<const uint8_t>> compiled_ptx_or =
-      CompileGpuAsmOrGetCached(
-          executor->GetDeviceDescription().cuda_compute_capability(),
-          redzone_checker_ptx, gpu_asm_opts);
-  if (compiled_ptx_or.ok()) {
-    compiled_ptx = compiled_ptx_or.value();
-  } else {
-    static absl::once_flag ptxas_not_found_logged;
-    absl::call_once(ptxas_not_found_logged, [&]() {
-      LOG(WARNING) << compiled_ptx_or.status()
-                   << "\nRelying on driver to perform ptx compilation. "
-                   << "\nModify $PATH to customize ptxas location."
-                   << "\nThis message will be only logged once.";
-    });
-  }
-
-  return LoadKernelOrGetPtr<DeviceMemory<uint8_t>, uint8_t, uint64_t,
-                            DeviceMemory<uint64_t>>(
-      executor, "redzone_checker", redzone_checker_ptx, compiled_ptx);
-}
-}  // namespace stream_executor
diff --git a/xla/stream_executor/rocm/hip_blas_lt.cc b/xla/stream_executor/rocm/hip_blas_lt.cc
index a432035ec6284..56e5932aab744 100644
--- a/xla/stream_executor/rocm/hip_blas_lt.cc
+++ b/xla/stream_executor/rocm/hip_blas_lt.cc
@@ -259,8 +259,7 @@ auto BlasLt::MatmulPlan::GetAlgorithms(const Stream* stream,
     // removed later when this limitation is gone.
     auto IsFP8 = [&](const MatrixLayout& layout) -> bool {
       return layout.type() == HIP_R_8F_E5M2_FNUZ ||
-             layout.type() == HIP_R_8F_E4M3_FNUZ ||
-             layout.type() == HIP_R_8F_E5M2 || layout.type() == HIP_R_8F_E4M3;
+             layout.type() == HIP_R_8F_E4M3_FNUZ;
     };
     if (IsFP8(a_desc_) && IsFP8(b_desc_)) {
       static int64_t dummy_pointer = 0xACEBALL;


Reviews:
user: login: i-chaochen
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T15:33:15Z
user: login: pemeliya
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-07T15:55:04Z
user: login: i-chaochen
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-08T13:00:54Z

Review comments:
diff_hunk: @@ -259,8 +259,7 @@ auto BlasLt::MatmulPlan::GetAlgorithms(const Stream* stream,
     // removed later when this limitation is gone.
     auto IsFP8 = [&](const MatrixLayout& layout) -> bool {
       return layout.type() == HIP_R_8F_E5M2_FNUZ ||
-             layout.type() == HIP_R_8F_E4M3_FNUZ ||
-             layout.type() == HIP_R_8F_E5M2 || layout.type() == HIP_R_8F_E4M3;
path: xla/stream_executor/rocm/hip_blas_lt.cc
user: login: i-chaochen
body: may I ask why we remove nano fp8? is it part of redzone checker?
created_at: 2025-05-06T15:32:32Z
author_association: COLLABORATOR
diff_hunk: @@ -259,8 +259,7 @@ auto BlasLt::MatmulPlan::GetAlgorithms(const Stream* stream,
     // removed later when this limitation is gone.
     auto IsFP8 = [&](const MatrixLayout& layout) -> bool {
       return layout.type() == HIP_R_8F_E5M2_FNUZ ||
-             layout.type() == HIP_R_8F_E4M3_FNUZ ||
-             layout.type() == HIP_R_8F_E5M2 || layout.type() == HIP_R_8F_E4M3;
path: xla/stream_executor/rocm/hip_blas_lt.cc
user: login: pemeliya
body: It does not compile under rocm-63. Besides, I think HIP_R_8F_E5M2 and HIP_R_8F_E4M3 are not supported on ROCM anyway, so should not see these types here. Supported ones are listed in [GemmRewriter](https://github.com/ROCm/xla/blob/ci_v0.5.0_large_workload/xla/service/gpu/transforms/gemm_rewriter.cc#L2213-L2242)
created_at: 2025-05-07T15:55:04Z
author_association: NONE

PR comments:
user: login: hsharsha
created_at: 2025-05-05T11:22:24Z
body: @pemeliya  Is it specific to rocm-jaxlib-v0.5.0 or do you plan to upstream this?
author_association: COLLABORATOR
user: login: pemeliya
created_at: 2025-05-05T12:32:22Z
body: > @pemeliya Is it specific to rocm-jaxlib-v0.5.0 or do you plan to upstream this?

yes actually I have created an upstream PR too: https://github.com/openxla/xla/pull/26057 
let's see what they say..
author_association: NONE
user: login: hsharsha
created_at: 2025-05-05T14:17:28Z
body: cuda has some ptx code in there. Not sure what it does, We just count the number of redzone violations on rocm side.
author_association: COLLABORATOR
user: login: pemeliya
created_at: 2025-05-06T08:26:52Z
body: > cuda has some ptx code in there. Not sure what it does, We just count the number of redzone violations on rocm side.

yeah, CUDA ptx code does exaclty the same, as they mentioned [here](https://github.com/openxla/xla/blob/main/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc#L73-L79), it's just translation of a simple CUDA kernel to PTX. Back to those days when StreamExecutor did not support in-process kernels.  
author_association: NONE