Using Eigen in custom tensorflow op

Question

How would one go about using Eigen's internal functions like Eigen::internal::scalar_product_op in a custom Tensorflow op? After going through the Tensorflow codebase, like cwise_ops_common.h and cwise_ops.h, I came up with the following (barebones) implementation:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"

#include "tensorflow/core/framework/numeric_types.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/platform/types.h"

using namespace tensorflow;

typedef Eigen::ThreadPoolDevice CPUDevice;

REGISTER_OP("CustomMul")
    .Input("to_zero: int8")
    .Output("zeroed: int8")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

class CustomMulOp : public OpKernel {
 public:
  explicit CustomMulOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& a_tensor = context->input(0);
    const Tensor& b_tensor = context->input(0);
    auto a = a_tensor.flat<int8>();
    auto b = b_tensor.flat<int8>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, b_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<int8>();
    const CPUDevice& d = context->eigen_device<CPUDevice>();
    // Multiply the tensors
    To32Bit(output).device(d) = 
        To32Bit(a).binaryExpr(b, typename Eigen::internal::scalar_product_op<int8>());
  }
};

REGISTER_KERNEL_BUILDER(Name("CustomMul").Device(DEVICE_CPU), CustomMulOp);

Compiling the above code with

TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')

g++ -std=c++11 -shared custom_mul.cc -o custom_mul.so -fPIC -I $TF_INC -O2

produces the following error:

In file included from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:112:0,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:4,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/partial_tensor_shape.h:21,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/attr_value_util.h:22,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/node_def_util.h:23,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/shape_inference.h:21,
                 from custom_mul.cc:2:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h: In instantiation of ‘bool Eigen::TensorEvaluator<Derived, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<Derived, Device>::CoeffReturnType*) [with Derived = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<Derived, Device>::CoeffReturnType = signed char]’:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorAssign.h:123:5:   required from ‘bool Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::Scalar*) [with LeftArgType = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; RightArgType = const Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> >; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::Scalar = signed char]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h:35:66:   required from ‘static void Eigen::internal::TensorExecutor<Expression, Device, Vectorizable>::run(const Expression&, const Device&) [with Expression = const Eigen::TensorAssignOp<Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> > >; Device = Eigen::ThreadPoolDevice; bool Vectorizable = false]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h:35:62:   required from ‘Eigen::TensorDevice<ExpressionType, DeviceType>& Eigen::TensorDevice<ExpressionType, DeviceType>::operator=(const OtherDerived&) [with OtherDerived = Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> >; ExpressionType = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; DeviceType = Eigen::ThreadPoolDevice]’
custom_mul.cc:41:85:   required from here
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h:59:16: error: invalid use of incomplete type ‘const struct Eigen::ThreadPoolDevice’
       m_device.memcpy((void*)dest, m_data, sizeof(Scalar) * m_dims.TotalSize());
       ~~~~~~~~~^~~~~~
In file included from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:92:0,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:4,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/partial_tensor_shape.h:21,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/attr_value_util.h:22,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/node_def_util.h:23,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/shape_inference.h:21,
                 from custom_mul.cc:2:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorForwardDeclarations.h:84:8: note: forward declaration of ‘struct Eigen::ThreadPoolDevice’
 struct ThreadPoolDevice;
        ^~~~~~~~~~~~~~~~
In file included from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:112:0,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:4,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/partial_tensor_shape.h:21,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/attr_value_util.h:22,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/node_def_util.h:23,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/shape_inference.h:21,
                 from custom_mul.cc:2:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h: In instantiation of ‘bool Eigen::TensorEvaluator<const Derived, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Derived, Device>::CoeffReturnType*) [with Derived = Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Derived, Device>::CoeffReturnType = const signed char]’:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h:404:5:   required from ‘bool Eigen::TensorEvaluator<const Eigen::TensorCwiseBinaryOp<BinaryOp, LeftArgType, RightArgType>, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Eigen::TensorCwiseBinaryOp<BinaryOp, LeftArgType, RightArgType>, Device>::CoeffReturnType*) [with BinaryOp = Eigen::internal::scalar_product_op<signed char, signed char>; LeftArgType = const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>; RightArgType = const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer>; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Eigen::TensorCwiseBinaryOp<BinaryOp, LeftArgType, RightArgType>, Device>::CoeffReturnType = signed char]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorAssign.h:128:62:   required from ‘bool Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::Scalar*) [with LeftArgType = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; RightArgType = const Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> >; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::Scalar = signed char]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h:35:66:   required from ‘static void Eigen::internal::TensorExecutor<Expression, Device, Vectorizable>::run(const Expression&, const Device&) [with Expression = const Eigen::TensorAssignOp<Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> > >; Device = Eigen::ThreadPoolDevice; bool Vectorizable = false]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h:35:62:   required from ‘Eigen::TensorDevice<ExpressionType, DeviceType>& Eigen::TensorDevice<ExpressionType, DeviceType>::operator=(const OtherDerived&) [with OtherDerived = Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> >; ExpressionType = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; DeviceType = Eigen::ThreadPoolDevice]’
