diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 3d1f6e42012fa7..e1a420ff607158 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -1654,6 +1654,7 @@ cc_library( local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]), deps = [ ":compiler", + "@com_google_absl//absl/status:statusor", "@llvm-project//llvm:Core", "@local_tsl//tsl/platform:denormal", "@local_tsl//tsl/profiler/lib:scoped_annotation", diff --git a/third_party/xla/xla/service/llvm_compiler.cc b/third_party/xla/xla/service/llvm_compiler.cc index 4bbcca0b484cd2..02afa91c5d9ea5 100644 --- a/third_party/xla/xla/service/llvm_compiler.cc +++ b/third_party/xla/xla/service/llvm_compiler.cc @@ -15,6 +15,11 @@ limitations under the License. #include "xla/service/llvm_compiler.h" +#include +#include +#include + +#include "absl/status/statusor.h" #include "tsl/platform/denormal.h" #include "tsl/profiler/lib/scoped_annotation.h" @@ -55,6 +60,6 @@ absl::StatusOr>> LLVMCompiler::Compile( result.push_back(std::move(executable)); } - return {std::move(result)}; + return std::move(result); } } // namespace xla diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 69c591200c9025..b841b47a70624f 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -12,10 +12,6 @@ load( "if_cuda_is_configured", ) load("//xla:xla.bzl", "tests_build_defs_bzl_deps", "xla_cc_binary", "xla_cc_test") -load( - "//xla/stream_executor:build_defs.bzl", - "if_gpu_is_configured", -) load("//xla/tests:build_defs.bzl", "generate_backend_suites", "generate_backend_test_macros", "xla_test", "xla_test_library") load("//xla/tsl:tsl.bzl", "internal_visibility", "tsl_copts") load("//xla/tsl:tsl.default.bzl", "filegroup") @@ -2462,32 +2458,30 @@ xla_test( xla_test( name = "llvm_compiler_test", srcs = ["llvm_compiler_test.cc"], - backends = ["gpu"], - local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]) + if_rocm_is_configured([ - "TENSORFLOW_USE_ROCM", - ]), - # TODO(b/317293391) Remove once Bazel test_suite handles tags correctly - tags = ["gpu"], - deps = if_gpu_is_configured([ - ":verified_hlo_module", + backend_tags = { + # TODO(b/317293391) Remove once Bazel test_suite handles tags correctly. + "gpu": ["gpu"], + }, + backends = [ + "cpu", + "gpu", + ], + deps = [ + ":hlo_test_base", "//xla:literal_util", "//xla:test_helpers", "//xla/hlo/ir:hlo", - "@com_google_absl//absl/status", + "//xla/hlo/ir:hlo_module_group", "//xla/service:backend", - "//xla/service:cpu_plugin", "//xla/service:llvm_compiler", - "//xla/service:platform_util", - "//xla/service/cpu:cpu_compiler", - "//xla/service/gpu:gpu_compiler", "//xla/stream_executor", + "@com_google_absl//absl/status", + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Core", + "@local_tsl//tsl/platform:casts", "@local_tsl//tsl/platform:env", "@local_tsl//tsl/platform:test_main", - ]) + if_cuda_is_configured([ - "//xla/stream_executor/cuda:cuda_platform_id", - ]) + if_rocm_is_configured([ - "//xla/stream_executor/rocm:rocm_platform_id", - ]), + ], ) xla_test( diff --git a/third_party/xla/xla/tests/llvm_compiler_test.cc b/third_party/xla/xla/tests/llvm_compiler_test.cc index dd15cbbeba4cc2..238482b650c3d6 100644 --- a/third_party/xla/xla/tests/llvm_compiler_test.cc +++ b/third_party/xla/xla/tests/llvm_compiler_test.cc @@ -21,199 +21,81 @@ limitations under the License. #include #include +#include #include "absl/status/status.h" +#include "llvm/IR/Module.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module_group.h" #include "xla/literal_util.h" #include "xla/service/backend.h" -#include "xla/service/cpu/cpu_compiler.h" -#include "xla/service/gpu/gpu_compiler.h" -#include "xla/service/platform_util.h" -#if GOOGLE_CUDA -#include "xla/stream_executor/cuda/cuda_platform_id.h" -#elif TENSORFLOW_USE_ROCM -#include "xla/stream_executor/rocm/rocm_platform_id.h" -#endif #include "xla/stream_executor/device_description.h" #include "xla/stream_executor/stream_executor.h" #include "xla/test_helpers.h" -#include "xla/tests/verified_hlo_module.h" +#include "xla/tests/hlo_test_base.h" +#include "tsl/platform/casts.h" #include "tsl/platform/threadpool.h" namespace xla { -namespace gpu { - -// Creating dummy data structure needed to initialize a GpuDummyCompiler -constexpr char kDummyTriple[] = "dummy-triple"; -constexpr char kDummyLayout[] = "e"; -const se::Platform::Id kGpuPlatformId = -#if GOOGLE_CUDA - se::cuda::kCudaPlatformId; -#elif TENSORFLOW_USE_ROCM - se::rocm::kROCmPlatformId; -#endif -// This class is a dummy implementation of GpuCompiler and is targeted for unit -// test only -class GpuDummyCompiler : public GpuCompiler { - public: - GpuDummyCompiler() - : GpuCompiler(kGpuPlatformId, kDummyTriple, kDummyLayout) {} - - int32_t GetToolkitVersion() const override { return 0; } - - absl::Status OptimizeHloConvolutionCanonicalization( - HloModule* hlo_module, se::GpuComputeCapability gpu_version, - se::dnn::VersionInfo dnn_version, - se::DeviceMemoryAllocator* device_allocator) { - return absl::OkStatus(); - } - - absl::Status OptimizeHloPostLayoutAssignment( - HloModule* hlo_module, se::StreamExecutor* stream_executor, - const CompileOptions& options, const TargetConfig& gpu_target_config, - tsl::thread::ThreadPool* thread_pool) override { - return absl::OkStatus(); - } +namespace { - absl::StatusOr CompileTargetBinary( - const HloModuleConfig& module_config, llvm::Module* llvm_module, - se::GpuComputeCapability gpu_version, bool relocatable, - const HloModule* debug_module, const CompileOptions& options) override { - return BackendCompileResult{}; - } -}; -} // namespace gpu +using LLVMCompilerTest = HloTestBase; -namespace { +const char* const kHloText = R"( +HloModule Constant -class LLVMCompilerTest : public ::testing::Test { - public: - void SetUp() override { - Platform* platform = FindPlatform(); - ASSERT_NE(platform, nullptr); - - BackendOptions backend_options; - backend_options.set_platform(platform); - absl::StatusOr> backend_or_status = - Backend::CreateBackend(backend_options); - ASSERT_IS_OK(backend_or_status.status()); - backend_ = std::move(backend_or_status).value(); - } - - ~LLVMCompilerTest() override {} - - protected: - using Platform = se::Platform; - - explicit LLVMCompilerTest(std::string platform_name) - : platform_name_(std::move(platform_name)) {} - - void TestCompilerHooks(LLVMCompiler* compiler) { - int pre_opt_hook_call_count = 0; - int post_opt_hook_call_count = 0; - - auto pre_opt_hook = [&pre_opt_hook_call_count](const llvm::Module&) { - ++pre_opt_hook_call_count; - return absl::OkStatus(); - }; - auto post_opt_hook = [&post_opt_hook_call_count](const llvm::Module&) { - ++post_opt_hook_call_count; - return absl::OkStatus(); - }; - - // Create HLO module, and run the compiler. - auto builder = HloComputation::Builder(TestName()); - builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - - auto hlo_module = CreateNewVerifiedModule(); - hlo_module->AddEntryComputation(builder.Build()); - - compiler->SetPreOptimizationHook(pre_opt_hook); - compiler->SetPostOptimizationHook(post_opt_hook); - - ASSERT_TRUE(compiler - ->RunBackend(std::move(hlo_module), - backend_->default_stream_executor(), - /*device_allocator=*/nullptr) - .ok()); - - // Test that hooks were called. - EXPECT_EQ(1, pre_opt_hook_call_count); - EXPECT_EQ(1, post_opt_hook_call_count); - } - - void TestMultiModuleCompilation(LLVMCompiler* compiler) { - HloComputation::Builder builder(TestName()); - builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - - std::unique_ptr hlo_module = CreateNewVerifiedModule(); - hlo_module->AddEntryComputation(builder.Build()); - - auto module_group = std::make_unique("test_module_group"); - module_group->push_back(hlo_module->Clone()); - module_group->push_back(std::move(hlo_module)); - - std::vector> executors; - executors.push_back({backend_->default_stream_executor()}); - executors.push_back({backend_->default_stream_executor()}); - - EXPECT_IS_OK(compiler->Compile(std::move(module_group), - std::move(executors), - /*device_allocator=*/nullptr)); - } - - private: - Platform* FindPlatform() { - auto status_or_platform = PlatformUtil::GetPlatform(platform_name_); - return status_or_platform.ok() ? status_or_platform.value() : nullptr; - } - - std::string platform_name_; - std::unique_ptr backend_; - - static std::string TestName() { - return ::testing::UnitTest::GetInstance()->current_test_info()->name(); - } - - std::unique_ptr CreateNewVerifiedModule() { - HloModuleConfig config; - config.set_debug_options(GetDebugOptionsFromFlags()); - return std::make_unique( - TestName(), config, /*verifier_layout_sensitive=*/false, - /*allow_mixed_precision_in_hlo_verifier=*/true, - backend_->compiler()->ShapeSizeBytesFunction()); - } -}; - -class CpuCompilerTest : public LLVMCompilerTest { - public: - CpuCompilerTest() : LLVMCompilerTest("Host") {} -}; - -class GpuCompilerTest : public LLVMCompilerTest { - public: - GpuCompilerTest() : LLVMCompilerTest("GPU") {} -}; - -TEST_F(CpuCompilerTest, HooksTest) { - cpu::CpuCompiler compiler; - TestCompilerHooks(&compiler); +ENTRY main { + ROOT constant = f32[] constant(42.0) } +)"; -TEST_F(GpuCompilerTest, HooksTest) { - gpu::GpuDummyCompiler compiler; - TestCompilerHooks(&compiler); -} +TEST_F(LLVMCompilerTest, HooksTest) { + int pre_opt_hook_call_count = 0; + int post_opt_hook_call_count = 0; -TEST_F(CpuCompilerTest, CpuMultiModuleCompilation) { - cpu::CpuCompiler compiler; - TestMultiModuleCompilation(&compiler); + auto pre_opt_hook = [&pre_opt_hook_call_count](const llvm::Module&) { + ++pre_opt_hook_call_count; + return absl::OkStatus(); + }; + auto post_opt_hook = [&post_opt_hook_call_count](const llvm::Module&) { + ++post_opt_hook_call_count; + return absl::OkStatus(); + }; + + // Create HLO module, and run the compiler. + auto hlo_module = ParseAndReturnVerifiedModule(kHloText).value(); + LLVMCompiler* compiler = + tensorflow::down_cast(backend().compiler()); + compiler->SetPreOptimizationHook(pre_opt_hook); + compiler->SetPostOptimizationHook(post_opt_hook); + + ASSERT_TRUE(compiler + ->RunBackend(std::move(hlo_module), + backend().default_stream_executor(), + /*device_allocator=*/nullptr) + .ok()); + + // Test that hooks were called. + EXPECT_EQ(1, pre_opt_hook_call_count); + EXPECT_EQ(1, post_opt_hook_call_count); } -TEST_F(GpuCompilerTest, GpuMultModuleCompilation) { - gpu::GpuDummyCompiler compiler; - TestMultiModuleCompilation(&compiler); +TEST_F(LLVMCompilerTest, DISABLED_MultiModuleCompilation) { + auto hlo_module = ParseAndReturnVerifiedModule(kHloText).value(); + auto hlo_module2 = ParseAndReturnVerifiedModule(kHloText).value(); + std::vector> modules; + modules.push_back(std::move(hlo_module)); + modules.push_back(std::move(hlo_module2)); + auto module_group = + std::make_unique("test_module_group", std::move(modules)); + + std::vector> executors; + executors.push_back({backend().default_stream_executor()}); + executors.push_back({backend().default_stream_executor()}); + + EXPECT_IS_OK(backend().compiler()->Compile(std::move(module_group), + std::move(executors), + backend().memory_allocator())); } + } // namespace } // namespace xla