Feature Request: plug-in support for new devices

[vrv]

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).

[aodhan-domhnaill]

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

[vrv]

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.

[aodhan-domhnaill]

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?

[vrv]

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);

[aodhan-domhnaill]

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

[aodhan-domhnaill]

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!

[vrv]

1. 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.

2. 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.

[aodhan-domhnaill]

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.

[vrv]

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?

[aodhan-domhnaill]

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?

[vrv]

Adding @petewarden

[aodhan-domhnaill]

@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));

[petewarden]

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.

[aodhan-domhnaill]

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