Skip to content

Cherry-pick openxla/xla#44389: [ROCm] Enable float and buffer checker #989

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
magaonka-amd wants to merge 1 commit into from
Closed

Cherry-pick openxla/xla#44389: [ROCm] Enable float and buffer checker #989

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions xla/backends/gpu/runtime/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5019,7 +5019,6 @@ xla_test(
},
backends = ["gpu"],
tags = [
"cuda-only",
"gpu",
],
deps = [
Expand All @@ -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",
Expand All @@ -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",
],
)
Expand Down Expand Up @@ -5102,7 +5103,6 @@ xla_test(
},
backends = ["gpu"],
tags = [
"cuda-only",
"gpu",
],
deps = [
Expand All @@ -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",
Expand All @@ -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",
],
)
Expand Down
10 changes: 3 additions & 7 deletions xla/backends/gpu/runtime/buffers_checksum_thunk.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
14 changes: 10 additions & 4 deletions xla/backends/gpu/runtime/buffers_checksum_thunk_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,15 @@ limitations under the License.
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#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"
Expand All @@ -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"
Expand Down Expand Up @@ -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_executor::StreamExecutorAddressAllocator>(
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 "
Expand Down
10 changes: 3 additions & 7 deletions xla/backends/gpu/runtime/buffers_float_check_thunk.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
14 changes: 10 additions & 4 deletions xla/backends/gpu/runtime/buffers_float_check_thunk_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,14 @@ limitations under the License.
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.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"
Expand All @@ -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"
Expand Down Expand Up @@ -100,17 +103,20 @@ constexpr PrimitiveType kPrimitiveTypeOf<Eigen::half> = 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_executor::StreamExecutorAddressAllocator>(
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 "
Expand Down
82 changes: 9 additions & 73 deletions xla/stream_executor/cuda/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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"],
Expand Down
Loading
Loading