| 55 | OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); |
| 56 | } |
| 57 | void Compute(OpKernelContext *context) override |
| 58 | { |
| 59 | |
| 60 | const Tensor &input = context->input(0); |
| 61 | auto input_flat = input.flat<float>(); |
| 62 | const TensorShape &input_shape = input.shape(); |
| 63 | const Tensor &real = context->input(1); |
| 64 | auto real_flat = real.flat<float>(); |
| 65 | const TensorShape &real_shape = real.shape(); |
| 66 | |
| 67 | TensorShape output_shape; |
| 68 | output_shape.AddDim(input_shape.dim_size(0) / times_); |
| 69 | |
| 70 | Tensor *output = NULL; |
| 71 | OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); |
| 72 | auto output_flat = output->flat<float>(); |
| 73 | |
| 74 | const int N = input_flat.size() / times_; |
| 75 | int M = input_shape.dim_size(0) / times_; |
| 76 | int C = input_shape.dim_size(1); |
| 77 | |
| 78 | unsigned long int eid_ = (eid_high_ << 32) + eid_low_; |
| 79 | typedef void (*function)(unsigned long int eid, float *input, int N, int M, int C, float *label, float *output); |
| 80 | dlerror(); |
| 81 | function softmax_cross_entropy_kernel = (function)dlsym(lib, "softmax_cross_entropy"); |
| 82 | const char *dlsym_error = dlerror(); |
| 83 | OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of softmax_cross_entropy failed: ", dlsym_error)); |
| 84 | softmax_cross_entropy_kernel(eid_, (float *)input_flat.data(), N, M, C, (float *)real_flat.data(), (float *)output_flat.data()); |
| 85 | }; |
| 86 | |
| 87 | private: |
| 88 | void *lib; |
nothing calls this directly
no outgoing calls
no test coverage detected