Custom Op written in C API compilation Issue #112

Activity

1. Rashed-MM commented on Dec 28, 2022

   

   Rashed-MM

   on Dec 28, 2022

   Last edited by Rashed-MM

   More actions

   @YogaVicky

   Sure! Here is an example of a custom op written using the C API that can be compiled and executed after being loaded with tf.load_op_library():

   #include "tensorflow/core/framework/op.h"
   #include "tensorflow/core/framework/shape_inference.h"
   #include "tensorflow/core/framework/op_kernel.h"
   
   using namespace tensorflow;
   
   REGISTER_OP("MyAdd")
       .Input("x: int32")
       .Input("y: int32")
       .Output("z: int32")
       .SetShapeFn([](shape_inference::InferenceContext* c) {
         c->set_output(0, c->input(0));
         return Status::OK();
       });
   
   class MyAddOp : public OpKernel {
    public:
     explicit MyAddOp(OpKernelConstruction* context) : OpKernel(context) {}
   
     void Compute(OpKernelContext* context) override {
       // Grab the input tensors
       const Tensor& x_tensor = context->input(0);
       const Tensor& y_tensor = context->input(1);
   
       // Create an output tensor
       Tensor* z_tensor = NULL;
       OP_REQUIRES_OK(context, context->allocate_output(0, x_tensor.shape(), &z_tensor));
   
       // Do the computation.
       int64 num_elements = x_tensor.NumElements();
       auto x_flat = x_tensor.flat<int32>();
       auto y_flat = y_tensor.flat<int32>();
       auto z_flat = z_tensor->flat<int32>();
       for (int64 i = 0; i < num_elements; ++i) {
         z_flat(i) = x_flat(i) + y_flat(i);
       }
     }
   };
   
   REGISTER_KERNEL_BUILDER(Name("MyAdd").Device(DEVICE_CPU), MyAddOp);`
   
2. Rashed-MM commented on Dec 28, 2022

   

   Rashed-MM

   on Dec 28, 2022

   Last edited by Rashed-MM

   More actions

   To compile this custom op, you will need to build it as a shared library. One way to do this is to use the TensorFlow C++ library and build the op as part of a TensorFlow binary. Here is an example of how you might do this using bazel:

   Create a file named BUILD in the same directory as your custom op code, with the following contents:

   load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_cc_binary")
   
   tf_cc_binary(
       name = "my_add_op",
       srcs = ["my_add_op.cc"],
       deps = ["@org_tensorflow//tensorflow:tensorflow"],
   )
   
3. Rashed-MM commented on Dec 28, 2022

   

   Rashed-MM

   on Dec 28, 2022

   Last edited by Rashed-MM

   More actions

   Run the following command to build the shared library:

   bazel build -c opt //path/to/custom/op:my_add_op

   This will build a shared library named libmy_add_op.so in the bazel-bin directory.

   To execute the custom op after loading it with `tf.load_