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