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;
}