void Compute()

in src/layer_norm_op.cc [89:170]
73 lines of code
15 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override
  {
    const Tensor& x = ctx->input(0);
    const Tensor& g = ctx->input(1);
    const Tensor& b = ctx->input(2);

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

    int N = 1, K = x.dim_size(axis_), last_dim = x.dims()-1;
    TensorShape shapeN;
    shapeN.AddDim(S_);
    for (int i = 0; i <= last_dim; i++)
    {
      if (i != axis_)
      {
        shapeN.AddDim(x.dim_size(i));
        N *= x.dim_size(i);
      }
    }
    OP_REQUIRES(ctx, axis_ != 0 || (N & 3) == 0, errors::InvalidArgument("Sum of non-feature axis dims needs to be multiple of 4 for feature axis=0."));

    if (K_ == 0)
    {
      OP_REQUIRES(ctx, K == g.shape().num_elements(), errors::InvalidArgument("Bad Gain Shape"));
      OP_REQUIRES(ctx, K == b.shape().num_elements(), errors::InvalidArgument("Bad Bias Shape"));
      OP_REQUIRES(ctx, (K % S_) == 0, errors::InvalidArgument("Shape not evenly devided by segments"));
      OP_REQUIRES(ctx, axis_ != 0 || S_ == 1, errors::InvalidArgument("Segments only implemented on axis=1 for now"));
      K_    = K / S_;
      rcpK_ = 1.0f / (float)K_;
      SMs_  = GetCountSMs();
    }

    Tensor*    y = nullptr;
    Tensor* mean = nullptr;
    Tensor* rstd = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(),    &y));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(1,    shapeN, &mean));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(2,    shapeN, &rstd));

    float* p1_ptr = nullptr;
    float* p2_ptr = nullptr;
    Tensor* P1 = nullptr;
    Tensor* P2 = nullptr;
    TensorShape shapeP;
    if (axis_ == 0)
    {
      shapeP.AddDim(SMs_*2);
      shapeP.AddDim(N);
    }
    OP_REQUIRES_OK(ctx, ctx->allocate_output(3,  shapeP, &P1));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(4,  shapeP, &P2));
    if (axis_ == 0)
    {
      p1_ptr = P1->flat<float>().data();
      p2_ptr = P2->flat<float>().data();
    }
             V1*    y_ptr = (V1*)y->flat<T>().data();
          float* mean_ptr = mean->flat<float>().data();
          float* rstd_ptr = rstd->flat<float>().data();
    const    V1*    x_ptr = (const V1*)x.flat<T>().data();
    const float*    g_ptr = g.flat<float>().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_) bench = new Benchmark(stream, "LayerNormForward", N*S_*K_*2*sizeof(T), 0, repeat_);

    for (int r = 0; r < repeat_; r++)
      if (axis_ == last_dim)
      {
        if (S_ > 1 || K_ <= 1024*8)
          LayerNormSegmentedForward_NC<V1,V4>(stream, SMs_, y_ptr, mean_ptr, rstd_ptr, x_ptr, g_ptr, b_ptr, epsilon_, N, S_, K_, rcpK_, relu_);
        else
          LayerNormForward_NC<V1,V4>(stream, SMs_, y_ptr, mean_ptr, rstd_ptr, x_ptr, g_ptr, b_ptr, epsilon_, K_, N, rcpK_, relu_);
      }
      else
        LayerNormForward_CN<V1,V4>(stream, SMs_, y_ptr, mean_ptr, rstd_ptr, p1_ptr, p2_ptr, x_ptr, g_ptr, b_ptr, epsilon_, K_, N, rcpK_, relu_);

    if (bench) delete bench;
  }