title: Adding a GPU Op in Tensorflow
body: <p>I am trying to add a new op to TensorFlow loosely following <a href="https://www.tensorflow.org/extend/adding_an_op" rel="nofollow noreferrer">this</a> document. The difference being that I am trying to implement a GPU based op. The op I'm trying to add is the cuda op from <a href="https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/adding_an_op" rel="nofollow noreferrer">here</a> (cuda_op.py, cuda_op_kernel.cc, cuda_op_kernel.cu.cc). I am trying to compile these outside of tensorflow and the use <code>tf.load_op_library</code> to pull them in. I have made some changes so here are my files:</p>
  
  <p>cuda_op_kernel.cc</p>
  
  <pre><code>#include "tensorflow/core/framework/op.h"
  #include "tensorflow/core/framework/shape_inference.h"
  #include "tensorflow/core/framework/op_kernel.h"
  
  using namespace tensorflow;  // NOLINT(build/namespaces)
  
  REGISTER_OP("AddOne")
      .Input("input: int32")
      .Output("output: int32")
      .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
        c-&gt;set_output(0, c-&gt;input(0));
        return Status::OK();
      });
  
  void AddOneKernelLauncher(const int* in, const int N, int* out);
  
  class AddOneOp : public OpKernel {
   public:
    explicit AddOneOp(OpKernelConstruction* context) : OpKernel(context) {}
  
    void Compute(OpKernelContext* context) override {
      // Grab the input tensor
      const Tensor&amp; input_tensor = context-&gt;input(0);
      auto input = input_tensor.flat&lt;int32&gt;();
  
      // Create an output tensor
      Tensor* output_tensor = NULL;
      OP_REQUIRES_OK(context, context-&gt;allocate_output(0, input_tensor.shape(),
                                                       &amp;output_tensor));
      auto output = output_tensor-&gt;template flat&lt;int32&gt;();
  
      // Set all but the first element of the output tensor to 0.
      const int N = input.size();
      // Call the cuda kernel launcher
      AddOneKernelLauncher(input.data(), N, output.data());
  
    }
  };
  
  REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_GPU), AddOneOp);
  </code></pre>
  
  <p>cuda_op_kernel.cu</p>
  
  <pre><code>#define EIGEN_USE_GPU
  #include &lt;cuda.h&gt;
  #include &lt;stdio.h&gt;
  
  __global__ void AddOneKernel(const int* in, const int N, int* out) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i &lt; N;
         i += blockDim.x * gridDim.x) {
      out[i] = in[i] + 1;
    }
  }
  
  void AddOneKernelLauncher(const int* in, const int N, int* out) {
    AddOneKernel&lt;&lt;&lt;32, 256&gt;&gt;&gt;(in, N, out);
  
    cudaError_t cudaerr = cudaDeviceSynchronize();
    if (cudaerr != cudaSuccess)
      printf("kernel launch failed with error \"%s\".\n", cudaGetErrorString(cudaerr));
  }
  </code></pre>
  
  <p>CMakeLists.txt</p>
  
  <pre><code>cmake_minimum_required(VERSION 3.5)
  
  #found from running python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())'
  include_directories(/usr/local/lib/python3.5/dist-packages/tensorflow/include)
  
  find_package(CUDA)
  
  #set flags based on tutorial
  set (CMAKE_CXX_FLAGS "--std=c++11 -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0")
  
  #pass flags to c++ compiler
  SET(CUDA_PROPAGATE_HOST_FLAGS ON)
  
  #create library
  cuda_add_library(
      cuda_op SHARED
      src/cuda_op_kernel.cu
      src/cuda_op_kernel.cc
      OPTIONS -gencode=arch=compute_20,code=sm_20)
  
  #copy test file to build folder
  configure_file(src/test.py test.py COPYONLY)
  </code></pre>
  
  <p>test.py</p>
  
  <pre><code>import tensorflow as tf
  mod = tf.load_op_library('./libcuda_op.so')
  with tf.Session() as sess:
      start = [5,4,3,2,1]
      print(start)
      print(mod.add_one(start).eval())
  </code></pre>
  
  <p>I am able to compile and run <code>test.py</code> successfully, but the output is always <code>[0 0 0 0 0]</code>. If I replace <code>AddOneKernel&lt;&lt;&lt;32, 256&gt;&gt;&gt;(in, N, out);</code> with <code>for (int i = 0; i &lt; N; i++) out[i] = in[i] + 1;</code> and <code>DEVICE_GPU</code> with <code>DEVICE_CPU</code>, the op outputs the right values <code>[6 5 4 3 2]</code> (with exact same <code>CMakeList.txt</code>).</p>
  
  <p>Any idea how to get the correct values to be returned?</p>