initial commit

pydiffvg_tensorflow/custom_ops/CMakeLists.txt

@@ -0,0 +1,29 @@
+cmake_minimum_required(VERSION 3.12)
+
+set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+project(diffvgTFCustomOp)
+
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+include_directories(SYSTEM ${TensorFlow_INCLUDE_DIR})
+
+# Compile two versions of the library
+add_library(diffvg_tf_data_ptr_cxx11_abi SHARED data_ptr.cc)
+set_target_properties(diffvg_tf_data_ptr_cxx11_abi PROPERTIES COMPILE_FLAGS -D_GLIBCXX_USE_CXX11_ABI=1)
+set_target_properties(diffvg_tf_data_ptr_cxx11_abi PROPERTIES LINK_FLAGS -D_GLIBCXX_USE_CXX11_ABI=1)
+if(APPLE)
+    # .so instead of .dylib
+    set_target_properties(diffvg_tf_data_ptr_cxx11_abi PROPERTIES SUFFIX .so)
+endif()
+target_link_libraries(diffvg_tf_data_ptr_cxx11_abi ${TensorFlow_LIBRARY})
+
+add_library(diffvg_tf_data_ptr_no_cxx11_abi SHARED data_ptr.cc)
+set_target_properties(diffvg_tf_data_ptr_no_cxx11_abi PROPERTIES COMPILE_FLAGS -D_GLIBCXX_USE_CXX11_ABI=0)
+set_target_properties(diffvg_tf_data_ptr_no_cxx11_abi PROPERTIES LINK_FLAGS -D_GLIBCXX_USE_CXX11_ABI=0)
+if(APPLE)
+    # .so instead of .dylib
+    set_target_properties(diffvg_tf_data_ptr_no_cxx11_abi PROPERTIES SUFFIX .so)
+endif()
+target_link_libraries(diffvg_tf_data_ptr_no_cxx11_abi ${TensorFlow_LIBRARY})

pydiffvg_tensorflow/custom_ops/data_ptr.cc

@@ -0,0 +1,88 @@
+// TODO: add back acknowledgement to the original author when release.
+
+#pragma warning(disable : 4003 4061 4100 4127 4242 4244 4267 4355 4365 4388 4464 4514 4574 4623 4625 4626 4647 4668 4710 4820 4946 5026 5027 5031 5039)
+
+// For windows
+#define NOMINMAX
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include <stdint.h>
+#include <climits>
+
+using namespace tensorflow;
+
+/* Tensorflow custom ops does not allow parameter types of list of 
+   various data types. Therefore, we can't pass a list but we have
+   to pass each objects individually. 
+
+   Consult Tensorflow source code: /tensorflow/core/framework/tensor.h
+   for what is supported by Tensorflow
+*/
+
+REGISTER_OP("DataPtr")
+    .Attr("T: {float, int32} = DT_INT32")  // To preserve backwards compatibility, you should specify a default value when adding an attr to an existing op:
+    .Input("input: T")  // Tensor
+    .Output("output: uint64")  // scalar
+    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+      c->set_output(0, {}); // scalar
+      return Status::OK();
+    });
+
+template <typename T>
+class DataPtrOp : public OpKernel {
+ public:
+  explicit DataPtrOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    const T *tensor = input_tensor.flat<T>().data();
+
+    // Create an output tensor
+    // NOTE: The output datatype must match the Ops definition!!!.
+    Tensor* output_tensor = NULL;
+    // Always allocate on CPU
+    AllocatorAttributes alloc_attr;
+    alloc_attr.set_on_host(true);
+    OP_REQUIRES_OK(context, 
+      context->allocate_output(0, {},  // Initialize a one-element scalar
+      &output_tensor,
+      alloc_attr)
+    );
+    auto output_flat = output_tensor->flat<uint64>();
+
+    // Cast pointer to unsigned long int
+    uintptr_t addr = (uintptr_t)tensor;
+
+    // Cast unsigned long int -> unsigned int64
+    uint64 addr_converted = addr;
+
+    output_flat(0) = addr_converted;
+  }
+};
+
+// Polymorphism: https://www.tensorflow.org/guide/extend/op#polymorphism
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_CPU)
+  .TypeConstraint<int32>("T"),
+  DataPtrOp<int32>);
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_CPU)
+  .TypeConstraint<float>("T"),
+  DataPtrOp<float>);
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_GPU)
+  .TypeConstraint<int32>("T")
+  .HostMemory("output"),
+  DataPtrOp<int32>);
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_GPU)
+  .TypeConstraint<float>("T")
+  .HostMemory("output"),
+  DataPtrOp<float>);