Method Compute

en_ops/e_batch_norm.cc:62–95

        OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!"));
    }
    void Compute(OpKernelContext *context) override
    {
        const Tensor &input = context->input(0);
        const Tensor &scale = context->input(1);
        const Tensor &offset = context->input(2);
        const Tensor &mean = context->input(3);
        const Tensor &variance = context->input(4);

        const TensorShape &input_shape = input.shape();

        Tensor *output = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &output));

        int M = input_shape.dim_size(0) / times_;
        int H = input_shape.dim_size(1);
        int W = input_shape.dim_size(2);

        auto input_flat = input.flat<float>();
        auto scale_flat = scale.flat<float>();
        auto offset_flat = offset.flat<float>();
        auto mean_flat = mean.flat<float>();
        auto variance_flat = variance.flat<float>();
        auto output_flat = output->flat<float>();

        int N = input_flat.size() / times_;

        unsigned long int eid_ = (eid_high_ << 32) + eid_low_;
        typedef void (*function)(unsigned long int eid, float *input, float *scale, float *offset, float *mean_x, float *variance, int N, int M, int H, int W, float *output, float epsilon);
        dlerror();
        function batchnorm_kernel = (function)dlsym(lib, "batchnorm");
        const char *dlsym_error = dlerror();
        OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of batchnorm failed: ", dlsym_error));
        batchnorm_kernel(eid_, (float *)input_flat.data(), (float *)scale_flat.data(), (float *)offset_flat.data(), (float *)mean_flat.data(), (float *)variance_flat.data(), N, M, H, W, (float *)output_flat.data(), epsilon_);
    };

private:
    void *lib;