[go: nahoru, domu]
Skip to content
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

Jump to bottom

Refactor llvm_compiler_test. #69947

Merged
merged 1 commit into from
Jun 19, 2024
Merged

Refactor llvm_compiler_test. #69947

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
1 change: 1 addition & 0 deletions third_party/xla/xla/service/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
7 changes: 6 additions & 1 deletion third_party/xla/xla/service/llvm_compiler.cc
View file
Open in desktop
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
Open in desktop
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 @@ -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(
Expand Down
232 changes: 57 additions & 175 deletions third_party/xla/xla/tests/llvm_compiler_test.cc
View file
Open in desktop
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