void Compute()

in src/optimize_op.cc [780:833]
43 lines of code
5 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override
  {
    if (SMs_ == 0)
      SMs_ = GetCountSMs();

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

    float grad_scale = ctx->input(0).scalar<float>()();
    float clip_norm  = ctx->input(1).scalar<float>()();

    OpInputList x_float, x_ehalf, x_bhalf;
    ctx->input_list("x_float", &x_float);
    ctx->input_list("x_ehalf", &x_ehalf);
    ctx->input_list("x_bhalf", &x_bhalf);

    uint tensor_idx = 0;
    uint tensor_cnt = x_float.size() + x_ehalf.size() + x_bhalf.size();

    Tensor *l2norm, *scale, *sum_sqr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0,  TensorShape({}), &l2norm));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(1,  TensorShape({}), &scale));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(2,  TensorShape({ tensor_cnt }), &sum_sqr));

    float* l2norm_ptr  = l2norm->flat<float>().data();
    float* scale_ptr   = scale->flat<float>().data();
    float* sum_sqr_ptr = sum_sqr->flat<float>().data();

    for (int i = 0; i < x_float.size(); i++)
    {
      uint size = x_float[i].shape().num_elements();
      const float* x_ptr = (const float*)x_float[i].flat<FLOAT>().data();

      ReduceSumSquared<float,float4>(stream, SMs_, sum_sqr_ptr, x_ptr, size, grad_scale, saturate_, zero_infs_, zero_nans_, tensor_idx, tensor_cnt);
      tensor_idx++;
    }
    for (int i = 0; i < x_ehalf.size(); i++)
    {
      uint size = x_ehalf[i].shape().num_elements();
      const ehalf* x_ptr = (const ehalf*)x_ehalf[i].flat<EHALF>().data();

      ReduceSumSquared<ehalf,ehalf4>(stream, SMs_, sum_sqr_ptr, x_ptr, size, grad_scale, saturate_, zero_infs_, zero_nans_, tensor_idx, tensor_cnt);
      tensor_idx++;
    }
    for (int i = 0; i < x_bhalf.size(); i++)
    {
      uint size = x_bhalf[i].shape().num_elements();
      const bhalf* x_ptr = (const bhalf*)x_bhalf[i].flat<BHALF>().data();

      ReduceSumSquared<bhalf,bhalf4>(stream, SMs_, sum_sqr_ptr, x_ptr, size, grad_scale, saturate_, zero_infs_, zero_nans_, tensor_idx, tensor_cnt);
      tensor_idx++;
    }

    ComputeClipNorm(stream, l2norm_ptr, scale_ptr, sum_sqr_ptr, clip_norm, tensor_cnt);
  }