in src/ew_op.cc [942:994]
void Compute(OpKernelContext* ctx) override {
const Tensor& dy = ctx->input(0);
const Tensor& b = ctx->input(1);
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* db; OP_REQUIRES_OK(ctx, ctx->allocate_output(0, 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(1, TensorShape({gridN_, K}), &p));
partial_ptr = p->flat<float>().data();
}
}
if (partial_ptr == nullptr)
{
Tensor* p; OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape(), &p));
}
float* db_ptr = ( float*)db->flat<float>().data();
const V1* dy_ptr = (const V1*)dy.flat<T>().data();
CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();
Benchmark* bench = nullptr;
if (bench_)
{
char bench_string[256];
sprintf(bench_string, "BiasGrad (%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, 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, nullptr, dy_ptr, nullptr, nullptr, axis_, gridN_, gridK_, vec_, width_, N, K, 0, !atomics_);
if (bench) delete bench;
}