| 113 | OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); |
| 114 | } |
| 115 | void Compute(OpKernelContext *context) override |
| 116 | { |
| 117 | |
| 118 | // get the input tensor |
| 119 | const Tensor &input = context->input(0); |
| 120 | auto input_flat = input.flat<float>(); |
| 121 | // check shapes of input |
| 122 | const TensorShape &input_shape = input.shape(); |
| 123 | TensorShape output_shape; |
| 124 | output_shape.AddDim(input_shape.dim_size(0) / times_); |
| 125 | output_shape.AddDim(num_ * (1 + dim_)); |
| 126 | // create output tensor |
| 127 | Tensor *output = NULL; |
| 128 | OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); |
| 129 | auto output_flat = output->flat<float>(); |
| 130 | int M = input_shape.dim_size(0) / times_; |
| 131 | int C = num_ * (1 + dim_); |
| 132 | unsigned long int eid_ = (eid_high_ << 32) + eid_low_; |
| 133 | typedef void (*function)(unsigned long int eid, float *x, float *y, int M, int C, int num, int dim); |
| 134 | |
| 135 | dlerror(); |
| 136 | function embedding_ss_grad_kernel = (function)dlsym(lib, "embedding_ss_grad"); |
| 137 | const char *dlsym_error = dlerror(); |
| 138 | OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of embedding_ss_grad failed: ", dlsym_error)); |
| 139 | |
| 140 | embedding_ss_grad_kernel(eid_, (float *)input_flat.data(), (float *)output_flat.data(), M, C, num_, dim_); |
| 141 | }; |
| 142 | |
| 143 | private: |
| 144 | void *lib; |
nothing calls this directly
no outgoing calls
no test coverage detected