1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // This program compiles an XLA program which computes 123 and writes the 17 // resulting object file to stdout. 18 19 #include <iostream> 20 #include <vector> 21 22 #include "llvm/ADT/Triple.h" 23 #include "tensorflow/compiler/xla/client/client_library.h" 24 #include "tensorflow/compiler/xla/client/computation_builder.h" 25 #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" 26 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" 27 #include "tensorflow/compiler/xla/types.h" 28 #include "tensorflow/compiler/xla/util.h" 29 #include "tensorflow/core/platform/init_main.h" 30 #include "tensorflow/core/platform/logging.h" 31 32 using xla::string; 33 34 xla::Computation Doubler(xla::Client* client) { 35 xla::ComputationBuilder builder(client, "doubler"); 36 auto r0f32 = xla::ShapeUtil::MakeShape(xla::F32, {}); 37 auto x = builder.Parameter(0, r0f32, "x"); 38 builder.Mul(x, builder.ConstantR0<float>(2.0)); 39 return std::move(builder.Build().ValueOrDie()); 40 } 41 42 int main(int argc, char** argv) { 43 tensorflow::port::InitMain(argv[0], &argc, &argv); 44 45 auto client = xla::ClientLibrary::GetOrCreateCompileOnlyClient().ValueOrDie(); 46 47 xla::ComputationBuilder builder(client, "aot_test_helper"); 48 auto opaque_shape = xla::ShapeUtil::MakeOpaqueShape(); 49 auto opaque_param = builder.Parameter(0, opaque_shape, "x"); 50 auto r0f32 = xla::ShapeUtil::MakeShape(xla::F32, {}); 51 auto sum = builder.CustomCall("SumStructElements", {opaque_param}, r0f32); 52 builder.Call(Doubler(client), {sum}); 53 54 if (argc != 2) { 55 LOG(FATAL) << "local_client_aot_test_helper TARGET_CPU"; 56 } 57 58 string triple_string; 59 string target_cpu = argv[1]; 60 if (target_cpu == "k8") { 61 triple_string = "x86_64-none-linux-gnu"; 62 } else if (target_cpu == "darwin") { 63 triple_string = "x86_64-apple-macosx"; 64 } else if (target_cpu == "arm") { 65 triple_string = "aarch64-none-linux-gnu"; 66 } else if (target_cpu == "local") { 67 triple_string = xla::llvm_ir::AsString(llvm::sys::getDefaultTargetTriple()); 68 } else { 69 LOG(FATAL) << "unsupported TARGET_CPU: " << target_cpu; 70 } 71 72 llvm::Triple triple(xla::llvm_ir::AsStringRef(triple_string)); 73 74 xla::Computation computation = builder.Build().ConsumeValueOrDie(); 75 xla::CompileOnlyClient::AotComputationInstance instance{ 76 &computation, /*argument_layouts=*/{&opaque_shape}, &r0f32}; 77 78 xla::cpu::CpuAotCompilationOptions options( 79 triple_string, 80 /*cpu_name=*/"", /*features=*/"", "SumAndDouble", 81 xla::cpu::CpuAotCompilationOptions::RelocationModel::Static); 82 83 auto results = 84 client->CompileAheadOfTime({instance}, options).ConsumeValueOrDie(); 85 auto result = xla::unique_ptr_static_cast<xla::cpu::CpuAotCompilationResult>( 86 std::move(results.front())); 87 // It's lame to hard-code the buffer assignments, but we need 88 // local_client_aot_test.cc to be able to easily invoke the function. 89 CHECK_EQ(result->result_buffer_index(), 1); 90 CHECK_EQ(result->buffer_sizes().size(), 2); 91 CHECK_EQ(result->buffer_sizes()[0], -1); // param buffer 92 CHECK_EQ(result->buffer_sizes()[1], sizeof(float)); // result buffer 93 if (triple.isOSBinFormatELF()) { 94 // Check the ELF magic. 95 CHECK_EQ(result->object_file_data()[0], 0x7F); 96 CHECK_EQ(result->object_file_data()[1], 'E'); 97 CHECK_EQ(result->object_file_data()[2], 'L'); 98 CHECK_EQ(result->object_file_data()[3], 'F'); 99 // Check the ELF class. 100 CHECK_EQ(result->object_file_data()[4], triple.isArch32Bit() ? 1 : 2); 101 // Check the ELF endianness: it should be little. 102 CHECK_EQ(result->object_file_data()[5], triple.isLittleEndian() ? 1 : 2); 103 // Check the ELF version: it should be 1. 104 CHECK_EQ(result->object_file_data()[6], 1); 105 } 106 107 const std::vector<char>& object_file_data = result->object_file_data(); 108 std::cout.write(object_file_data.data(), object_file_data.size()); 109 110 return 0; 111 } 112