custom_mul.cc:41:85:   required from here
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h:186:16: error: invalid use of incomplete type ‘const struct Eigen::ThreadPoolDevice’
       m_device.memcpy((void*)data, m_data, m_dims.TotalSize() * sizeof(Scalar));
       ~~~~~~~~~^~~~~~
In file included from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:92:0,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:4,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/partial_tensor_shape.h:21,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/attr_value_util.h:22,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/node_def_util.h:23,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/shape_inference.h:21,
                 from custom_mul.cc:2:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorForwardDeclarations.h:84:8: note: forward declaration of ‘struct Eigen::ThreadPoolDevice’
 struct ThreadPoolDevice;
        ^~~~~~~~~~~~~~~~
In file included from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:112:0,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:4,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/partial_tensor_shape.h:21,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/attr_value_util.h:22,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/node_def_util.h:23,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/shape_inference.h:21,
                 from custom_mul.cc:2:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h: In instantiation of ‘bool Eigen::TensorEvaluator<const Derived, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Derived, Device>::CoeffReturnType*) [with Derived = Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer>; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Derived, Device>::CoeffReturnType = const signed char]’:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h:405:5:   required from ‘bool Eigen::TensorEvaluator<const Eigen::TensorCwiseBinaryOp<BinaryOp, LeftArgType, RightArgType>, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Eigen::TensorCwiseBinaryOp<BinaryOp, LeftArgType, RightArgType>, Device>::CoeffReturnType*) [with BinaryOp = Eigen::internal::scalar_product_op<signed char, signed char>; LeftArgType = const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>; RightArgType = const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer>; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Eigen::TensorCwiseBinaryOp<BinaryOp, LeftArgType, RightArgType>, Device>::CoeffReturnType = signed char]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorAssign.h:128:62:   required from ‘bool Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::evalSubExprsIfNeeded(Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::Scalar*) [with LeftArgType = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; RightArgType = const Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> >; Device = Eigen::ThreadPoolDevice; Eigen::TensorEvaluator<const Eigen::TensorAssignOp<LhsXprType, RhsXprType>, Device>::Scalar = signed char]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h:35:66:   required from ‘static void Eigen::internal::TensorExecutor<Expression, Device, Vectorizable>::run(const Expression&, const Device&) [with Expression = const Eigen::TensorAssignOp<Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> > >; Device = Eigen::ThreadPoolDevice; bool Vectorizable = false]’
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h:35:62:   required from ‘Eigen::TensorDevice<ExpressionType, DeviceType>& Eigen::TensorDevice<ExpressionType, DeviceType>::operator=(const OtherDerived&) [with OtherDerived = Eigen::TensorCwiseBinaryOp<Eigen::internal::scalar_product_op<signed char, signed char>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, int>, 16, Eigen::MakePointer>, const Eigen::TensorMap<Eigen::Tensor<const signed char, 1, 1, long int>, 16, Eigen::MakePointer> >; ExpressionType = Eigen::TensorMap<Eigen::Tensor<signed char, 1, 1, int>, 16, Eigen::MakePointer>; DeviceType = Eigen::ThreadPoolDevice]’
custom_mul.cc:41:85:   required from here
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h:186:16: error: invalid use of incomplete type ‘const struct Eigen::ThreadPoolDevice’
       m_device.memcpy((void*)data, m_data, m_dims.TotalSize() * sizeof(Scalar));
       ~~~~~~~~~^~~~~~
In file included from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:92:0,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:4,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/partial_tensor_shape.h:21,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/attr_value_util.h:22,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/node_def_util.h:23,
                 from /home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/tensorflow/core/framework/shape_inference.h:21,
                 from custom_mul.cc:2:
/home/user123/miniconda3/envs/ml/lib/python2.7/site-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorForwardDeclarations.h:84:8: note: forward declaration of ‘struct Eigen::ThreadPoolDevice’
 struct ThreadPoolDevice;

I'm using TensorFlow v.1.0.1 and GCC 6.3.1.

Answer 1

You need to add #define EIGEN_USE_THREADS before including third_party/eigen3/unsupported/Eigen/CXX11/Tensor, otherwise the ThreadPoolDevice will not be available.

Answer 2

struct ThreadPoolDevice is only forward-declared initially (at custom_mul.cc:2) but it must be fully defined. This means you miss some include. Move this include

#include <third_party/eigen3/unsupported/Eigen/CXX11/Tensor>

up to the beginning of your code. I think that should fully define ThreadPoolDevice.