Feature Request: plug-in support for new devices

We need to be able to support using new devices in TensorFlow without requiring editing the TensorFlow binary source code.

This requires a few changes. Among many others:

The ability to dynamically register an implementation / factories for alternative devices
New or better APIs for allowing OpKernel's to use arbitrary device resources (e.g., right now we assume either the use of StreamExecutor or EigenDevice implementations, but that's obviously not general).

Few questions.

About Feature Request

By dynamically, I assume you are referring to at compile time, but self-contained code changes. If you mean at runtime via dynamic module loading, I don't know off hand how to do this, but I think we can figure it out.

How do you wish to handle the kernel code? Should I plan to have new hardware supported by developers modifying the tensorflow/core/kernels/* and registering the new hardware? Or do you want this also separated. Although this currently is handled with macros and registration systems in the C++, but I ask from a build perspective. We could make a new folder that contains all the code for extra devices and contains all the build scripts over there. It this is done, I can write template files that have all the basic structural code (many projects do this for configuration scripts) that a user could basically rename and be ready to go.

About TensorFlow

TensorFlow Device Contexts, Streams, and Context Switching

Yes, I do mean self-contained code changes: a user should be able to use an existing binary install of TensorFlow that doesn't contain the custom device code, but then call a function to "load a device" and then another function to "load the kernels" for that device.

The latter is already possible for existing devices via the tf.load_op_library() mechanism, so theoretically something similar could be done for a new tf.load_device().

I'll answer your other question on SO, but I don't think the answer there will be instructive for new devices. Every device has its own execution model, and so the way the device code is written needs to take into account the execution model. On CPU, the ThreadPoolDevice is just a device that uses Eigen's multi-threaded CPU support to implement operations. On GPU, we have to manage streams and execution order carefully because GPU devices are asynchronous, not synchronous.

If you told us more about the execution model of your hardware, we might be able to suggest something more relevant -- it's likely that copying the CPU or the GPU way of doing things is not the right solution.

I need to talk with my supervisor before giving too many details. Here is the public info, but as a high level understanding, the chip is basically a cluster computer reduced to fit on a single chip. The many DSPs need to work together to do anything useful.

I didn't know about the operation interface! Its pretty awesome and I definitely think that is what I want to build.

It would seem that at a minimum, a developer would need to write an Allocator, Device, DeviceFactory, and DeviceContext. This would give a non-functional device because there are no kernels registered to it. As I was developing, I noticed that some kernels seemed to be core functions like ConstOp, AssignOp, NoOp, etc. that are needed for other things to work. It would seem that a user wouldn't want to code these explicitly as they are kind of obvious and redundant. Do you think these can/should be automatically built into the framework so that every device at least has these working out of the box?

Yes, those four pieces are the minimal requirements just to get the basics working, and then you'd have to register kernels for every op that was supported.

You're right that some ops are probably trivially implementable as long as some of the basics above are implemented. We'd have to think about how to 'auto register kernels' for all devices. However, for things like NoOp, it shouldn't be too hard: it would just be an include and a registration like in

tensorflow/tensorflow/core/kernels/no_op.cc

Line 21 in 5df4c71

REGISTER_KERNEL_BUILDER(Name("NoOp").Device(DEVICE_GPU), NoOp);

How do I add headers to be installed into the TensorFlow include path? And how about for the .proto files? I need to "compile" them to headers and have them installed

I need several framework and common_runtime headers to compile device code. I tried to look through the Bazel files, but couldn't find anything obvious.

EDIT:

I tried a workaround of just setting my include path to the root of the Git repo, but found that things like #include "unsupported/Eigen/CXX11/Tensor" are being included. These would seem to be unnecessary for a "general" device. What should I do about this? I figure my options are: just include them (what I will do for the mean time), or write stripped down headers that contain only the minimal necessary to construct a device. IMO the second option is nicer for long term, but certainly more annoying to do.

EDIT 2:

Also, the Eigen files don't appear to have header guards and cause recurrent includes

FYI. I got a fake device working with contexts, factories, kernels and all, so I will begin to try to make it self contained. I posted various questions to StackOverflow and things. Those are more educational. The library problem above is more important for the purposes of this issue.

Thanks a bunch!

@keveman

At the moment, unsupported/Eigen/CXX11/Tensor is required by the framework -- it's how you can map our tensor.h Tensor into typed and shaped tensors.
Adding @keveman about how to add TF headers to include path. I thought https://github.com/tensorflow/tensorflow/blob/master/tensorflow/g3doc/get_started/os_setup.md#setting-up-tensorflow-for-development might be relevant, but I'm not sure.

I got all the code down to one file with no changes to the main code base except the below. I list my solution ideas, but they involve modification of the main code so I wanted to make sure its okay.

diff --git a/tensorflow/core/util/device_name_utils.cc b/tensorflow/core/util/device_name_utils.cc
index 5816dbd..50309ed 100644
--- a/tensorflow/core/util/device_name_utils.cc
+++ b/tensorflow/core/util/device_name_utils.cc
@@ -162,6 +162,16 @@ bool DeviceNameUtils::ParseFullName(StringPiece fullname, ParsedName* p) {
       }
       progress = true;
     }
+    if (str_util::ConsumePrefix(&fullname, "/kpu:") ||
+        str_util::ConsumePrefix(&fullname, "/KPU:")) {
+      p->has_type = true;
+      p->type = "KPU";  // Treat '/kpu:..' as uppercase '/device:KPU:...'
+      p->has_id = !str_util::ConsumePrefix(&fullname, "*");
+      if (p->has_id && !ConsumeNumber(&fullname, &p->id)) {
+        return false;
+      }
+      progress = true;
+    }

     if (!progress) {
       return false;

To get rid of this one, I was thinking of using the existing registration, but that means exposing this function. I could just make the underlying std::unordered_map a private/protected member of the DeviceFactory class and then make DeviceNameUtils a friend. Or just expose the function handle in the header.

diff --git a/tensorflow/python/framework/device.py b/tensorflow/python/framework/device.py
index 8f5125d..505c2a9 100644
--- a/tensorflow/python/framework/device.py
+++ b/tensorflow/python/framework/device.py
@@ -155,7 +155,9 @@ class DeviceSpec(object):
         elif ly == 2 and y[0] == "task":
           self.task = y[1]
         elif ((ly == 1 or ly == 2) and
-              ((y[0].upper() == "GPU") or (y[0].upper() == "CPU"))):
+              ((y[0].upper() == "GPU") or 
+               (y[0].upper() == "CPU") or
+               (y[0].upper() == "KPU"))):
           if self.device_type is not None:
             raise ValueError("Cannot specify multiple device types: %s" % spec)
           self.device_type = y[0].upper()

If there is Python access to the same interface as above, then I can just use that. I am a fan of less configuration.

I think if you name your device "/device:KPU:0" instead of just "/KPU:0", it should work without any additional edits. "/cpu:0" and "/gpu:0" were shortcuts before we realized it was helpful to have the "/device" qualifier like we do for the other parts of a full device name. Can you give that a try and let me know?

Yup that worked!

Device mapping:
/job:localhost/replica:0/task:0/device:FPU:0 -> physical description
/job:localhost/replica:0/task:0/device:FPU:1 -> physical description
/job:localhost/replica:0/task:0/device:FPU:2 -> physical description
/job:localhost/replica:0/task:0/device:FPU:3 -> physical description

Once I figure out how to get those includes and compile it outside of TF, I will submit a pull request within documentation a la "Adding an Op"

After that. I would like to discuss how to best support my particular hardware. My company probably wants this to be private as I will need to tell you details, so would you be willing to do it over email and not as a GitHub Issue?

@petewarden

Adding @petewarden

@keveman

@keveman

I have been manually adding headers to my include path and I need,

tensorflow/core/framework/device_attributes.pb_text.h
tensorflow/core/framework/op_segment.h
tensorflow/core/framework/resource_mgr.h
tensorflow/core/framework/tensor.pb_text.h
tensorflow/core/framework/type_index.h
tensorflow/core/graph/edgeset.h
tensorflow/core/graph/graph.h
tensorflow/core/graph/types.h
tensorflow/core/arena.h
tensorflow/core/lib/gtl/int_type.h
tensorflow/core/lib/gtl/iterator_range.h
tensorflow/core/public/session_options.h

Manually adding symlinks to for all those above headers tothe TensorFlow source, I still get this error,

g++ -std=c++11 -shared fpu_device.cc -o fpu.so -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include
In file included from /usr/local/lib/python2.7/dist-packages/tensorflow/include/tensorflow/core/common_runtime/local_device.h:19:0,
                 from fpu_device.cc:10:
/usr/local/lib/python2.7/dist-packages/tensorflow/include/tensorflow/core/common_runtime/device.h: In member function ‘std::string tensorflow::Device::DebugString() const’:
/usr/local/lib/python2.7/dist-packages/tensorflow/include/tensorflow/core/common_runtime/device.h:139:74: error: ‘ProtoDebugString’ was not declared in this scope
   string DebugString() const { return ProtoDebugString(device_attributes_); }
                                                                          ^
fpu_device.cc: In member function ‘virtual tensorflow::Status tensorflow::FPUDevice::MakeTensorFromProto(const tensorflow::TensorProto&, tensorflow::AllocatorAttributes, tensorflow::Tensor*)’:
fpu_device.cc:84:41: error: ‘ProtoDebugString’ was not declared in this scope
            ProtoDebugString(tensor_proto));

ProtoDebugString() is generated by tensorflow/tools/proto_text, which is a bit like protoc but for a minimal footprint class instead of a full protobuf one.

Thanks. Do you know how I can get it to be recognized when I am compiling it into an external module?

I am currently running

g++ -std=c++11 -shared fpu_device.cc -o fpu.so -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include

@RichDubielzig

@CUinNYC indeed, that doc is a great attempt to show it's done, and that's the basic idea of how the scaffolding works, but the implementation details for every new hardware device means it takes care to figure out exactly how to implement the device code.

We're working with a few external folks such as those in this thread to provide some early guidance as to how it's done, and at some point we'll have some nice examples to point at (beyond say, our CPU and GPU implementations).

@RichDubielzig: the dataflow model of TF execution means that any operations that can be run at a given time will be run by the executor (assuming there are enough inter-op-parallelism threads to run them). So yes, it is possible to have multiple "matmul" nodes running at once if both can be run at the same time.

On CPU, we use Eigen for a lot of the computation, and each op has a shared threadpool on which to execute -- so those threads will typically always be busy if there's lots of op work to do (though I admit, it probably won't be super efficient to context switch all the time, but I digress).

On GPU, we use GPU streams to enqueue asyncrhronous ops: execution of an OpKernel just queues the work to be done, and we use streams to enforce execution order at the actual device.

@hawkinsp might have more to say here.

@vrv

@vrv I really appreciate the initial documentation. Thanks to that, I successfully created a "fake CPU" (really easy indeed). It required minor changes because of the TensorFlow code evolution (r1.0). So having an external-module approach would be more portable. Let me know if you need a help or feedback comments on a preliminary documentation.

In particular, I am interested in how to allocate memory for accelerators. My goal is to reduce as much as possible memory copies.

If anybody has comments or examples I will really appreciate that.

Just following up on my question: I am seeing results with this approach:

Define my operation as an AsyncOpKernel. Kernel pseudocode:

ComputeAsync( context  input, DoneCallback done):
  Get input tensors A & B from context
  Also get a Knureon command queue handle Q from the DeviceContext object.
  Figure out the dimensions of the tensors
  Assemble a call f(A,B,dims) to my Knureon matrix multipllier algorithm
  Queue the tuple (f(), Q, done) in the launch queue.
  Return immediately

In launch queue, in a separate thread:

(as part of init routine, we have created our own dedicated ThreadPool, which includes a thread to run this wait/spawn loop)

  Wait for resources R to run f()
  schedule a task to run f() on R resources using Q and calling done().

In the task scheduled above:

  deploy f() to R and kick it off.
  wait on Q to finish
  release R back to the Knureon environment
  call the AsyncOpKernel done() callback
  exit

Asked a new question on StackOverflow about the ConstTensor, which doesn't map to the DMAHelper::base() function:

http://stackoverflow.com/questions/42707600/tensorflow-how-to-get-pointer-to-data-contents-of-consttensor

Posted another question on StackOverflow regarding my confusion with DeviceContext and when it needs to be used. I am revisiting the issue because it turns out we have a smaller limit than I thought on open queues with the system and I'm wondering if I should

redesign with the idea of mapping one DeviceContext per queue
keep with the current design of one DeviceContext per OpKernelContext and manage queues in a separate pool as they are needed.

http://stackoverflow.com/questions/42869340/tensorflow-why-does-a-gpu-device-only-have-one-device-context-and-does-it-rea

Another issue we have run into in attempting to debug XLA: It doesn't seem like we are able to exercise any backend compiler when running in a debugger. The issue is here:

#9194

@RichDubielzig

@RichDubielzig I've been told that XLA-related development discussion has been moved to https://groups.google.com/forum/#!forum/xla-dev -- you might get more help / feedback there now.

Asked a new question related to this on stackoverflow. Will also post to the xla-dev forum

https://stackoverflow.com/questions/44271775/can-tensorflow-xla-be-supported-by-plug-ins

It has been 14 days with no activity and this issue has an assignee.Please update the label and/or status accordingly.

Automatically closing due to lack of recent activity. Please update the issue when new information becomes available, and we will reopen the issue. Thanks!

Navigation Menu

Repository navigation

Feature Request: plug-in support for new devices #4359

About Feature Request

About TensorFlow

58 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Feature Request: plug-in support for new devices #4359

Description

Activity

aodhan-domhnaill commented on Sep 14, 2016

About Feature Request

About TensorFlow

vrv commented on Sep 14, 2016

aodhan-domhnaill commented on Sep 14, 2016

vrv commented on Sep 14, 2016

aodhan-domhnaill commented on Sep 14, 2016

aodhan-domhnaill commented on Sep 15, 2016

vrv commented on Sep 15, 2016

aodhan-domhnaill commented on Sep 16, 2016

vrv commented on Sep 16, 2016

aodhan-domhnaill commented on Sep 16, 2016

vrv commented on Sep 16, 2016

aodhan-domhnaill commented on Sep 19, 2016

petewarden commented on Sep 19, 2016

aodhan-domhnaill commented on Sep 19, 2016

58 remaining items

vrv commented on Feb 24, 2017

GiuseppeDiGuglielmo commented on Feb 24, 2017

RichDubielzig commented on Feb 28, 2017

RichDubielzig commented on Mar 9, 2017

RichDubielzig commented on Mar 18, 2017

RichDubielzig commented on Apr 13, 2017

vrv commented on Apr 13, 2017

RichDubielzig commented on May 30, 2017

tensorflowbutler commented on Dec 20, 2017

tensorflowbutler commented on Jan 4, 2018

shivaniag commented on Jan 8, 2018

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions