不再让CPU和总线拖后腿：Exafunction让GPU跑的更快！

tensorflow::SavedModelBundle bundle;
TF_RETURN_IF_ERROR(bundle, tensorflow::LoadSavedModel(
    session_options, run_options, saved_model_dir, tags, &bundle
));

tensorflow::Session* session = bundle.GetSession();
// Create a tensor in CPU memory
tensorflow::Tensor tensor(tensorflow::DT_FLOAT, {1, 2, 3});
// Pairs of feed name and tensor to pass into the model
std::vector<std::pair<std::string, tensorflow::Tensor>> inputs{"input", std::move(tensor)};
// The outputs will be written here. These will also be on the CPU.
std::vector<tensorflow::Tensor> outputs;
// Tensor names to fetch for the output.
std::vector<std::string> fetch_names;
// Run the model!
session->Run(inputs, fetch_names, {}, &outputs);

tensorflow::CallableOptions callable_options;
std::string gpu_device_name = FirstGpuDeviceName(session);
// Names of input tensors.
std::vector<std::string> feed_names;
*callable_options.mutable_feed() = {feed_names.begin(), feed_names.end()};
// Names of output tensors.
std::vector<std::string> fetch_names;
*callable_options.mutable_fetch() = {fetch_names.begin(), fetch_names.end()};
auto& feed_devices = *callable_options.mutable_feed_devices();
for (auto& input_name : feed_names) {
  feed_devices[input_name] = gpu_device_name;
}
auto& fetch_devices = *callable_options.mutable_fetch_devices();
for (auto& output_name : fetch_names) {
  fetch_devices[output_name] = gpu_device_name;
}

std::string FirstGpuDeviceName(tensorflow::Session* session) {
  // Gets device name for the first GPU in the session.
  std::vector<tensorflow::DeviceAttributes> devices;
  auto status = session->ListDevices(&devices);
  assert(status.ok());
  for (const tensorflow::DeviceAttributes& d : devices) {
    if (d.device_type() == "GPU" || d.device_type() == "gpu") {
      return d.name();
    }
  }
  CHECK(false) << "GPU not found";
}

TF_RETURN_IF_ERROR(session->MakeCallable(callable_options, &callable));
// Before the session is destroyed:
// session->ReleaseCallable(callable);

// Get TensorFlow's GPU allocator for device 0
// This needs to match the device placement used when loading the SavedModel
// and creating the session.
tensorflow::TfDeviceId gpu_device_id(0);
tensorflow::Allocator* gpu_allocator =
    tensorflow::GPUProcessState::singleton()->GetGPUAllocator(gpu_device_id);
// Synchronize to ensure memory can be safely allocated & overwritten
cudaDeviceSynchronize();
// The input tensors are now allocated in GPU memory using TensorFlow's
// allocator. They must be in the same order as the feed names.
std::vector<tensorflow::Tensor> inputs;
for (int i = 0; i < 10; i++) {
  tensorflow::Tensor tensor(gpu_allocator, tensorflow::DT_FLOAT, {1, 2, 3});
  // Fill the input here
  inputs.push_back(std::move(tensor));
}
// Synchronize to ensure the inputs are valid
cudaDeviceSynchronize();

// The outputs will also be placed on the GPU thanks to the fetch_devices
// setting above.
std::vector<tensorflow::Tensor> outputs;
TF_RETURN_IF_ERROR(session->RunCallable(callable, inputs, &outputs, nullptr));

callable_options.set_fetch_skip_sync(true);

cudaStream_t stream = GetTfGpuStream(session, gpu_device_name);
// Returns the tensorflow::BaseGPUDevice for a given device name
tensorflow::BaseGPUDevice* GetTfGpuDevice(tensorflow::Session* session,
                                          const std::string& gpu_device_name) {
  const tensorflow::DeviceMgr* device_mgr;
  auto status = session->LocalDeviceManager(&device_mgr);
  CHECK(status.ok()) << status;
  tensorflow::Device* device;
  status = device_mgr->LookupDevice(gpu_device_name, &device);
  CHECK(status.ok()) << status;
  auto* gpu_device = dynamic_cast<tensorflow::BaseGPUDevice*>(device);
  return CHECK_NOTNULL(gpu_device);
}
// Returns the compute stream for a given TensorFlow GPU device.
cudaStream_t GetTfGpuStream(tensorflow::Session* session,
                            const std::string& gpu_device_name) {
  auto* device = GetTfGpuDevice(session, gpu_device_name);
  const tensorflow::DeviceBase::GpuDeviceInfo* device_info =
      device->tensorflow_gpu_device_info();
  CHECK_NOTNULL(device_info);
  CUstream tf_stream =
      stream_executor::gpu::AsGpuStreamValue(device_info->stream);
  return tf_stream;
}

cudaStream_t stream = GetTfGpuStream(session, gpu_device_name);
std::vector<tensorflow::Tensor> inputs;
for (int i = 0; i < 10; i++) {
  tensorflow::Tensor tensor(gpu_allocator, tensorflow::DT_FLOAT, {1, 2, 3});
  // Fill input buffers here, using the stream in kernel calls
  // or calls to cudaMemcpyAsync. Alternatively, synchronize with another
  // GPU stream using events.
  inputs.push_back(std::move(tensor));
}
std::vector<tensorflow::Tensor> outputs;
TF_RETURN_IF_ERROR(session->RunCallable(callable, inputs, &outputs, nullptr));
cudaStreamSynchronize(stream);