[go: nahoru, domu]
Skip to content

Commit

Permalink
Refactor llvm_compiler_test.
Browse files Browse the repository at this point in the history 
We can run the CpuCompiler and GPUCompiler related tests in separate test
targets.

PiperOrigin-RevId: 644307904
  • Loading branch information
@akuegel @tensorflower-gardener
akuegel authored and tensorflower-gardener committed Jun 19, 2024
1 parent 51f415c commit 18a7c95
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 198 deletions.
1 change: 1 addition & 0 deletions third_party/xla/xla/service/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1651,6 +1651,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",
Expand Down
7 changes: 6 additions & 1 deletion third_party/xla/xla/service/llvm_compiler.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,11 @@ limitations under the License.

#include "xla/service/llvm_compiler.h"

#include <memory>
#include <utility>
#include <vector>

#include "absl/status/statusor.h"
#include "tsl/platform/denormal.h"
#include "tsl/profiler/lib/scoped_annotation.h"

Expand Down Expand Up @@ -55,6 +60,6 @@ absl::StatusOr<std::vector<std::unique_ptr<Executable>>> LLVMCompiler::Compile(
result.push_back(std::move(executable));
}

return {std::move(result)};
return std::move(result);
}
} // namespace xla
38 changes: 16 additions & 22 deletions third_party/xla/xla/tests/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -2461,32 +2457,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(
Expand Down
232 changes: 57 additions & 175 deletions third_party/xla/xla/tests/llvm_compiler_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,199 +21,81 @@ limitations under the License.
#include <utility>
#include <vector>

#include <gtest/gtest.h>
#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<GpuCompiler::BackendCompileResult> 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<std::unique_ptr<Backend>> 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<float>(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<float>(42.0)));

std::unique_ptr<HloModule> hlo_module = CreateNewVerifiedModule();
hlo_module->AddEntryComputation(builder.Build());

auto module_group = std::make_unique<HloModuleGroup>("test_module_group");
module_group->push_back(hlo_module->Clone());
module_group->push_back(std::move(hlo_module));

std::vector<std::vector<se::StreamExecutor*>> 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> backend_;

static std::string TestName() {
return ::testing::UnitTest::GetInstance()->current_test_info()->name();
}

std::unique_ptr<HloModule> CreateNewVerifiedModule() {
HloModuleConfig config;
config.set_debug_options(GetDebugOptionsFromFlags());
return std::make_unique<VerifiedHloModule>(
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<xla::LLVMCompiler*>(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<std::unique_ptr<HloModule>> modules;
modules.push_back(std::move(hlo_module));
modules.push_back(std::move(hlo_module2));
auto module_group =
std::make_unique<HloModuleGroup>("test_module_group", std::move(modules));

std::vector<std::vector<se::StreamExecutor*>> 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

0 comments on commit 18a7c95

Please sign in to comment.