void Compute()

in src/ew_op.cc [769:808]

• 29 lines of code
• 10 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override {

    const Tensor& x = ctx->input(0);
    const Tensor& b = ctx->input(1);

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

    OP_REQUIRES(ctx, axis_ < x.dims() && (axis_ == 0 || axis_ == x.dims()-1), errors::InvalidArgument("BiasRelu bad axis"));

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

    OP_REQUIRES(ctx, K == b.shape().num_elements(), errors::InvalidArgument("BiasRelu missmatched channels"));

    Tensor* y = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y));

             V1* y_ptr = (V1*)y->flat<T>().data();
    const    V1* x_ptr = (const V1*)x.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, "BiasRelu     (%7d,%7d,%d,%d)", N, K, (uint)sizeof(V1), axis_);
      bench = new Benchmark(stream, bench_string, 2*N*K*sizeof(V1) + K*sizeof(float), 0, bench_);
    }

    int repeat = bench_ ? bench_ : 1;
    for (int i = 0; i < repeat; i++)
      EW_Bias_Relu<V1,V4>(stream, y_ptr, x_ptr, b_ptr, axis_, N, K, relu_);

    if (bench) delete bench;
  }