89 lines
2.7 KiB
C++
89 lines
2.7 KiB
C++
// 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>);
|