void Compute()

in src/quantize_op.cc [227:296]
59 lines of code
13 McCabe index (conditional complexity)

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

    int step = stp.dtype() == DT_INT64 ? (int)stp.scalar<int64>()() : stp.scalar<int32>()();

    ctx->set_output(0, x);

    if (freq_)
    {
      bool test = false;
      if (step != prev_step_)
      {
        prev_step_ = step;
        if (step < freq_)
        {
          for (std::vector<int32>::iterator it = first_steps_.begin() ; it != first_steps_.end(); ++it)
            if (step == *it)
            {
              test = true;
              break;
            }
        }
        else
          test = (step & (freq_-1)) == 0;
      }

      if (test)
      {
        if (SMs_ == 0)
          SMs_ = GetCountSMs();

        const V* x_ptr = (const V*)x.flat<T>().data();
        uint size = x.NumElements();

        Tensor s; OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_FLOAT, TensorShape({sizeof(QuantStats)/4}), &s));
        float* stats_ptr = s.flat<float>().data();

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

        QuantStats stats = QuantizationStats<V>(stream, SMs_, stats_ptr, x_ptr, sat_val_, ftz_val_, size);

        if (stats.max_val < max_stat_lo_) max_stat_lo_ = stats.max_val;
        if (stats.max_val > max_stat_hi_) max_stat_hi_ = stats.max_val;

        if (!logfile_.empty())
        {
          if (FILE* log = fopen(logfile_.c_str(), "a"))
          {
            float mean_stdv5 = stats.mean + stats.stdv * 5.0f;

            fprintf(log, "%.6f\t%.6f\t%3d\t%3d\t%3d\t%3d\t%3d\t%3d\t%d\t%s\n",
              stats.sat_pct,
              stats.ftz_pct,
              (*(int*)&stats.max_val >> 23)-127,
              (*(int*)&stats.mean    >> 23)-127,
              (*(int*)&stats.stdv    >> 23)-127,
              (*(int*)&mean_stdv5    >> 23)-127,
              (*(int*)&max_stat_lo_  >> 23)-127,
              (*(int*)&max_stat_hi_  >> 23)-127,
              step,
              this->name().c_str()
            );
            fclose(log);
          }
        }
      }
    }
  }