| 94 | REGISTER_KERNEL_BUILDER(Name("EClip").Device(DEVICE_CPU), EClipOp); |
| 95 | |
| 96 | class EClipGradOp : public OpKernel |
| 97 | { |
| 98 | public: |
| 99 | explicit EClipGradOp(OpKernelConstruction *context) : OpKernel(context) |
| 100 | { |
| 101 | OP_REQUIRES_OK(context, context->GetAttr("eid_low", &eid_low_)); |
| 102 | OP_REQUIRES_OK(context, context->GetAttr("eid_high", &eid_high_)); |
| 103 | OP_REQUIRES_OK(context, context->GetAttr("times", ×_)); |
| 104 | lib = dlopen("/home/zhangyan/jhrsgx/privacy_test/privacy_tf/sgx_tf_ops/sgx.so", RTLD_LAZY); |
| 105 | OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); |
| 106 | } |
| 107 | void Compute(OpKernelContext *context) override |
| 108 | { |
| 109 | |
| 110 | const Tensor &grad = context->input(0); |
| 111 | const Tensor &input = context->input(1); |
| 112 | |
| 113 | const Tensor &down = context->input(2); |
| 114 | auto down_flat = down.flat<float>(); |
| 115 | |
| 116 | const Tensor &up = context->input(3); |
| 117 | auto up_flat = up.flat<float>(); |
| 118 | |
| 119 | auto grad_flat = grad.flat<float>(); |
| 120 | auto input_flat = input.flat<float>(); |
| 121 | |
| 122 | const TensorShape &input_shape = input.shape(); |
| 123 | |
| 124 | Tensor *output = NULL; |
| 125 | OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &output)); |
| 126 | auto output_flat = output->flat<float>(); |
| 127 | |
| 128 | int N = input_flat.size() / times_; |
| 129 | |
| 130 | int n1 = 1; |
| 131 | int n2 = 1; |
| 132 | if (down_flat.size() > 1) |
| 133 | { |
| 134 | n1 = down_flat.size() / times_; |
| 135 | } |
| 136 | if (up_flat.size() > 1) |
| 137 | { |
| 138 | n2 = up_flat.size() / times_; |
| 139 | } |
| 140 | |
| 141 | unsigned long int eid_ = (eid_high_ << 32) + eid_low_; |
| 142 | typedef void (*function)(unsigned long int eid, float *input, float *grad, float *down, float *up, int n1, int n2, int N, float *output); |
| 143 | dlerror(); |
| 144 | function clip_grad_kernel = (function)dlsym(lib, "clip_grad"); |
| 145 | const char *dlsym_error = dlerror(); |
| 146 | OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of clip_grad failed: ", dlsym_error)); |
| 147 | clip_grad_kernel(eid_, (float *)input_flat.data(), (float *)grad_flat.data(), (float *)down_flat.data(), (float *)up_flat.data(), n1, n2, N, (float *)output_flat.data()); |
| 148 | }; |
| 149 | |
| 150 | private: |
| 151 | void *lib; |
| 152 | int64 eid_low_; |
| 153 | int64 eid_high_; |
nothing calls this directly
no outgoing calls
no test coverage detected