Class EmbConvertGradOp

en_ops/emb_convert.cc:102–150

REGISTER_KERNEL_BUILDER(Name("EmbConvert").Device(DEVICE_CPU), EmbConvertOp);

class EmbConvertGradOp : public OpKernel
{
public:
    explicit EmbConvertGradOp(OpKernelConstruction *context) : OpKernel(context)
    {
        OP_REQUIRES_OK(context, context->GetAttr("eid_low", &eid_low_));
        OP_REQUIRES_OK(context, context->GetAttr("eid_high", &eid_high_));
        OP_REQUIRES_OK(context, context->GetAttr("num", &num_));
        OP_REQUIRES_OK(context, context->GetAttr("dim", &dim_));
        OP_REQUIRES_OK(context, context->GetAttr("times", &times_));
        lib = dlopen("/home/zhangyan/jhrsgx/privacy_test/privacy_tf/sgx_tf_ops/sgx.so", RTLD_LAZY);
        OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!"));
    }
    void Compute(OpKernelContext *context) override
    {

        // get the input tensor
        const Tensor &input = context->input(0);
        auto input_flat = input.flat<float>();
        // check shapes of input
        const TensorShape &input_shape = input.shape();
        TensorShape output_shape;
        output_shape.AddDim(input_shape.dim_size(0) / times_);
        output_shape.AddDim(num_ * (1 + dim_));
        // create output tensor
        Tensor *output = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
        auto output_flat = output->flat<float>();
        int M = input_shape.dim_size(0) / times_;
        int C = num_ * (1 + dim_);
        unsigned long int eid_ = (eid_high_ << 32) + eid_low_;
        typedef void (*function)(unsigned long int eid, float *x, float *y, int M, int C, int num, int dim);

        dlerror();
        function embedding_ss_grad_kernel = (function)dlsym(lib, "embedding_ss_grad");
        const char *dlsym_error = dlerror();
        OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of embedding_ss_grad failed: ", dlsym_error));

        embedding_ss_grad_kernel(eid_, (float *)input_flat.data(), (float *)output_flat.data(), M, C, num_, dim_);
    };

private:
    void *lib;
    int64 eid_low_;
    int64 eid_high_;
    int64 num_;
    int64 dim_;
    int64 times_;
};
REGISTER_KERNEL_BUILDER(Name("EmbConvertGrad").Device(DEVICE_CPU), EmbConvertGradOp);