se_gpu_pjrt_compiler_aot_test.cc
imports tsl/platform/casts.himports xla/tsl/platform/statusor.h
SHA-256df7a98d75c216cfa31f750a586acfd570d0717785a30b307d6ea5cd66b843545
MaleculeMd₂(In₂Pa)
Evidence
1/* Copyright 2023 The OpenXLA Authors.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is …
11:42… IND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include <memory>
17#include <optional>
18#include <string>
19#include <utility>
20#include <vector>
21
22#include <gmock/gmock.h>
23#include <gtest/gtest.h>
24#include "absl/memory/memory.h"
25#include "absl/status/status_matchers.h"
26#include "absl/status/statusor.h"
27#include "absl/strings/string_view.h"
28#include "mlir/Dialect/Func/IR/FuncOps.h"
29#include "mlir/IR/BuiltinOps.h"
30#include "mlir/IR/MLIRContext.h"
31#include "mlir/Parser/Parser.h"
32#include "xla/hlo/builder/xla_computation.h"
33#include "xla/hlo/ir/hlo_module.h"
34#include "xla/hlo/parser/hlo_parser.h"
35#include "xla/layout.h"
36#include "xla/literal.h"
37#include "xla/literal_util.h"
38#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
39#include "xla/pjrt/compiled_memory_stats.h"
40#include "xla/pjrt/gpu/se_gpu_pjrt_client.h"
41#include "xla/pjrt/gpu/se_gpu_pjrt_compiler.h"
42#include "xla/pjrt/maybe_owning_mlir_module.h"
43#include "xla/pjrt/pjrt_client.h"
44#include "xla/pjrt/pjrt_compiler.h"
45#include "xla/pjrt/pjrt_executable.h"
46#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
47#include "xla/service/compiler.h"
48#include "xla/shape.h"
49#include "xla/shape_util.h"
50#include "xla/tests/literal_test_util.h"
51#include "xla/tsl/platform/statusor.h"
52#include "xla/xla_data.pb.h"
53#include "tsl/platform/casts.h"
54
55namespace xla {
56namespace {
57
58using ::absl_testing::IsOkAndHolds;
84:8… ::vector<std::vector<std::unique_ptr<PjRtBuffer>>>& result) {
85 ASSERT_EQ(result.size(), 1);
86 std::vector<std::unique_ptr<xla::PjRtBuffer>>& result_buffers = result[0];
87 ASSERT_EQ(result_buffers.size …
89:6… result_buffers[0]->ToLiteral().Await());
90 EXPECT_TRUE(
91 LiteralTestUtil::Equal(LiteralUtil::CreateR0(2), *result_literal));
92}
93
94TEST(StreamExecutorGpuCompilerTest, SuccessAotCom …
226:78… ,
227 LoadOptions()));
228
229 EXPECT_EQ(deserialized_executable->GetExecutable()->name(), "Identity");
230}
231
232constexpr char const* kD2HProgramTupleOutput = R"(
233 HloModul …