MCPcopy Create free account
hub / github.com/conflow-dev/ConFlow / EClipGradOp

Class EClipGradOp

en_ops/e_clip.cc:96–155  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

94REGISTER_KERNEL_BUILDER(Name("EClip").Device(DEVICE_CPU), EClipOp);
95
96class EClipGradOp : public OpKernel
97{
98public:
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", &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
150private:
151 void *lib;
152 int64 eid_low_;
153 int64 eid_high_;

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected