void Compute()

in src/ew_op.cc [848:905]
48 lines of code
13 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override {

    const Tensor& dy = ctx->input(0);
    const Tensor&  y = ctx->input(1);
    const Tensor&  b = ctx->input(2);

    if (axis_ < 0)
      axis_ += dy.dims();

    int N = 1, K = dy.dim_size(axis_), rank = dy.dims();
    for (int i = 0; i < rank; i++)
      if (i != axis_)
        N *= dy.dim_size(i);

    Tensor* dx; OP_REQUIRES_OK(ctx, ctx->allocate_output(0, dy.shape(), &dx));
    Tensor* db; OP_REQUIRES_OK(ctx, ctx->allocate_output(1,  b.shape(), &db));

    float* partial_ptr = nullptr;
    if (axis_ != 0)
    {
      if (N_ != N)
      {
        EW_Bias_Relu_Grad_Partial(!atomics_, N, K, &gridN_, &gridK_, &vec_, &width_);
        N_ = N;
      }
      if (gridN_ > 1 && !atomics_)
      {
        Tensor* p; OP_REQUIRES_OK(ctx, ctx->allocate_output(2,  TensorShape({gridN_, K}), &p));
        partial_ptr = p->flat<float>().data();
      }
    }
    if (partial_ptr == nullptr)
    {
      Tensor* p; OP_REQUIRES_OK(ctx, ctx->allocate_output(2,  TensorShape(), &p));
    }

          float* db_ptr = (   float*)db->flat<float>().data();
             V1* dx_ptr = (      V1*)dx->flat<T>().data();
    const    V1* dy_ptr = (const V1*)dy.flat<T>().data();
    const    V1*  y_ptr = (const V1*)y.flat<T>().data();
    const float*  b_ptr = b.flat<float>().data();

    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    Benchmark* bench = nullptr;
    if (bench_)
    {
      char bench_string[256];
      sprintf(bench_string, "BiasReluGrad (%7d,%7d,%d,%d) (gn:%3d gk:%3d v:%d w:%2d)", N, K, (uint)sizeof(V1), axis_, gridN_, gridK_, vec_, width_);
      bench = new Benchmark(stream, bench_string, 3*N*K*sizeof(V1) + K*sizeof(float), 0, bench_);
    }

    int repeat = bench_ ? bench_ : 1;
    for (int i = 0; i < repeat; i++)
      EW_Bias_Relu_Grad<V1,V4>(stream, db_ptr, partial_ptr, dx_ptr, dy_ptr, y_ptr, b_ptr, axis_, gridN_, gridK_, vec_, width_, N, K, relu_, !atomics_);

    if (bench) delete bench;
  }