| 34 | using namespace tensorflow; |
| 35 | |
| 36 | class EAbsOp : public OpKernel |
| 37 | { |
| 38 | public: |
| 39 | explicit EAbsOp(OpKernelConstruction *context) : OpKernel(context) |
| 40 | { |
| 41 | OP_REQUIRES_OK(context, context->GetAttr("eid_low", &eid_low_)); |
| 42 | OP_REQUIRES_OK(context, context->GetAttr("eid_high", &eid_high_)); |
| 43 | OP_REQUIRES_OK(context, context->GetAttr("times", ×_)); |
| 44 | lib = dlopen("/home/zhangyan/jhrsgx/privacy_test/privacy_tf/sgx_tf_ops/sgx.so", RTLD_LAZY); |
| 45 | OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); |
| 46 | } |
| 47 | void Compute(OpKernelContext *context) override |
| 48 | { |
| 49 | const Tensor &input = context->input(0); |
| 50 | auto input_flat = input.flat<float>(); |
| 51 | |
| 52 | const TensorShape &input_shape = input.shape(); |
| 53 | |
| 54 | Tensor *output = NULL; |
| 55 | OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &output)); |
| 56 | auto output_flat = output->flat<float>(); |
| 57 | int N = input_flat.size() / times_; |
| 58 | |
| 59 | unsigned long int eid_ = (eid_high_ << 32) + eid_low_; |
| 60 | typedef void (*function)(unsigned long int eid, float *input, int N, float *output); |
| 61 | dlerror(); |
| 62 | function abs_kernel = (function)dlsym(lib, "abs1"); |
| 63 | const char *dlsym_error = dlerror(); |
| 64 | OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of abs failed: ", dlsym_error)); |
| 65 | abs_kernel(eid_, (float *)input_flat.data(), N, (float *)output_flat.data()); |
| 66 | }; |
| 67 | |
| 68 | private: |
| 69 | void *lib; |
| 70 | int64 eid_low_; |
| 71 | int64 eid_high_; |
| 72 | int64 times_; |
| 73 | }; |
| 74 | REGISTER_KERNEL_BUILDER(Name("EAbs").Device(DEVICE_CPU), EAbsOp); |
| 75 | |
| 76 | class EAbsGradOp : public OpKernel |
nothing calls this directly
no outgoing calls
no test coverage detected