void Compute()

in src/layer_norm_op.cc [227:304]
68 lines of code
13 McCabe index (conditional complexity)

  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& b  = ctx->input(3);
    const Tensor& mean = ctx->input(4);
    const Tensor& rstd = ctx->input(5);

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

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

    if (K_ == 0)
    {
      K_    = K / S_;
      rcpK_ = 1.0f / (float)K_;
      SMs_  = GetCountSMs();
    }

    Tensor* dx = nullptr;
    Tensor* dg = nullptr;
    Tensor* db = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &dx));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(1, g.shape(), &dg));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, b.shape(), &db));

    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));
    float* p1_ptr = nullptr;
    float* p2_ptr = nullptr;
    if (axis_ == 0)
    {
      p1_ptr = P1->flat<float>().data();
      p2_ptr = P2->flat<float>().data();
    }

             V1*   dx_ptr = (V1*)dx->flat<T>().data();
          float*   dg_ptr = dg->flat<float>().data();
          float*   db_ptr = db->flat<float>().data();
    const    V1*   dy_ptr = (const V1*)dy.flat<T>().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();
    const float* mean_ptr = mean.flat<float>().data();
    const float* rstd_ptr = rstd.flat<float>().data();

    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    Benchmark* bench = nullptr;
    if (bench_) bench = new Benchmark(stream, "LayerNormBackward", N*S_*K_*5*sizeof(T), 0, repeat_);

    for (int r = 0; r < repeat_; r++)
      if (axis_ == 0)
        LayerNormBackward_CN<V1,V4>(stream, SMs_, dx_ptr, dg_ptr, db_ptr, p1_ptr, p2_ptr, dy_ptr, x_ptr, g_ptr, b_ptr, mean_ptr, rstd_ptr, epsilon_, K_, N, rcpK_, relu_);
      else
      {
        if (S_ > 1 || K_ <= 1024*8)
          LayerNormSegmentedBackward_NC<V1,V4>(stream, SMs_, dx_ptr, dg_ptr, db_ptr, dy_ptr, x_ptr, g_ptr, b_ptr, mean_ptr, rstd_ptr, epsilon_, N, S_, K_, rcpK_, relu_, atomics_);
        else
          LayerNormBackward_NC<V1,V4>(stream, SMs_, dx_ptr, dg_ptr, db_ptr, dy_ptr, x_ptr, g_ptr, b_ptr, mean_ptr, rstd_ptr, epsilon_, K_, N, rcpK_, relu_);
      }

    if (bench) delete bench;
  }