how the tensorflow call the device?

0
votes

how the tensorflow call the device? And how the tesorflow assign the task to the device for automatically? Are there interface source code in tensorflow? But I can't find the source code in details.

just like :

with tf.device('/gpu:2'):

a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')

b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')

c = tf.matmul(a, b)

why when we run the with tf.device('/gpu:2'):,the ops(a,b,c) can assign to the gpu:2 for automatically,how can I find the detailed descriptions in the source code of tensorflow?

Thanks in advance.

1
tensorflowdeviceassign

1 Answers

2
votes

There are several parts of the TensorFlow codebase that are relevant to what device is invoked.

  1. SimplePlacer::Run() will assign a fully-qualified device to every node in a TensorFlow graph.

  2. Partition() will split a graph (in which each node is labelled with a fully-qualified device) into multiple per-device subgraphs.

  3. Once an OpKernel instance has been created for each of your ops (e.g. MatMulOp for tf.matmul()), BaseGPUDevice::Compute() is called to invoke that kernel on a particular GPU. This calls BaseGPUDevice::ComputeHelper(), which obtains a CUDA stream for the computation, and calls OpKernel::Compute().

  4. In the OpKernel::Compute() method for your kernel (e.g. MatMulOp::Compute()), some device-specific code will be invoked to issue the computation on the CUDA stream (e.g. for MatMulOp, we invoke cuBLAS GEMM). Most invocations on the CUDA stream go via the tensorflow/stream_executor/stream.h interface.