void Compute()

in src/edge_bias_op.cc [55:117]


  void Compute(OpKernelContext* ctx) override {

    const Tensor& x   = ctx->input(0);
    const Tensor& g   = ctx->input(1);
    const Tensor& b   = ctx->input(2);
    const Tensor& lut = ctx->input(3);

    uint rank = x.dims();
    uint N    = x.dim_size(0);
    uint MPQ  = 1, K, edges;
    // NCHW
    if (layout_ == 0)
    {
      K = x.dim_size(1);
      for (int i = 2; i < rank; i++)
        MPQ *= x.dim_size(i);

      edges = b.dim_size(1);
    }
    // NHWC
    else
    {
      K = x.dim_size(rank-1);
      for (int i = 1; i < rank-1; i++)
        MPQ *= x.dim_size(i);

      edges = b.dim_size(0);
    }
    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    const V*     x_ptr = (V*)x.flat<T>().data();
    const float* g_ptr = g.flat<float>().data();
    const float* b_ptr = b.flat<float>().data();
    const int* lut_ptr = lut.flat<int32>().data();

    V* y_ptr;
    if (inference_)
    {
      // in place
      ctx->set_output(0, x);
      y_ptr = (V*)x_ptr;
    }
    else
    {
      Tensor* y;
      OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y));
      y_ptr = (V*)y->flat<T>().data();
    }

    Benchmark* bench = nullptr;
    if (bench_)
    {
      char bench_string[256];
      sprintf(bench_string, "EdgeBias     N:%3d,K:%3d,E:%2d L:%d", N, K, edges, layout_);
      bench = new Benchmark(stream, bench_string, 2*N*K*MPQ*sizeof(V) + 2*N*K*entries_*sizeof(V) + 2*K*edges*sizeof(float), 0, bench_);
    }

    int repeat = bench_ ? bench_ : 1;
    for (int i = 0; i < repeat; i++)
      EdgeBiasForward<V>(stream, y_ptr, x_ptr, g_ptr, b_ptr, lut_ptr, edges, MPQ, K, N, layout_, inference_);

    if (bench) delete bench;
  }