void Compute()

in src/edge_bias_op.cc [162:220]


  void Compute(OpKernelContext* ctx) override {

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

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

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

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

    // in place
    ctx->set_output(0, dy);

    Tensor* dg = nullptr;
    Tensor* db = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(1, g.shape(), &dg));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, g.shape(), &db));

              V*  dy_ptr = (V*)dy.flat<T>().data();
          float*  dg_ptr = dg->flat<float>().data();
          float*  db_ptr = db->flat<float>().data();
    const     V*   x_ptr = (const V*)x.flat<T>().data();
    const float*   g_ptr = g.flat<float>().data();
    const   int* lut_ptr = lut.flat<int32>().data();

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

    int repeat = bench_ ? bench_ : 1;
    for (int i = 0; i < repeat; i++)
      EdgeBiasBackward<V>(stream, dy_ptr, dg_ptr, db_ptr, x_ptr, g_ptr, lut_ptr, edges, MPQ, K, N, layout_);

    if (bench) delete bench;
  }