Class EDotOp

en_ops/e_dot.cc:52–114

    .Output("output2: float");

class EDotOp : public OpKernel
{
public:
    explicit EDotOp(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("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
    {
        const Tensor &input1 = context->input(0);
        auto input1_flat = input1.flat<float>();

        const Tensor &input2 = context->input(1);
        auto input2_flat = input2.flat<float>();

        const TensorShape &input1_shape = input1.shape();
        const TensorShape &input2_shape = input2.shape();

        int shape1[input1_shape.dims()];
        for (int i = 0; i < input1_shape.dims(); i++)
        {
            shape1[i] = input1_shape.dim_size(i);
        }

        int shape2[input2_shape.dims()];
        for (int i = 0; i < input2_shape.dims(); i++)
        {
            shape2[i] = input2_shape.dim_size(i);
        }

        TensorShape output_shape;
        for (int i = 0; i < input1_shape.dims(); i++)
        {
            int dim_tmp = shape1[i] >= shape2[i] ? shape1[i] : shape2[i];
            output_shape.AddDim(dim_tmp);
        }

        Tensor *output = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
        auto output_flat = output->flat<float>();
        int N1 = input1_flat.size() / times_;
        int N2 = input2_flat.size() / times_;
        int ndims = input1_shape.dims();

        unsigned long int eid_ = (eid_high_ << 32) + eid_low_;
        typedef void (*function)(unsigned long int eid, float *input1, float *input2, int N1, int N2, int *shape1, int *shape2, int ndims, float *output);
        dlerror();
        function dot_kernel = (function)dlsym(lib, "dot");
        const char *dlsym_error = dlerror();
        OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of dot failed: ", dlsym_error));
        dot_kernel(eid_, (float *)input1_flat.data(), (float *)input2_flat.data(), N1, N2, (int *)shape1, (int *)shape2, ndims, (float *)output_flat.data());
    };

private: