diff --git a/xla/backends/gpu/runtime/BUILD b/xla/backends/gpu/runtime/BUILD index 33cff897288fe..d47dd5e6798bf 100644 --- a/xla/backends/gpu/runtime/BUILD +++ b/xla/backends/gpu/runtime/BUILD @@ -5019,7 +5019,6 @@ xla_test( }, backends = ["gpu"], tags = [ - "cuda-only", "gpu", ], deps = [ @@ -5037,6 +5036,7 @@ xla_test( "//xla/runtime:device_id", "//xla/service:buffer_assignment", "//xla/service:executable", + "//xla/service:platform_util", "//xla/service/gpu:buffer_allocations", "//xla/stream_executor:device_address", "//xla/stream_executor:device_description", @@ -5050,6 +5050,7 @@ xla_test( "//xla/tsl/platform:statusor", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", ], ) @@ -5102,7 +5103,6 @@ xla_test( }, backends = ["gpu"], tags = [ - "cuda-only", "gpu", ], deps = [ @@ -5123,6 +5123,7 @@ xla_test( "//xla/runtime:device_id", "//xla/service:buffer_assignment", "//xla/service:executable", + "//xla/service:platform_util", "//xla/service/gpu:buffer_allocations", "//xla/stream_executor:device_address", "//xla/stream_executor:platform", @@ -5134,6 +5135,7 @@ xla_test( "//xla/tsl/platform:status_macros", "//xla/tsl/platform:statusor", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", ], ) diff --git a/xla/backends/gpu/runtime/buffers_checksum_thunk.cc b/xla/backends/gpu/runtime/buffers_checksum_thunk.cc index 916973676a960..e8b44b38620ce 100644 --- a/xla/backends/gpu/runtime/buffers_checksum_thunk.cc +++ b/xla/backends/gpu/runtime/buffers_checksum_thunk.cc @@ -46,15 +46,11 @@ namespace se = stream_executor; absl::Status BuffersDebugChecksumThunk::Initialize( const InitializeParams& params) { - if (params.executor->GetPlatform()->id() != se::cuda::kCudaPlatformId) { - VLOG(1) - << "Buffer checksumming not supported on non-CUDA platforms, skipping"; - return absl::OkStatus(); - } - if (!params.executor->GetDeviceDescription() + if (params.executor->GetPlatform()->id() == se::cuda::kCudaPlatformId && + !params.executor->GetDeviceDescription() .cuda_compute_capability() .IsAtLeastPascal()) { - VLOG(1) + LOG_FIRST_N(WARNING, 1) << "Buffer checksumming not supported on CUDA architectures older than " "Pascal due to missing atomic fetch_add with system scope, skipping"; return absl::OkStatus(); diff --git a/xla/backends/gpu/runtime/buffers_checksum_thunk_test.cc b/xla/backends/gpu/runtime/buffers_checksum_thunk_test.cc index 6934a147ef857..8a2c367915170 100644 --- a/xla/backends/gpu/runtime/buffers_checksum_thunk_test.cc +++ b/xla/backends/gpu/runtime/buffers_checksum_thunk_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -26,6 +27,7 @@ limitations under the License. #include #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "xla/tsl/platform/status_macros.h" #include "xla/backends/gpu/runtime/buffer_debug_log.pb.h" #include "xla/backends/gpu/runtime/buffer_debug_log_entry_metadata_store.h" @@ -40,6 +42,7 @@ limitations under the License. #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/buffer_allocations.h" +#include "xla/service/platform_util.h" #include "xla/service/service_executable_run_options.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/device_description.h" @@ -103,17 +106,20 @@ class FakeThunk : public Thunk { class BuffersDebugChecksumThunkTest : public ::testing::Test { protected: void SetUp() override { + std::string name = absl::AsciiStrToUpper( + xla::PlatformUtil::CanonicalPlatformName("gpu").value()); TF_ASSERT_OK_AND_ASSIGN(platform_, - se::PlatformManager::PlatformWithName("CUDA")); + se::PlatformManager::PlatformWithName(name)); TF_ASSERT_OK_AND_ASSIGN(executor_, platform_->ExecutorForDevice(0)); TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream(std::nullopt)); allocator_ = std::make_unique( stream_->parent()); - if (!executor_->GetDeviceDescription() - .cuda_compute_capability() - .IsAtLeastPascal()) { + if (const auto* cc = executor_->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + cc != nullptr && !cc->IsAtLeastPascal()) { GTEST_SKIP() << "buffer checksumming is not supported on CUDA architectures " "older than Pascal due to missing atomic fetch_add with " diff --git a/xla/backends/gpu/runtime/buffers_float_check_thunk.cc b/xla/backends/gpu/runtime/buffers_float_check_thunk.cc index f393eb80d3be3..69cbbc0ff00c9 100644 --- a/xla/backends/gpu/runtime/buffers_float_check_thunk.cc +++ b/xla/backends/gpu/runtime/buffers_float_check_thunk.cc @@ -81,15 +81,11 @@ BuffersDebugFloatCheckThunk::BuffersDebugFloatCheckThunk( } absl::Status BuffersDebugFloatCheckThunk::Initialize( const InitializeParams& params) { - if (params.executor->GetPlatform()->id() != se::cuda::kCudaPlatformId) { - VLOG(1) << "Buffer float checking not supported on non-CUDA platforms, " - "skipping"; - return absl::OkStatus(); - } - if (!params.executor->GetDeviceDescription() + if (params.executor->GetPlatform()->id() == se::cuda::kCudaPlatformId && + !params.executor->GetDeviceDescription() .cuda_compute_capability() .IsAtLeastPascal()) { - VLOG(1) + LOG_FIRST_N(WARNING, 1) << "Buffer float checking not supported on CUDA architectures older " "than Pascal due to missing atomic fetch_add with system scope, " "skipping"; diff --git a/xla/backends/gpu/runtime/buffers_float_check_thunk_test.cc b/xla/backends/gpu/runtime/buffers_float_check_thunk_test.cc index 4557bb8ac8bb9..b4de9d5034a68 100644 --- a/xla/backends/gpu/runtime/buffers_float_check_thunk_test.cc +++ b/xla/backends/gpu/runtime/buffers_float_check_thunk_test.cc @@ -22,12 +22,14 @@ limitations under the License. #include #include #include +#include #include #include #include #include #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "xla/tsl/platform/status_macros.h" #include "xla/backends/gpu/runtime/buffer_debug_log.pb.h" #include "xla/backends/gpu/runtime/buffer_debug_log_entry_metadata_store.h" @@ -42,6 +44,7 @@ limitations under the License. #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/buffer_allocations.h" +#include "xla/service/platform_util.h" #include "xla/service/service_executable_run_options.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/gpu/buffer_debug_log.h" @@ -100,17 +103,20 @@ constexpr PrimitiveType kPrimitiveTypeOf = PrimitiveType::F16; class BuffersDebugFloatCheckThunkTest : public ::testing::Test { protected: void SetUp() override { + std::string name = absl::AsciiStrToUpper( + xla::PlatformUtil::CanonicalPlatformName("gpu").value()); TF_ASSERT_OK_AND_ASSIGN(platform_, - se::PlatformManager::PlatformWithName("CUDA")); + se::PlatformManager::PlatformWithName(name)); TF_ASSERT_OK_AND_ASSIGN(executor_, platform_->ExecutorForDevice(0)); TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream(std::nullopt)); allocator_ = std::make_unique( stream_->parent()); - if (!executor_->GetDeviceDescription() - .cuda_compute_capability() - .IsAtLeastPascal()) { + if (const auto* cc = executor_->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + cc != nullptr && !cc->IsAtLeastPascal()) { GTEST_SKIP() << "buffer float checking is not supported on CUDA architectures " "older than Pascal due to missing atomic fetch_add with " diff --git a/xla/stream_executor/cuda/BUILD b/xla/stream_executor/cuda/BUILD index ea2728c8eedb5..b03562eaa18d3 100644 --- a/xla/stream_executor/cuda/BUILD +++ b/xla/stream_executor/cuda/BUILD @@ -431,61 +431,34 @@ cuda_library( cuda_library( name = "buffer_debug_xor_checksum_kernel_cuda", - srcs = ["buffer_debug_xor_checksum_kernel_cuda.cu.cc"], + srcs = [ + "buffer_debug_xor_checksum_kernel_cuda.cu.cc", + "//xla/stream_executor/gpu:buffer_debug_xor_checksum_kernel_lib.cu.h", + ], # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"], tags = [ "cuda-only", "gpu", ], deps = [ + ":cuda_platform", ":cuda_platform_id", "//xla/backends/gpu/runtime:buffer_debug_log_structs", "//xla/stream_executor:kernel_spec", "//xla/stream_executor/gpu:buffer_debug_xor_checksum_kernel", "//xla/stream_executor/gpu:gpu_kernel_registry", - "//xla/tsl/platform:logging", "@com_google_absl//absl/base", "@local_config_cuda//cuda:cuda_headers", ], alwayslink = True, ) -xla_test( - name = "buffer_debug_xor_checksum_kernel_cuda_test", - srcs = ["buffer_debug_xor_checksum_kernel_cuda_test.cc"], - backends = ["gpu"], - tags = ["cuda-only"], - deps = [ - ":buffer_debug_xor_checksum_kernel_cuda", - "//xla/backends/gpu/runtime:buffer_debug_log_structs", - "//xla/stream_executor:device_address", - "//xla/stream_executor:launch_dim", - "//xla/stream_executor:platform", - "//xla/stream_executor:platform_manager", - "//xla/stream_executor:stream", - "//xla/stream_executor:stream_executor_h", - "//xla/stream_executor:stream_executor_memory_allocator", - "//xla/stream_executor:typed_kernel_factory", - "//xla/stream_executor/gpu:buffer_debug_log", - "//xla/stream_executor/gpu:buffer_debug_xor_checksum_kernel", - "//xla/stream_executor/gpu:gpu_kernel_registry", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:status_macros", - "//xla/tsl/platform:statusor", - "@com_google_absl//absl/cleanup", - "@com_google_absl//absl/log", - "@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_googletest//:gtest_main", - ], -) - cuda_library( name = "buffer_debug_float_check_kernel_cuda", - srcs = ["buffer_debug_float_check_kernel_cuda.cu.cc"], + srcs = [ + "buffer_debug_float_check_kernel_cuda.cu.cc", + "//xla/stream_executor/gpu:buffer_debug_float_check_kernel_lib.cu.h", + ], # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"], tags = [ "cuda-only", @@ -506,43 +479,6 @@ cuda_library( alwayslink = True, ) -xla_test( - name = "buffer_debug_float_check_kernel_cuda_test", - srcs = ["buffer_debug_float_check_kernel_cuda_test.cc"], - backends = ["gpu"], - tags = ["cuda-only"], - deps = [ - ":buffer_debug_float_check_kernel_cuda", - "//xla:types", - "//xla/backends/gpu/runtime:buffer_debug_log_structs", - "//xla/backends/gpu/runtime:buffer_debug_log_structs_test_matchers", - "//xla/backends/gpu/runtime:thunk_id", - "//xla/stream_executor:device_address", - "//xla/stream_executor:device_description", - "//xla/stream_executor:launch_dim", - "//xla/stream_executor:platform", - "//xla/stream_executor:platform_manager", - "//xla/stream_executor:stream", - "//xla/stream_executor:stream_executor_address_allocator", - "//xla/stream_executor:stream_executor_h", - "//xla/stream_executor:typed_kernel_factory", - "//xla/stream_executor/gpu:buffer_debug_float_check_kernel", - "//xla/stream_executor/gpu:buffer_debug_log", - "//xla/stream_executor/gpu:gpu_kernel_registry", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:status_macros", - "//xla/tsl/platform:statusor", - "@com_google_absl//absl/cleanup", - "@com_google_absl//absl/log", - "@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_googletest//:gtest_main", - ], -) - cc_library( name = "cudnn_plugin", srcs = ["cuda_dnn.cc"], diff --git a/xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda.cu.cc b/xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda.cu.cc index 31f481506977f..06d9d2b9cf147 100644 --- a/xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda.cu.cc +++ b/xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda.cu.cc @@ -13,50 +13,25 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include -#include #include -#include -#include #include -#include -#include -#include #include "absl/base/casts.h" -#include "Eigen/Core" #include "third_party/gpus/cuda/include/cuda/atomic" -#include "xla/backends/gpu/runtime/buffer_debug_log_structs.h" +#include "third_party/gpus/cuda/include/cuda_fp16.h" +#include "Eigen/Core" #include "xla/stream_executor/cuda/cuda_platform_id.h" #include "xla/stream_executor/gpu/buffer_debug_float_check_kernel.h" #include "xla/stream_executor/gpu/gpu_kernel_registry.h" #include "xla/stream_executor/kernel_spec.h" -#include "xla/util.h" - -namespace se = stream_executor; - -namespace { -using xla::gpu::FloatCheckResult; +namespace stream_executor::gpu { +__device__ static constexpr uint64_t kWarpSize = 32; +} // namespace stream_executor::gpu -// https://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/: -// > CUDA architecture limits the numbers of threads per block (1024 threads -// > per block limit). -static constexpr uint64_t kBlockSize = 1024; -// warpSize is not a compile time constant on all OSS CI builds, but we need it -// to be one for static array initialization. We assert this value matches -// warpSize at runtime. -static constexpr uint64_t kWarpSize = 32; -static constexpr uint64_t kMaxWarpsPerBlock = kBlockSize / kWarpSize; -template -static constexpr uint64_t kElementsPerMemoryAccess = - std::max(16 / sizeof(T), 1); -template -using Chunk = std::array>; +#include "xla/stream_executor/gpu/buffer_debug_float_check_kernel_lib.cu.h" -template -__host__ __device__ static constexpr T kInfinity = - std::numeric_limits::infinity(); +namespace stream_executor::gpu { // __nv_bfloat16 is a distinct type different from Eigen::bfloat16. // CUDA doesn't provide std::numeric_limits<__nv_bfloat16>::infinity(), which @@ -70,329 +45,103 @@ __host__ __device__ static constexpr T kInfinity = // - __nv_bfloat16 and Eigen::bfloat16 are bitwise identical // So we can get a constant this way. template <> -__host__ __device__ constexpr __nv_bfloat16 kInfinity<__nv_bfloat16> = +__device__ constexpr __nv_bfloat16 kInfinity<__nv_bfloat16> = absl::bit_cast<__nv_bfloat16>(kInfinity); // - __half lacks std::numeric_limits specialization, and Eigen::half is not a // literal type (non-constexpr constructors), so we construct infinity from raw // bits. template <> -__host__ __device__ constexpr __half kInfinity<__half> = +__device__ constexpr __half kInfinity<__half> = absl::bit_cast<__half>(uint16_t{0x7C00}); -__device__ unsigned int ThreadIdx() { - return threadIdx.z * blockDim.y * blockDim.x + threadIdx.y * blockDim.x + - threadIdx.x; +template <> +__device__ inline bool IsNan(__nv_bfloat16 v) { + return __isnan(v); } - -__device__ unsigned int BlockIdx() { - return blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + - blockIdx.x; +template <> +__device__ inline bool IsInf(__nv_bfloat16 v) { + return __isinf(v); } - -__device__ inline bool IsNan(float v) { return isnan(v); } -__device__ inline bool IsNan(double v) { return isnan(v); } -__device__ inline bool IsNan(__nv_bfloat16 v) { return __isnan(v); } -__device__ inline bool IsNan(__half v) { return __isnan(v); } -__device__ inline bool IsInf(float v) { return isinf(v); } -__device__ inline bool IsInf(double v) { return isinf(v); } -__device__ inline bool IsInf(__nv_bfloat16 v) { return __isinf(v); } -__device__ inline bool IsInf(__half v) { return __isinf(v); } -__device__ inline bool IsZero(float v) { return v == 0.0f; } -__device__ inline bool IsZero(double v) { return v == 0.0; } +template <> __device__ inline bool IsZero(__nv_bfloat16 v) { return v == __nv_bfloat16(0.0f); } -__device__ inline bool IsZero(__half v) { return v == __half(0.0f); } - -// Reduce a warp worth of values into a single one and have the 0th thread in -// the warp return it. -__device__ uint32_t WarpReduceSum(uint32_t value) { - static constexpr uint32_t kFullMask = ~0; - for (unsigned int offset = 1; offset < kWarpSize; offset <<= 1) { - value += __shfl_down_sync(kFullMask, value, offset); - } - return value; -} - -template -__device__ T FiniteOr(T val, T fallback) { - if (IsNan(val) || IsInf(val)) { - return fallback; - } - return val; -} - -template -__device__ T WarpReduceMin(T value) { - value = FiniteOr(value, kInfinity); - - static constexpr uint32_t kFullMask = ~0; - for (unsigned int offset = 1; offset < kWarpSize; offset <<= 1) { - value = std::min(value, __shfl_down_sync(kFullMask, value, offset)); - } - return value; -} - -template -__device__ T WarpReduceMax(T value) { - value = FiniteOr(value, -kInfinity); - - static constexpr uint32_t kFullMask = ~0; - for (unsigned int offset = 1; offset < kWarpSize; offset <<= 1) { - value = std::max(value, __shfl_down_sync(kFullMask, value, offset)); - } - return value; -} - -// Sum up a block worth of FloatCheckResults into a single one and have the 0th -// thread in the block return it. -__device__ FloatCheckResult BlockReduceResults(uint32_t tid, - FloatCheckResult value) { - assert(kWarpSize == warpSize); - static_assert(kBlockSize == kWarpSize * kMaxWarpsPerBlock); - // Required to do the second warp reduction. - static_assert(kMaxWarpsPerBlock == kWarpSize); - - const size_t warp_idx = tid / kWarpSize; - const size_t lane_idx = tid % kWarpSize; - - value.nan_count = WarpReduceSum(value.nan_count); - value.inf_count = WarpReduceSum(value.inf_count); - value.zero_count = WarpReduceSum(value.zero_count); - value.min_value = WarpReduceMin(value.min_value); - value.max_value = WarpReduceMax(value.max_value); - - __shared__ uint32_t scratch_nan[kMaxWarpsPerBlock]; - __shared__ uint32_t scratch_inf[kMaxWarpsPerBlock]; - __shared__ uint32_t scratch_zero[kMaxWarpsPerBlock]; - __shared__ double scratch_min[kMaxWarpsPerBlock]; - __shared__ double scratch_max[kMaxWarpsPerBlock]; - if (lane_idx == 0) { - scratch_nan[warp_idx] = value.nan_count; - scratch_inf[warp_idx] = value.inf_count; - scratch_zero[warp_idx] = value.zero_count; - scratch_min[warp_idx] = value.min_value; - scratch_max[warp_idx] = value.max_value; - } - - __syncthreads(); - // The first warp reduces the results from all warps. - if (warp_idx == 0) { - value.nan_count = scratch_nan[lane_idx]; - value.inf_count = scratch_inf[lane_idx]; - value.zero_count = scratch_zero[lane_idx]; - value.min_value = scratch_min[lane_idx]; - value.max_value = scratch_max[lane_idx]; - - value.nan_count = WarpReduceSum(value.nan_count); - value.inf_count = WarpReduceSum(value.inf_count); - value.zero_count = WarpReduceSum(value.zero_count); - value.min_value = WarpReduceMin(value.min_value); - value.max_value = WarpReduceMax(value.max_value); - } else { - value.nan_count = 0; - value.inf_count = 0; - value.zero_count = 0; - value.min_value = kInfinity; - value.max_value = -kInfinity; - } - - return value; -} - -// Get a part of the input buffer current thread block is responsible for -// processing, assuming the load is spread up to max_blocks across the entire -// grid. If max_blocks is not provided, the entire grid is used. -template -__device__ inline std::tuple GetBlockInput( - const T* input, uint64_t input_size, - std::optional max_blocks = std::nullopt) { - size_t grid_size = gridDim.x * gridDim.y * gridDim.z; - if (max_blocks.has_value()) { - grid_size = std::min(grid_size, *max_blocks); - } - const uint64_t max_block_input_size = xla::RoundUpTo( - xla::CeilOfRatio(input_size, grid_size), kElementsPerMemoryAccess); - const uint64_t block_input_offset = BlockIdx() * max_block_input_size; - const uint64_t block_input_size = std::min( - max_block_input_size, - input_size >= block_input_offset ? input_size - block_input_offset : 0); - return {input + block_input_offset, block_input_size}; +template <> +__device__ inline bool IsNan(__half v) { + return __isnan(v); } - -template -__device__ FloatCheckResult CheckFloats(const T* input, uint64_t input_size, - uint64_t max_blocks) { - const unsigned int tid = ThreadIdx(); - const auto [block_input, block_input_size] = - GetBlockInput(input, input_size, max_blocks); - - const Chunk* chunked_input = - reinterpret_cast*>(block_input); - const uint64_t input_chunks = - xla::FloorOfRatio(block_input_size, kElementsPerMemoryAccess); - // This may be less than block_input_size only for the last block. - const uint64_t chunked_input_size = - xla::RoundDownTo(block_input_size, kElementsPerMemoryAccess); - - FloatCheckResult result{}; - result.min_value = kInfinity; - result.max_value = -kInfinity; - - for (uint64_t i = tid; i < input_chunks; i += kBlockSize) { - Chunk values = chunked_input[i]; - for (const T value : values) { - result.nan_count += IsNan(value); - result.inf_count += IsInf(value); - result.zero_count += IsZero(value); - result.min_value = std::min( - result.min_value, static_cast(FiniteOr(value, kInfinity))); - result.max_value = - std::max(result.max_value, - static_cast(FiniteOr(value, -kInfinity))); - } - } - - if (tid == 0 && chunked_input_size < block_input_size) { - const size_t rest = block_input_size - chunked_input_size; - for (uint64_t j = 0; j < rest; ++j) { - const T value = block_input[chunked_input_size + j]; - result.nan_count += IsNan(value); - result.inf_count += IsInf(value); - result.zero_count += IsZero(value); - result.min_value = std::min( - result.min_value, static_cast(FiniteOr(value, kInfinity))); - result.max_value = - std::max(result.max_value, - static_cast(FiniteOr(value, -kInfinity))); - } - } - - return BlockReduceResults(tid, result); +template <> +__device__ inline bool IsInf(__half v) { + return __isinf(v); } - -__device__ FloatCheckResult ReduceResults(const FloatCheckResult* input, - uint64_t input_size) { - const unsigned int tid = ThreadIdx(); - const auto [block_input, block_input_size] = GetBlockInput(input, input_size); - - FloatCheckResult result{}; - result.min_value = kInfinity; - result.max_value = -kInfinity; - - for (uint64_t i = tid; i < input_size; i += kBlockSize) { - const FloatCheckResult value = block_input[i]; - result.nan_count += value.nan_count; - result.inf_count += value.inf_count; - result.zero_count += value.zero_count; - result.min_value = std::min(result.min_value, value.min_value); - result.max_value = std::max(result.max_value, value.max_value); - } - - // Now reduce a block worth of values into a single one. - return BlockReduceResults(tid, result); +template <> +__device__ inline bool IsZero(__half v) { + return v == __half(0.0f); } -// Count the number of floats for NaNs, Infs and zeros in input buffer and store -// partially accumulated results in the tmp array. template -__global__ void FloatCheck(const T* input, uint64_t input_size, - xla::gpu::FloatCheckResult* tmp, uint64_t tmp_size) { - assert(blockDim.x * blockDim.y * blockDim.z == kBlockSize); - assert(BlockIdx() < tmp_size); - if (BlockIdx() >= tmp_size) { - return; - } - - FloatCheckResult result = CheckFloats(input, input_size, tmp_size); - if (ThreadIdx() == 0) { - tmp[BlockIdx()] = result; - } +__device__ inline T WarpShuffleDown(T value, unsigned int offset) { + static constexpr uint32_t kFullMask = ~0u; + return __shfl_down_sync(kFullMask, value, offset); } -// Reduce the partially accumulated results from `FloatCheck` invocations and -// append the result to the buffer debug log. -__global__ void ReduceFloatCheckResults( - xla::gpu::FloatCheckResult* tmp, uint64_t tmp_size, - xla::gpu::BufferDebugLogEntryId entry_id, - xla::gpu::BufferDebugLogHeader* log_header, - xla::gpu::BufferDebugFloatCheckEntry* log_entries) { - assert(blockDim.x * blockDim.y * blockDim.z == kBlockSize); - assert(BlockIdx() == 0); - if (BlockIdx() >= 1) { - return; - } - - assert(tmp_size > 0); - FloatCheckResult total = ReduceResults(tmp, tmp_size); - - if (BlockIdx() == 0 && ThreadIdx() == 0) { - cuda::atomic_ref log_write_idx( - log_header->write_idx); +__device__ inline uint32_t AtomicIncSystem(uint32_t* write_idx) { #if __CUDA_ARCH__ >= 600 - const uint32_t write_idx = log_write_idx.fetch_add(1); - if (write_idx < log_header->capacity) { - log_entries[write_idx] = - xla::gpu::BufferDebugFloatCheckEntry{entry_id, total}; - } + ::cuda::atomic_ref log_write_idx( + *write_idx); + return log_write_idx.fetch_add(1); #else - // Our toolchains generate a fetch_add PTX instructions with system scope, - // which is not supported on pre-Pascal architectures. - (void)total; - assert(false); + // Our toolchains generate a fetch_add PTX instructions with system scope, + // which is not supported on pre-Pascal architectures. + assert(false); + return ~0u; #endif - } -} - -se::KernelLoaderSpec GetFloatCheckF32KernelSpec(int arity) { - return se::KernelLoaderSpec::CreateInProcessSymbolSpec( - absl::bit_cast(&FloatCheck), - "BufferDebugFloatCheckF32Kernel", arity); -} - -se::KernelLoaderSpec GetFloatCheckBf16KernelSpec(int arity) { - return se::KernelLoaderSpec::CreateInProcessSymbolSpec( - absl::bit_cast(&FloatCheck<__nv_bfloat16>), - "BufferDebugFloatCheckBf16Kernel", arity); -} - -se::KernelLoaderSpec GetFloatCheckF64KernelSpec(int arity) { - return se::KernelLoaderSpec::CreateInProcessSymbolSpec( - absl::bit_cast(&FloatCheck), - "BufferDebugFloatCheckF64Kernel", arity); -} - -se::KernelLoaderSpec GetFloatCheckF16KernelSpec(int arity) { - return se::KernelLoaderSpec::CreateInProcessSymbolSpec( - absl::bit_cast(&FloatCheck<__half>), - "BufferDebugFloatCheckF16Kernel", arity); -} - -se::KernelLoaderSpec GetReduceFloatCheckResultsKernelSpec(int arity) { - return se::KernelLoaderSpec::CreateInProcessSymbolSpec( - absl::bit_cast(&ReduceFloatCheckResults), - "BufferDebugReduceFloatCheckResultsKernel", arity); } -} // namespace +} // namespace stream_executor::gpu GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( - BufferDebugFloatCheckF32Kernel, se::gpu::BufferDebugFloatCheckF32Kernel, - se::cuda::kCudaPlatformId, GetFloatCheckF32KernelSpec); + BufferDebugFloatCheckF32Kernel, + stream_executor::gpu::BufferDebugFloatCheckF32Kernel, + stream_executor::cuda::kCudaPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::FloatCheck), + "BufferDebugFloatCheckF32Kernel", arity); + })); GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( - BufferDebugFloatCheckBf16Kernel, se::gpu::BufferDebugFloatCheckBf16Kernel, - se::cuda::kCudaPlatformId, GetFloatCheckBf16KernelSpec); + BufferDebugFloatCheckBf16Kernel, + stream_executor::gpu::BufferDebugFloatCheckBf16Kernel, + stream_executor::cuda::kCudaPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast( + &stream_executor::gpu::FloatCheck<__nv_bfloat16>), + "BufferDebugFloatCheckBf16Kernel", arity); + })); GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( - BufferDebugFloatCheckF64Kernel, se::gpu::BufferDebugFloatCheckF64Kernel, - se::cuda::kCudaPlatformId, GetFloatCheckF64KernelSpec); + BufferDebugFloatCheckF64Kernel, + stream_executor::gpu::BufferDebugFloatCheckF64Kernel, + stream_executor::cuda::kCudaPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::FloatCheck), + "BufferDebugFloatCheckF64Kernel", arity); + })); GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( - BufferDebugFloatCheckF16Kernel, se::gpu::BufferDebugFloatCheckF16Kernel, - se::cuda::kCudaPlatformId, GetFloatCheckF16KernelSpec); + BufferDebugFloatCheckF16Kernel, + stream_executor::gpu::BufferDebugFloatCheckF16Kernel, + stream_executor::cuda::kCudaPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::FloatCheck<__half>), + "BufferDebugFloatCheckF16Kernel", arity); + })); GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( BufferDebugReduceFloatCheckResultsKernel, - se::gpu::BufferDebugAppendReducedFloatCheckResultsKernel, - se::cuda::kCudaPlatformId, GetReduceFloatCheckResultsKernelSpec); + stream_executor::gpu::BufferDebugAppendReducedFloatCheckResultsKernel, + stream_executor::cuda::kCudaPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::ReduceFloatCheckResults), + "BufferDebugReduceFloatCheckResultsKernel", arity); + })); diff --git a/xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda.cu.cc b/xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda.cu.cc index 412441afb6192..8aac6186d7066 100644 --- a/xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda.cu.cc +++ b/xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda.cu.cc @@ -13,196 +13,39 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include -#include -#include #include "absl/base/casts.h" #include "third_party/gpus/cuda/include/cuda/atomic" -#include "xla/backends/gpu/runtime/buffer_debug_log_structs.h" #include "xla/stream_executor/cuda/cuda_platform_id.h" #include "xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel.h" +#include "xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_lib.cu.h" #include "xla/stream_executor/gpu/gpu_kernel_registry.h" #include "xla/stream_executor/kernel_spec.h" -#include "xla/tsl/platform/logging.h" -namespace se = stream_executor; +namespace stream_executor::gpu { -namespace { - -__device__ unsigned int ThreadIdx() { - return threadIdx.z * blockDim.y * blockDim.x + threadIdx.y * blockDim.x + - threadIdx.x; -} - -__device__ unsigned int BlockIdx() { - return blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + - blockIdx.x; -} - -// Based on -// https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf -template -__device__ void WarpReduceXor(unsigned int tid, volatile uint32_t* data) { - if (BLOCK_SIZE >= 64) data[tid] ^= data[tid + 32]; - if (BLOCK_SIZE >= 32) data[tid] ^= data[tid + 16]; - if (BLOCK_SIZE >= 16) data[tid] ^= data[tid + 8]; - if (BLOCK_SIZE >= 8) data[tid] ^= data[tid + 4]; - if (BLOCK_SIZE >= 4) data[tid] ^= data[tid + 2]; - if (BLOCK_SIZE >= 2) data[tid] ^= data[tid + 1]; -} - -// Calculates a XOR of all elements of `input` and puts the result in `output`. -// -// Optimized implementation based on -// https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf -// that takes advantage of `BLOCK_SIZE` threads. -// -// `BLOCK_SIZE` must be a power of 2 no larger than 1024. -template -__device__ void ReduceXor(const uint32_t* input, uint64_t input_size, - uint32_t* output) { - __shared__ uint32_t scratch[BLOCK_SIZE]; - - assert(BlockIdx() == 0); - const unsigned int tid = ThreadIdx(); - - scratch[tid] = 0; - for (unsigned int i = tid; i < input_size; i += BLOCK_SIZE) { - scratch[tid] ^= input[i]; - } - - __syncthreads(); - - if (BLOCK_SIZE >= 1024) { - if (tid < 512) { - scratch[tid] ^= scratch[tid + 512]; - } - __syncthreads(); - } - if (BLOCK_SIZE >= 512) { - if (tid < 256) { - scratch[tid] ^= scratch[tid + 256]; - } - __syncthreads(); - } - if (BLOCK_SIZE >= 256) { - if (tid < 128) { - scratch[tid] ^= scratch[tid + 128]; - } - __syncthreads(); - } - if (BLOCK_SIZE >= 128) { - if (tid < 64) { - scratch[tid] ^= scratch[tid + 64]; - } - __syncthreads(); - } - if (tid < 32) WarpReduceXor(tid, scratch); - if (tid == 0) *output = scratch[0]; -} - -// Attempts to append the checksum of the `input` buffer to the `log_entries`, -// using `log_header` to track available capacity and used space. -// -// The log entry is tagged with `entry_id`. The checksum is parallelized as -// much as block dimensions allow it. -// -// If the log does not have enough space for the new entry, the entry is -// discarded. -// -// `input_size` is the size of the input buffer in bytes. -// -// LIMITATIONS: -// - Only a single thread block is supported. -// - Block dimensions must be a power of 2. -__global__ void AppendChecksum(xla::gpu::BufferDebugLogEntryId entry_id, - const uint8_t* input, uint64_t input_size, - xla::gpu::BufferDebugLogHeader* log_header, - xla::gpu::BufferDebugLogEntry* log_entries) { - const uint32_t block_size = blockDim.x * blockDim.y * blockDim.z; - const uint32_t* input_u32 = reinterpret_cast(input); - const uint64_t input_u32_size = input_size / sizeof(uint32_t); - uint32_t checksum = 0; - - assert(gridDim.x == 1 && gridDim.y == 1 && gridDim.z == 1); - if (BlockIdx() != 0) { - return; - } - - // https://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/: - // > CUDA architecture limits the numbers of threads per block (1024 threads - // > per block limit). - switch (block_size) { - case 1024: - ReduceXor<1024>(input_u32, input_u32_size, &checksum); - break; - case 512: - ReduceXor<512>(input_u32, input_u32_size, &checksum); - break; - case 256: - ReduceXor<256>(input_u32, input_u32_size, &checksum); - break; - case 128: - ReduceXor<128>(input_u32, input_u32_size, &checksum); - break; - case 64: - ReduceXor<64>(input_u32, input_u32_size, &checksum); - break; - case 32: - ReduceXor<32>(input_u32, input_u32_size, &checksum); - break; - case 16: - ReduceXor<16>(input_u32, input_u32_size, &checksum); - break; - case 8: - ReduceXor<8>(input_u32, input_u32_size, &checksum); - break; - case 4: - ReduceXor<4>(input_u32, input_u32_size, &checksum); - break; - case 2: - ReduceXor<2>(input_u32, input_u32_size, &checksum); - break; - case 1: - ReduceXor<1>(input_u32, input_u32_size, &checksum); - break; - default: - // Unsupported block size. - assert(false); - return; - } - - if (ThreadIdx() == 0) { - const size_t last_chunk_size = input_size % sizeof(uint32_t); - uint32_t last_chunk = 0; - memcpy(&last_chunk, input + input_u32_size * sizeof(uint32_t), - last_chunk_size); - checksum ^= last_chunk; - - cuda::atomic_ref - checksum_log_write_idx(log_header->write_idx); +__device__ inline uint32_t AtomicIncSystem(uint32_t* write_idx) { #if __CUDA_ARCH__ >= 600 - const uint32_t write_idx = checksum_log_write_idx.fetch_add(1); - if (write_idx < log_header->capacity) { - log_entries[write_idx] = {entry_id, checksum}; - } + ::cuda::atomic_ref log_write_idx( + *write_idx); + return log_write_idx.fetch_add(1); #else - // Our toolchains generate a fetch_add PTX instructions with system scope, - // which is not supported on pre-Pascal architectures. - assert(false); + // Our toolchains generate a fetch_add PTX instructions with system scope, + // which is not supported on pre-Pascal architectures. + assert(false); + return ~0u; #endif - } -} - -se::KernelLoaderSpec GetChecksumKernelSpec(int arity) { - return se::KernelLoaderSpec::CreateInProcessSymbolSpec( - absl::bit_cast(&AppendChecksum), "BufferDebugXorChecksumKernel", - arity); } -} // namespace +} // namespace stream_executor::gpu GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( - BufferDebugXorChecksumKernel, se::gpu::BufferDebugXorChecksumKernel, - se::cuda::kCudaPlatformId, GetChecksumKernelSpec); + BufferDebugXorChecksumKernel, + stream_executor::gpu::BufferDebugXorChecksumKernel, + stream_executor::cuda::kCudaPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::AppendChecksum), + "BufferDebugXorChecksumKernel", arity); + })); diff --git a/xla/stream_executor/gpu/BUILD b/xla/stream_executor/gpu/BUILD index 4add57a4704a2..74f811d3d0583 100644 --- a/xla/stream_executor/gpu/BUILD +++ b/xla/stream_executor/gpu/BUILD @@ -780,6 +780,8 @@ cc_library( exports_files([ "buffer_comparator_kernel_lib.cu.h", + "buffer_debug_float_check_kernel_lib.cu.h", + "buffer_debug_xor_checksum_kernel_lib.cu.h", "collective_signal.cu.h", "all_reduce_kernel_lib.cu.h", "multi_gpu_barrier_kernel.cu.h", @@ -962,6 +964,40 @@ cc_library( ], ) +xla_test( + name = "buffer_debug_xor_checksum_kernel_test", + srcs = ["buffer_debug_xor_checksum_kernel_test.cc"], + backends = ["gpu"], + tags = ["no-oneapi"], + deps = [ + ":buffer_debug_log", + ":buffer_debug_xor_checksum_kernel", + ":gpu_kernel_registry", + "//xla/backends/gpu/runtime:buffer_debug_log_structs", + "//xla/service:platform_util", + "//xla/stream_executor:device_address", + "//xla/stream_executor:device_description", + "//xla/stream_executor:launch_dim", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream", + "//xla/stream_executor:stream_executor_address_allocator", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor:typed_kernel_factory", + "//xla/tsl/lib/core:status_test_util", + "//xla/tsl/platform:errors", + "//xla/tsl/platform:status_macros", + "//xla/tsl/platform:statusor", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "buffer_debug_float_check_kernel", hdrs = ["buffer_debug_float_check_kernel.h"], @@ -973,6 +1009,44 @@ cc_library( ], ) +xla_test( + name = "buffer_debug_float_check_kernel_test", + srcs = ["buffer_debug_float_check_kernel_test.cc"], + backends = ["gpu"], + tags = ["no-oneapi"], + deps = [ + ":buffer_debug_float_check_kernel", + ":buffer_debug_log", + ":gpu_kernel_registry", + "//xla:types", + "//xla/backends/gpu/runtime:buffer_debug_log_structs", + "//xla/backends/gpu/runtime:buffer_debug_log_structs_test_matchers", + "//xla/backends/gpu/runtime:thunk_id", + "//xla/service:platform_util", + "//xla/stream_executor:device_address", + "//xla/stream_executor:device_description", + "//xla/stream_executor:launch_dim", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream", + "//xla/stream_executor:stream_executor_address_allocator", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor:typed_kernel_factory", + "//xla/tsl/lib/core:status_test_util", + "//xla/tsl/platform:errors", + "//xla/tsl/platform:status_macros", + "//xla/tsl/platform:statusor", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "prefix_sum_kernel", hdrs = ["prefix_sum_kernel.h"], diff --git a/xla/stream_executor/gpu/buffer_debug_float_check_kernel_lib.cu.h b/xla/stream_executor/gpu/buffer_debug_float_check_kernel_lib.cu.h new file mode 100644 index 0000000000000..1fdac7ad574b3 --- /dev/null +++ b/xla/stream_executor/gpu/buffer_debug_float_check_kernel_lib.cu.h @@ -0,0 +1,313 @@ +/* Copyright 2025 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_BUFFER_DEBUG_FLOAT_CHECK_KERNEL_LIB_CU_H_ +#define XLA_STREAM_EXECUTOR_GPU_BUFFER_DEBUG_FLOAT_CHECK_KERNEL_LIB_CU_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "xla/backends/gpu/runtime/buffer_debug_log_structs.h" +#include "xla/util.h" + +namespace stream_executor::gpu { + +using xla::gpu::FloatCheckResult; + +// https://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/: +// > CUDA architecture limits the numbers of threads per block (1024 threads +// > per block limit). +static constexpr uint64_t kBlockSize = 1024; +// `kWarpSize` must be defined by the platform-specific translation unit (in +// namespace stream_executor::gpu) before including this header. warpSize is not +// a compile time constant on all OSS CI builds, but we need it to be one for +// static array initialization, and it differs between CUDA and ROCm. +__device__ static constexpr uint64_t kMaxWarpsPerBlock = kBlockSize / kWarpSize; +template +static constexpr uint64_t kElementsPerMemoryAccess = + std::max(16 / sizeof(T), 1); +template +using Chunk = std::array>; + +template +__device__ static constexpr T kInfinity = std::numeric_limits::infinity(); + +template +__device__ inline bool IsNan(T v) { + return isnan(v); +} +template +__device__ inline bool IsInf(T v) { + return isinf(v); +} +template +__device__ inline bool IsZero(T v) { + return v == T(0); +} + +__device__ inline uint32_t AtomicIncSystem(uint32_t* write_idx); + +__device__ inline unsigned int ThreadIdx() { + return threadIdx.z * blockDim.y * blockDim.x + threadIdx.y * blockDim.x + + threadIdx.x; +} + +__device__ inline unsigned int BlockIdx() { + return blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + + blockIdx.x; +} + +template +__device__ T FiniteOr(T val, T fallback) { + if (IsNan(val) || IsInf(val)) { + return fallback; + } + return val; +} + +template +__device__ inline T WarpShuffleDown(T value, unsigned int offset); + +// Reduce `subgroup_size` lanes worth of values into a single one and have the +// 0th lane of the subgroup return it. +__device__ inline uint32_t WarpSubgroupReduceSum(uint32_t value, + uint64_t subgroup_size) { + for (unsigned int offset = 1; offset < subgroup_size; offset <<= 1) { + value += WarpShuffleDown(value, offset); + } + return value; +} + +template +__device__ T WarpSubgroupReduceMin(T value, uint64_t subgroup_size) { + value = FiniteOr(value, kInfinity); + + for (unsigned int offset = 1; offset < subgroup_size; offset <<= 1) { + value = std::min(value, WarpShuffleDown(value, offset)); + } + return value; +} + +template +__device__ T WarpSubgroupReduceMax(T value, uint64_t subgroup_size) { + value = FiniteOr(value, -kInfinity); + + for (unsigned int offset = 1; offset < subgroup_size; offset <<= 1) { + value = std::max(value, WarpShuffleDown(value, offset)); + } + return value; +} + +// Sum up a block worth of FloatCheckResults into a single one and have the 0th +// thread in the block return it. +__device__ inline FloatCheckResult BlockReduceResults(uint32_t tid, + FloatCheckResult value) { + assert(kWarpSize == warpSize); + static_assert(kBlockSize == kWarpSize * kMaxWarpsPerBlock); + + const size_t warp_idx = tid / kWarpSize; + const size_t lane_idx = tid % kWarpSize; + + value.nan_count = WarpSubgroupReduceSum(value.nan_count, kWarpSize); + value.inf_count = WarpSubgroupReduceSum(value.inf_count, kWarpSize); + value.zero_count = WarpSubgroupReduceSum(value.zero_count, kWarpSize); + value.min_value = WarpSubgroupReduceMin(value.min_value, kWarpSize); + value.max_value = WarpSubgroupReduceMax(value.max_value, kWarpSize); + + __shared__ uint32_t scratch_nan[kMaxWarpsPerBlock]; + __shared__ uint32_t scratch_inf[kMaxWarpsPerBlock]; + __shared__ uint32_t scratch_zero[kMaxWarpsPerBlock]; + __shared__ double scratch_min[kMaxWarpsPerBlock]; + __shared__ double scratch_max[kMaxWarpsPerBlock]; + if (lane_idx == 0) { + scratch_nan[warp_idx] = value.nan_count; + scratch_inf[warp_idx] = value.inf_count; + scratch_zero[warp_idx] = value.zero_count; + scratch_min[warp_idx] = value.min_value; + scratch_max[warp_idx] = value.max_value; + } + + __syncthreads(); + // The first kMaxWarpsPerBlock lanes of first warp reduces the results from + // all warps. + if (warp_idx == 0) { + const size_t scratch_idx = lane_idx % kMaxWarpsPerBlock; + value.nan_count = scratch_nan[scratch_idx]; + value.inf_count = scratch_inf[scratch_idx]; + value.zero_count = scratch_zero[scratch_idx]; + value.min_value = scratch_min[scratch_idx]; + value.max_value = scratch_max[scratch_idx]; + + value.nan_count = WarpSubgroupReduceSum(value.nan_count, kMaxWarpsPerBlock); + value.inf_count = WarpSubgroupReduceSum(value.inf_count, kMaxWarpsPerBlock); + value.zero_count = + WarpSubgroupReduceSum(value.zero_count, kMaxWarpsPerBlock); + value.min_value = WarpSubgroupReduceMin(value.min_value, kMaxWarpsPerBlock); + value.max_value = WarpSubgroupReduceMax(value.max_value, kMaxWarpsPerBlock); + } else { + value.nan_count = 0; + value.inf_count = 0; + value.zero_count = 0; + value.min_value = kInfinity; + value.max_value = -kInfinity; + } + + return value; +} + +// Get a part of the input buffer current thread block is responsible for +// processing, assuming the load is spread up to max_blocks across the entire +// grid. If max_blocks is not provided, the entire grid is used. +template +__device__ inline std::tuple GetBlockInput( + const T* input, uint64_t input_size, + std::optional max_blocks = std::nullopt) { + size_t grid_size = gridDim.x * gridDim.y * gridDim.z; + if (max_blocks.has_value()) { + grid_size = std::min(grid_size, *max_blocks); + } + const uint64_t max_block_input_size = xla::RoundUpTo( + xla::CeilOfRatio(input_size, grid_size), kElementsPerMemoryAccess); + const uint64_t block_input_offset = BlockIdx() * max_block_input_size; + const uint64_t block_input_size = std::min( + max_block_input_size, + input_size >= block_input_offset ? input_size - block_input_offset : 0); + return {input + block_input_offset, block_input_size}; +} + +template +__device__ FloatCheckResult CheckFloats(const T* input, uint64_t input_size, + uint64_t max_blocks) { + const unsigned int tid = ThreadIdx(); + const auto [block_input, block_input_size] = + GetBlockInput(input, input_size, max_blocks); + + const Chunk* chunked_input = + reinterpret_cast*>(block_input); + const uint64_t input_chunks = + xla::FloorOfRatio(block_input_size, kElementsPerMemoryAccess); + // This may be less than block_input_size only for the last block. + const uint64_t chunked_input_size = + xla::RoundDownTo(block_input_size, kElementsPerMemoryAccess); + + FloatCheckResult result{}; + result.min_value = kInfinity; + result.max_value = -kInfinity; + + for (uint64_t i = tid; i < input_chunks; i += kBlockSize) { + Chunk values = chunked_input[i]; + for (const T value : values) { + result.nan_count += IsNan(value); + result.inf_count += IsInf(value); + result.zero_count += IsZero(value); + result.min_value = std::min( + result.min_value, static_cast(FiniteOr(value, kInfinity))); + result.max_value = + std::max(result.max_value, + static_cast(FiniteOr(value, -kInfinity))); + } + } + + if (tid == 0 && chunked_input_size < block_input_size) { + const size_t rest = block_input_size - chunked_input_size; + for (uint64_t j = 0; j < rest; ++j) { + const T value = block_input[chunked_input_size + j]; + result.nan_count += IsNan(value); + result.inf_count += IsInf(value); + result.zero_count += IsZero(value); + result.min_value = std::min( + result.min_value, static_cast(FiniteOr(value, kInfinity))); + result.max_value = + std::max(result.max_value, + static_cast(FiniteOr(value, -kInfinity))); + } + } + + return BlockReduceResults(tid, result); +} + +__device__ inline FloatCheckResult ReduceResults(const FloatCheckResult* input, + uint64_t input_size) { + const unsigned int tid = ThreadIdx(); + const auto [block_input, block_input_size] = GetBlockInput(input, input_size); + + FloatCheckResult result{}; + result.min_value = kInfinity; + result.max_value = -kInfinity; + + for (uint64_t i = tid; i < input_size; i += kBlockSize) { + const FloatCheckResult value = block_input[i]; + result.nan_count += value.nan_count; + result.inf_count += value.inf_count; + result.zero_count += value.zero_count; + result.min_value = std::min(result.min_value, value.min_value); + result.max_value = std::max(result.max_value, value.max_value); + } + + // Now reduce a block worth of values into a single one. + return BlockReduceResults(tid, result); +} + +// Count the number of floats for NaNs, Infs and zeros in input buffer and store +// partially accumulated results in the tmp array. +template +__global__ void FloatCheck(const T* input, uint64_t input_size, + xla::gpu::FloatCheckResult* tmp, uint64_t tmp_size) { + assert(blockDim.x * blockDim.y * blockDim.z == kBlockSize); + assert(BlockIdx() < tmp_size); + if (BlockIdx() >= tmp_size) { + return; + } + + FloatCheckResult result = CheckFloats(input, input_size, tmp_size); + if (ThreadIdx() == 0) { + tmp[BlockIdx()] = result; + } +} + +// Reduce the partially accumulated results from `FloatCheck` invocations and +// append the result to the buffer debug log. +__global__ void ReduceFloatCheckResults( + xla::gpu::FloatCheckResult* tmp, uint64_t tmp_size, + xla::gpu::BufferDebugLogEntryId entry_id, + xla::gpu::BufferDebugLogHeader* log_header, + xla::gpu::BufferDebugFloatCheckEntry* log_entries) { + assert(blockDim.x * blockDim.y * blockDim.z == kBlockSize); + assert(BlockIdx() == 0); + if (BlockIdx() >= 1) { + return; + } + + assert(tmp_size > 0); + FloatCheckResult total = ReduceResults(tmp, tmp_size); + + if (BlockIdx() == 0 && ThreadIdx() == 0) { + const uint32_t write_idx = AtomicIncSystem(&log_header->write_idx); + if (write_idx < log_header->capacity) { + log_entries[write_idx] = + xla::gpu::BufferDebugFloatCheckEntry{entry_id, total}; + } + } +} + +} // namespace stream_executor::gpu + +#endif // XLA_STREAM_EXECUTOR_GPU_BUFFER_DEBUG_FLOAT_CHECK_KERNEL_LIB_CU_H_ diff --git a/xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda_test.cc b/xla/stream_executor/gpu/buffer_debug_float_check_kernel_test.cc similarity index 96% rename from xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda_test.cc rename to xla/stream_executor/gpu/buffer_debug_float_check_kernel_test.cc index 868e866b0e7e1..32c807874e273 100644 --- a/xla/stream_executor/cuda/buffer_debug_float_check_kernel_cuda_test.cc +++ b/xla/stream_executor/gpu/buffer_debug_float_check_kernel_test.cc @@ -13,12 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/stream_executor/gpu/buffer_debug_float_check_kernel.h" + #include #include #include #include #include #include +#include #include #include @@ -27,15 +30,16 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "xla/tsl/platform/status_macros.h" #include "xla/backends/gpu/runtime/buffer_debug_log_structs.h" #include "xla/backends/gpu/runtime/buffer_debug_log_structs_test_matchers.h" #include "xla/backends/gpu/runtime/thunk_id.h" +#include "xla/service/platform_util.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/device_description.h" -#include "xla/stream_executor/gpu/buffer_debug_float_check_kernel.h" #include "xla/stream_executor/gpu/buffer_debug_log.h" #include "xla/stream_executor/gpu/gpu_kernel_registry.h" #include "xla/stream_executor/launch_dim.h" @@ -52,7 +56,7 @@ limitations under the License. namespace se = stream_executor; -namespace stream_executor::cuda { +namespace stream_executor { namespace { using xla::gpu::BufferDebugFloatCheckEntry; @@ -67,16 +71,19 @@ using xla::gpu::ZeroCountIs; class FloatCheckKernelTest : public ::testing::Test { protected: void SetUp() override { + std::string name = absl::AsciiStrToUpper( + xla::PlatformUtil::CanonicalPlatformName("gpu").value()); TF_ASSERT_OK_AND_ASSIGN(platform_, - se::PlatformManager::PlatformWithName("CUDA")); + se::PlatformManager::PlatformWithName(name)); TF_ASSERT_OK_AND_ASSIGN(executor_, platform_->ExecutorForDevice(0)); TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream(std::nullopt)); allocator_ = std::make_unique(stream_->parent()); - if (!executor_->GetDeviceDescription() - .cuda_compute_capability() - .IsAtLeastPascal()) { + if (const auto* cc = executor_->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + cc != nullptr && !cc->IsAtLeastPascal()) { GTEST_SKIP() << "Buffer checking is not supported on CUDA architectures older " "than Pascal due to missing atomic fetch_add with system scope"; @@ -360,4 +367,4 @@ TYPED_TEST(FloatCheckKernelTypedTest, ReduceFloatCheckResults) { } } // namespace -} // namespace stream_executor::cuda +} // namespace stream_executor diff --git a/xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_lib.cu.h b/xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_lib.cu.h new file mode 100644 index 0000000000000..0f86eb6e74e47 --- /dev/null +++ b/xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_lib.cu.h @@ -0,0 +1,191 @@ +/* Copyright 2025 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_BUFFER_DEBUG_XOR_CHECKSUM_KERNEL_LIB_CU_H_ +#define XLA_STREAM_EXECUTOR_GPU_BUFFER_DEBUG_XOR_CHECKSUM_KERNEL_LIB_CU_H_ + +#include +#include +#include +#include + +#include "xla/backends/gpu/runtime/buffer_debug_log_structs.h" + +namespace stream_executor::gpu { + +// Atomically increments `*write_idx` with system scope and returns the previous +// value. Defined by the platform-specific translation unit, as the intrinsic +// differs between CUDA and ROCm. +__device__ inline uint32_t AtomicIncSystem(uint32_t* write_idx); + +__device__ inline unsigned int ThreadIdx() { + return threadIdx.z * blockDim.y * blockDim.x + threadIdx.y * blockDim.x + + threadIdx.x; +} + +__device__ inline unsigned int BlockIdx() { + return blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + + blockIdx.x; +} + +// Based on +// https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf +template +__device__ void WarpReduceXor(unsigned int tid, volatile uint32_t* data) { + if (BLOCK_SIZE >= 64) data[tid] ^= data[tid + 32]; + if (BLOCK_SIZE >= 32) data[tid] ^= data[tid + 16]; + if (BLOCK_SIZE >= 16) data[tid] ^= data[tid + 8]; + if (BLOCK_SIZE >= 8) data[tid] ^= data[tid + 4]; + if (BLOCK_SIZE >= 4) data[tid] ^= data[tid + 2]; + if (BLOCK_SIZE >= 2) data[tid] ^= data[tid + 1]; +} + +// Calculates a XOR of all elements of `input` and puts the result in `output`. +// +// Optimized implementation based on +// https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf +// that takes advantage of `BLOCK_SIZE` threads. +// +// `BLOCK_SIZE` must be a power of 2 no larger than 1024. +template +__device__ void ReduceXor(const uint32_t* input, uint64_t input_size, + uint32_t* output) { + __shared__ uint32_t scratch[BLOCK_SIZE]; + + assert(BlockIdx() == 0); + const unsigned int tid = ThreadIdx(); + + scratch[tid] = 0; + for (unsigned int i = tid; i < input_size; i += BLOCK_SIZE) { + scratch[tid] ^= input[i]; + } + + __syncthreads(); + + if (BLOCK_SIZE >= 1024) { + if (tid < 512) { + scratch[tid] ^= scratch[tid + 512]; + } + __syncthreads(); + } + if (BLOCK_SIZE >= 512) { + if (tid < 256) { + scratch[tid] ^= scratch[tid + 256]; + } + __syncthreads(); + } + if (BLOCK_SIZE >= 256) { + if (tid < 128) { + scratch[tid] ^= scratch[tid + 128]; + } + __syncthreads(); + } + if (BLOCK_SIZE >= 128) { + if (tid < 64) { + scratch[tid] ^= scratch[tid + 64]; + } + __syncthreads(); + } + if (tid < 32) WarpReduceXor(tid, scratch); + if (tid == 0) *output = scratch[0]; +} + +// Attempts to append the checksum of the `input` buffer to the `log_entries`, +// using `log_header` to track available capacity and used space. +// +// The log entry is tagged with `entry_id`. The checksum is parallelized as +// much as block dimensions allow it. +// +// If the log does not have enough space for the new entry, the entry is +// discarded. +// +// `input_size` is the size of the input buffer in bytes. +// +// LIMITATIONS: +// - Only a single thread block is supported. +// - Block dimensions must be a power of 2. +__global__ void AppendChecksum(xla::gpu::BufferDebugLogEntryId entry_id, + const uint8_t* input, uint64_t input_size, + xla::gpu::BufferDebugLogHeader* log_header, + xla::gpu::BufferDebugLogEntry* log_entries) { + const uint32_t block_size = blockDim.x * blockDim.y * blockDim.z; + const uint32_t* input_u32 = reinterpret_cast(input); + const uint64_t input_u32_size = input_size / sizeof(uint32_t); + uint32_t checksum = 0; + + assert(gridDim.x == 1 && gridDim.y == 1 && gridDim.z == 1); + if (BlockIdx() != 0) { + return; + } + + // https://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/: + // > CUDA architecture limits the numbers of threads per block (1024 threads + // > per block limit). + switch (block_size) { + case 1024: + ReduceXor<1024>(input_u32, input_u32_size, &checksum); + break; + case 512: + ReduceXor<512>(input_u32, input_u32_size, &checksum); + break; + case 256: + ReduceXor<256>(input_u32, input_u32_size, &checksum); + break; + case 128: + ReduceXor<128>(input_u32, input_u32_size, &checksum); + break; + case 64: + ReduceXor<64>(input_u32, input_u32_size, &checksum); + break; + case 32: + ReduceXor<32>(input_u32, input_u32_size, &checksum); + break; + case 16: + ReduceXor<16>(input_u32, input_u32_size, &checksum); + break; + case 8: + ReduceXor<8>(input_u32, input_u32_size, &checksum); + break; + case 4: + ReduceXor<4>(input_u32, input_u32_size, &checksum); + break; + case 2: + ReduceXor<2>(input_u32, input_u32_size, &checksum); + break; + case 1: + ReduceXor<1>(input_u32, input_u32_size, &checksum); + break; + default: + // Unsupported block size. + assert(false); + return; + } + + if (ThreadIdx() == 0) { + const size_t last_chunk_size = input_size % sizeof(uint32_t); + uint32_t last_chunk = 0; + memcpy(&last_chunk, input + input_u32_size * sizeof(uint32_t), + last_chunk_size); + checksum ^= last_chunk; + + const uint32_t write_idx = AtomicIncSystem(&log_header->write_idx); + if (write_idx < log_header->capacity) { + log_entries[write_idx] = {entry_id, checksum}; + } + } +} + +} // namespace stream_executor::gpu + +#endif // XLA_STREAM_EXECUTOR_GPU_BUFFER_DEBUG_XOR_CHECKSUM_KERNEL_LIB_CU_H_ diff --git a/xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda_test.cc b/xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_test.cc similarity index 94% rename from xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda_test.cc rename to xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_test.cc index 2a5afaa9d7650..0d5ffca099850 100644 --- a/xla/stream_executor/cuda/buffer_debug_xor_checksum_kernel_cuda_test.cc +++ b/xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_test.cc @@ -13,31 +13,35 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel.h" + #include #include #include #include +#include #include #include #include "absl/cleanup/cleanup.h" -#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "xla/tsl/platform/status_macros.h" #include "xla/backends/gpu/runtime/buffer_debug_log_structs.h" +#include "xla/service/platform_util.h" #include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/device_description.h" #include "xla/stream_executor/gpu/buffer_debug_log.h" -#include "xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel.h" #include "xla/stream_executor/gpu/gpu_kernel_registry.h" #include "xla/stream_executor/launch_dim.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_memory_allocator.h" +#include "xla/stream_executor/stream_executor_address_allocator.h" #include "xla/stream_executor/typed_kernel_factory.h" // IWYU pragma: keep, required for KernelType::FactoryType::Create #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/errors.h" @@ -45,7 +49,7 @@ limitations under the License. namespace se = stream_executor; -namespace stream_executor::cuda { +namespace stream_executor { namespace { using xla::gpu::BufferDebugLogEntry; @@ -55,16 +59,19 @@ using xla::gpu::BufferDebugLogHeader; class ChecksumKernelTest : public ::testing::Test { protected: void SetUp() override { + std::string name = absl::AsciiStrToUpper( + xla::PlatformUtil::CanonicalPlatformName("gpu").value()); TF_ASSERT_OK_AND_ASSIGN(platform_, - se::PlatformManager::PlatformWithName("CUDA")); + se::PlatformManager::PlatformWithName(name)); TF_ASSERT_OK_AND_ASSIGN(executor_, platform_->ExecutorForDevice(0)); TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream(std::nullopt)); allocator_ = std::make_unique(stream_->parent()); - if (!executor_->GetDeviceDescription() - .cuda_compute_capability() - .IsAtLeastPascal()) { + if (const auto* cc = executor_->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + cc != nullptr && !cc->IsAtLeastPascal()) { GTEST_SKIP() << "Buffer checksumming is not supported on CUDA architectures older " "than Pascal due to missing atomic fetch_add with system scope"; @@ -258,4 +265,4 @@ TEST_F(ChecksumKernelTest, DiscardsOverflowingChecksums) { } } // namespace -} // namespace stream_executor::cuda +} // namespace stream_executor diff --git a/xla/stream_executor/rocm/BUILD b/xla/stream_executor/rocm/BUILD index 6604edb9ced1b..26428e4f45a66 100644 --- a/xla/stream_executor/rocm/BUILD +++ b/xla/stream_executor/rocm/BUILD @@ -778,6 +778,8 @@ cc_library( ":all_reduce_kernel_rocm", ":amdhipblaslt_plugin", ":buffer_comparator_kernel_rocm", + ":buffer_debug_float_check_kernel_rocm", + ":buffer_debug_xor_checksum_kernel_rocm", ":collective_signal_rocm", ":cub_scan_kernel_rocm", ":cub_sort_kernel_rocm", @@ -1428,3 +1430,51 @@ xla_cc_test( "@tsl//tsl/platform:test", ], ) + +rocm_library( + name = "buffer_debug_float_check_kernel_rocm", + srcs = [ + "buffer_debug_float_check_kernel_rocm.cu.cc", + "//xla/stream_executor/gpu:buffer_debug_float_check_kernel_lib.cu.h", + ], + tags = [ + "gpu", + "rocm-only", + ], + deps = [ + ":rocm_platform", + ":rocm_platform_id", + "//xla:util", + "//xla/backends/gpu/runtime:buffer_debug_log_structs", + "//xla/stream_executor:kernel_spec", + "//xla/stream_executor/gpu:buffer_debug_float_check_kernel", + "//xla/stream_executor/gpu:gpu_kernel_registry", + "@com_google_absl//absl/base", + "@eigen_archive//:eigen3", + "@local_config_rocm//rocm:rocm_headers", + ], + alwayslink = True, +) + +rocm_library( + name = "buffer_debug_xor_checksum_kernel_rocm", + srcs = [ + "buffer_debug_xor_checksum_kernel_rocm.cu.cc", + "//xla/stream_executor/gpu:buffer_debug_xor_checksum_kernel_lib.cu.h", + ], + tags = [ + "gpu", + "rocm-only", + ], + deps = [ + ":rocm_platform", + ":rocm_platform_id", + "//xla/backends/gpu/runtime:buffer_debug_log_structs", + "//xla/stream_executor:kernel_spec", + "//xla/stream_executor/gpu:buffer_debug_xor_checksum_kernel", + "//xla/stream_executor/gpu:gpu_kernel_registry", + "@com_google_absl//absl/base", + "@local_config_rocm//rocm:rocm_headers", + ], + alwayslink = True, +) diff --git a/xla/stream_executor/rocm/buffer_debug_float_check_kernel_rocm.cu.cc b/xla/stream_executor/rocm/buffer_debug_float_check_kernel_rocm.cu.cc new file mode 100644 index 0000000000000..1b698dfc97a82 --- /dev/null +++ b/xla/stream_executor/rocm/buffer_debug_float_check_kernel_rocm.cu.cc @@ -0,0 +1,133 @@ +/* Copyright 2025 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 + +#include "absl/base/casts.h" +#include +#include +#include "xla/stream_executor/gpu/buffer_debug_float_check_kernel.h" +#include "xla/stream_executor/gpu/gpu_kernel_registry.h" +#include "xla/stream_executor/kernel_spec.h" +#include "xla/stream_executor/rocm/rocm_platform_id.h" + +namespace stream_executor::gpu { +#if defined(__GFX9__) +__device__ static constexpr uint64_t kWarpSize = 64; +#else +__device__ static constexpr uint64_t kWarpSize = 32; +#endif +} // namespace stream_executor::gpu + +#include "xla/stream_executor/gpu/buffer_debug_float_check_kernel_lib.cu.h" + +namespace stream_executor::gpu { + +template <> +__device__ constexpr hip_bfloat16 kInfinity = + absl::bit_cast(static_cast( + absl::bit_cast(kInfinity) >> 16)); + +template <> +__device__ constexpr _Float16 kInfinity<_Float16> = + absl::bit_cast<_Float16>(uint16_t{0x7C00}); + +template <> +__device__ inline bool IsNan(hip_bfloat16 v) { + return __isnan(v); +} +template <> +__device__ inline bool IsInf(hip_bfloat16 v) { + return __isinf(v); +} +template <> +__device__ inline bool IsZero(hip_bfloat16 v) { + return v == hip_bfloat16(0.0f); +} +template <> +__device__ inline bool IsNan(_Float16 v) { + return isnan(static_cast(v)); +} +template <> +__device__ inline bool IsInf(_Float16 v) { + return isinf(static_cast(v)); +} +template <> +__device__ inline bool IsZero(_Float16 v) { + return v == 0.0f; +} + +template <> +__device__ _Float16 WarpShuffleDown(_Float16 value, unsigned int offset) { + return absl::bit_cast<_Float16>(static_cast( + __shfl_down(absl::bit_cast(value), offset))); +} + +template +__device__ inline T WarpShuffleDown(T value, unsigned int offset) { + return __shfl_down(value, offset); +} + +__device__ inline uint32_t AtomicIncSystem(uint32_t* write_idx) { + return atomicAdd_system(write_idx, 1); +} + +} // namespace stream_executor::gpu + +GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( + BufferDebugFloatCheckF32Kernel, + stream_executor::gpu::BufferDebugFloatCheckF32Kernel, + stream_executor::rocm::kROCmPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::FloatCheck), + "BufferDebugFloatCheckF32Kernel", arity); + })); + +GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( + BufferDebugFloatCheckBf16Kernel, + stream_executor::gpu::BufferDebugFloatCheckBf16Kernel, + stream_executor::rocm::kROCmPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast( + &stream_executor::gpu::FloatCheck), + "BufferDebugFloatCheckBf16Kernel", arity); + })); + +GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( + BufferDebugFloatCheckF16Kernel, + stream_executor::gpu::BufferDebugFloatCheckF16Kernel, + stream_executor::rocm::kROCmPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::FloatCheck<_Float16>), + "BufferDebugFloatCheckF16Kernel", arity); + })); + +GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( + BufferDebugFloatCheckF64Kernel, + stream_executor::gpu::BufferDebugFloatCheckF64Kernel, + stream_executor::rocm::kROCmPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::FloatCheck), + "BufferDebugFloatCheckF64Kernel", arity); + })); + +GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( + BufferDebugReduceFloatCheckResultsKernel, + stream_executor::gpu::BufferDebugAppendReducedFloatCheckResultsKernel, + stream_executor::rocm::kROCmPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::ReduceFloatCheckResults), + "BufferDebugReduceFloatCheckResultsKernel", arity); + })); diff --git a/xla/stream_executor/rocm/buffer_debug_xor_checksum_kernel_rocm.cu.cc b/xla/stream_executor/rocm/buffer_debug_xor_checksum_kernel_rocm.cu.cc new file mode 100644 index 0000000000000..6120c9f4724b4 --- /dev/null +++ b/xla/stream_executor/rocm/buffer_debug_xor_checksum_kernel_rocm.cu.cc @@ -0,0 +1,41 @@ +/* Copyright 2025 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 + +#include "absl/base/casts.h" +#include +#include "xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel.h" +#include "xla/stream_executor/gpu/buffer_debug_xor_checksum_kernel_lib.cu.h" +#include "xla/stream_executor/gpu/gpu_kernel_registry.h" +#include "xla/stream_executor/kernel_spec.h" +#include "xla/stream_executor/rocm/rocm_platform_id.h" + +namespace stream_executor::gpu { + +__device__ inline uint32_t AtomicIncSystem(uint32_t* write_idx) { + return atomicAdd_system(write_idx, 1); +} + +} // namespace stream_executor::gpu + +GPU_KERNEL_REGISTRY_REGISTER_KERNEL_STATICALLY( + BufferDebugXorChecksumKernel, + stream_executor::gpu::BufferDebugXorChecksumKernel, + stream_executor::rocm::kROCmPlatformId, ([](int arity) { + return stream_executor::KernelLoaderSpec::CreateInProcessSymbolSpec( + absl::bit_cast(&stream_executor::gpu::AppendChecksum), + "BufferDebugXorChecksumKernel", arity); + })); diff --git a/xla/stream_executor/rocm/rocm_executor.cc b/xla/stream_executor/rocm/rocm_executor.cc index 0542ec8619aa6..8882d50ef5fdf 100644 --- a/xla/stream_executor/rocm/rocm_executor.cc +++ b/xla/stream_executor/rocm/rocm_executor.cc @@ -676,7 +676,7 @@ absl::StatusOr> RocmExecutor::LoadKernel( VLOG(1) << "Resolve ROCM kernel " << kernel_name << " from symbol pointer: " << symbol; - + ScopedActivateContext activation(&rocm_context_); hipFunction_t func; RETURN_IF_ERROR( ToStatus(hipGetFuncBySymbol(&func, spec.in_process_symbol()->symbol),