1 /* Copyright 2019 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 // Before you start, make sure libtpu.so, libtpu.h and libtpu_client.c are in
17 // the same working directory.
18 //
19 // To compile: gcc -o libtpu_client libtpu_client.c -ldl
20 // To run: sudo ./libtpu_client
21
22 #include <dlfcn.h>
23 #include <stdio.h>
24 #include <stdlib.h>
25
26 #include "libtpu.h"
27
LoadAndInitializeDriver(const char * shared_lib,struct TpuDriverFn * driver_fn)28 void* LoadAndInitializeDriver(const char* shared_lib,
29 struct TpuDriverFn* driver_fn) {
30 void* handle;
31 handle = dlopen(shared_lib, RTLD_NOW);
32 if (!handle) {
33 fprintf(stderr, "Error: %s\n", dlerror());
34 exit(EXIT_FAILURE);
35 }
36
37 PrototypeTpuDriver_Initialize* initialize_fn;
38 *(void**)(&initialize_fn) = dlsym(handle, "TpuDriver_Initialize");
39 initialize_fn(driver_fn, true);
40
41 return handle;
42 }
43
main(int argc,char ** argv)44 int main(int argc, char** argv) {
45 char* api_path = "libtpu.so";
46 if (argc == 2) {
47 api_path = argv[1];
48 }
49
50 struct TpuDriverFn driver_fn;
51 void* handle = LoadAndInitializeDriver(api_path, &driver_fn);
52
53 fprintf(stdout, "------ Going to Query Version ------\n");
54 fprintf(stdout, "TPU Driver Version: %s\n", driver_fn.TpuDriver_Version());
55
56 fprintf(stdout, "------ Going to Open a TPU Driver ------\n");
57 struct TpuDriver* driver = driver_fn.TpuDriver_Open("local://");
58
59 fprintf(stdout, "------ Going to Query for System Information ------\n");
60 struct TpuSystemInfo* info = driver_fn.TpuDriver_QuerySystemInfo(driver);
61 driver_fn.TpuDriver_FreeSystemInfo(info);
62
63 // An example of simple program to sum two parameters.
64 const char* hlo_module_text = R"(HloModule add_vec_module
65 ENTRY %add_vec (a: s32[256], b: s32[256]) -> s32[256] {
66 %a = s32[256] parameter(0)
67 %b = s32[256] parameter(1)
68 ROOT %sum = s32[256] add(%a, %b)
69 }
70 )";
71
72 fprintf(stdout, "------ Going to Compile a TPU program ------\n");
73 struct TpuCompiledProgramHandle* cph =
74 driver_fn.TpuDriver_CompileProgramFromText(driver, hlo_module_text,
75 /*num_replicas=*/1, /*eventc=*/0, /*eventv*/NULL);
76
77 TpuEvent* compile_events[] = {cph->event};
78 fprintf(stdout, "------ Going to Load a TPU program ------\n");
79 struct TpuLoadedProgramHandle* lph =
80 driver_fn.TpuDriver_LoadProgram(driver, /*core_id=*/0, cph,
81 /*eventc=*/1, /*eventv=*/compile_events);
82
83 const int size = 1024;
84
85 fprintf(stdout, "------ Going to Allocate a TPU Buffer ------\n");
86 struct TpuBufferHandle* buf_a_handle =
87 driver_fn.TpuDriver_Allocate(driver, /*core-id=*/0, /*memory_region=*/1,
88 /*bytes=*/size, /*eventc=*/0, /*eventv=*/NULL);
89 fprintf(stdout, "------ Going to Allocate a TPU Buffer ------\n");
90 struct TpuBufferHandle* buf_b_handle =
91 driver_fn.TpuDriver_Allocate(driver, /*core-id=*/0, /*memory_region=*/1,
92 /*bytes=*/size, /*eventc=*/0, /*eventv=*/NULL);
93 fprintf(stdout, "------ Going to Allocate a TPU Buffer ------\n");
94 struct TpuBufferHandle* buf_sum_handle =
95 driver_fn.TpuDriver_Allocate(driver, /*core-id=*/0, /*memory_region=*/1,
96 /*bytes=*/size, /*eventc=*/0, /*eventv=*/NULL);
97
98 char a_src[size], b_src[size], sum_src[size];
99 for (int i = 0; i < size; ++i) {
100 a_src[i] = 1;
101 b_src[i] = 2;
102 sum_src[i] = 0;
103 }
104
105 TpuEvent* allocate_buf_a_events[] = {buf_a_handle->event};
106 fprintf(stdout, "------ Going to Transfer To Device ------\n");
107 struct TpuEvent* transfer_ev1 =
108 driver_fn.TpuDriver_TransferToDevice(driver, a_src, buf_a_handle,
109 /*eventc=*/1, /*eventv=*/allocate_buf_a_events);
110 TpuEvent* allocate_buf_b_events[] = {buf_a_handle->event};
111 fprintf(stdout, "------ Going to Transfer To Device ------\n");
112 struct TpuEvent* transfer_ev2 =
113 driver_fn.TpuDriver_TransferToDevice(driver, b_src, buf_b_handle,
114 /*eventc=*/1, /*eventv=*/allocate_buf_b_events);
115
116 fprintf(stdout, "------ Going to Execute a TPU program ------\n");
117 DeviceAssignment device_assignment = {NULL, 0};
118 TpuBufferHandle* input_buffer_handle[] = {buf_a_handle, buf_b_handle};
119 TpuBufferHandle* output_buffer_handle[] = {buf_sum_handle};
120 TpuEvent* transfer_events[] = {transfer_ev1, transfer_ev2};
121 struct TpuEvent* execute_event =
122 driver_fn.TpuDriver_ExecuteProgram(driver, lph,
123 /*inputc=*/2, /*input_buffer_handle=*/input_buffer_handle,
124 /*outputc=*/1, /*output_buffer_handle=*/output_buffer_handle,
125 device_assignment,
126 /*eventc=*/2, /*eventv*/transfer_events);
127
128 fprintf(stdout, "------ Going to Transfer From Device ------\n");
129 TpuEvent* execute_events[] = {execute_event};
130 struct TpuEvent* transfer_sum_event =
131 driver_fn.TpuDriver_TransferFromDevice(driver, buf_sum_handle, sum_src,
132 /*eventc=*/1, /*eventv=*/execute_events);
133
134 TpuStatus* status = driver_fn.TpuDriver_EventAwait(transfer_sum_event,
135 10000000);
136 if (status->code != 0) {
137 fprintf(stdout, "Transfer Event Await: Code: %d, Message: %s\n",
138 status->code, status->msg);
139 }
140
141 fprintf(stdout, "------ Going to Unload a TPU program ------\n");
142 struct TpuEvent* unload_program_event = driver_fn.TpuDriver_UnloadProgram(
143 driver, lph, /*eventc=*/1, /*eventv=*/execute_events);
144
145 fprintf(stdout, "------ Going to Deallocate a TPU Buffer ------\n");
146 struct TpuEvent* dealloc_ev1 = driver_fn.TpuDriver_Deallocate(driver,
147 buf_a_handle, /*eventc=*/0, /*eventv=*/NULL);
148 driver_fn.TpuDriver_FreeEvent(dealloc_ev1);
149
150 fprintf(stdout, "------ Going to Deallocate a TPU Buffer ------\n");
151 struct TpuEvent* dealloc_ev2 = driver_fn.TpuDriver_Deallocate(driver,
152 buf_b_handle, /*eventc=*/0, /*eventv=*/NULL);
153 driver_fn.TpuDriver_FreeEvent(dealloc_ev2);
154
155 fprintf(stdout, "------ Going to Deallocate a TPU Buffer ------\n");
156 struct TpuEvent* dealloc_ev3 = driver_fn.TpuDriver_Deallocate(driver,
157 buf_sum_handle, /*eventc=*/0, /*eventv=*/NULL);
158 driver_fn.TpuDriver_FreeEvent(dealloc_ev3);
159
160 fprintf(stdout, "sum:\n");
161 for (size_t i = 0; i < size; ++i) {
162 fprintf(stdout, "%d ", sum_src[i]);
163 }
164
165 dlclose(handle);
166 exit(EXIT_SUCCESS);
167 }
168