void Compute()

in src/quantize_op.cc [112:188]
61 lines of code
13 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();

    const Tensor&  x = ctx->input(0);
    const Tensor& em = ctx->input(1);
    OpInputList    e;  ctx->input_list("entropy", &e);

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

          V*  y_ptr = (      V*)y->flat<T>().data();
    const V*  x_ptr = (const V*)x.flat<T>().data();
    uint*     e_ptr = e.size() > 0 ? (uint*)e[0].flat<float>().data() : NULL;
    int64* emax_ptr = (int64*)em.flat<int64>().data();

    uint size = x.NumElements();

    int emax  = (int)*emax_ptr;
    //printf("emax: %d\n", emax);
    if (emax != last_exp_)
      UpdateExponent(emax);

    if (freq_ && ((count_ & (pow2_-1)) == 0))
    {
      if ((pow2_ << 1) <= freq_)
      {
        if (pow2_count_ == freq2_)
        {
          pow2_     <<= 1;
          pow2_count_ = 0;
        }
        pow2_count_  += 1;
      }
      Tensor s; OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_FLOAT, TensorShape({sizeof(QuantStats)/4}), &s));
      float* stats_ptr = s.flat<float>().data();

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

      float mode_max = mode_ ? stats.mean + stats.stdv * stdv_mul_ : stats.max_val;

      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;

      *emax_ptr = UpdateExponent((*(int*)&mode_max >> 23)-127 + bias_pad_);

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

          fprintf(log, "%.3f\t%.3f\t%3d\t%3d\t%3d\t%3d\t%3d\t%3d\t%3d\t%3d\t%d\t%s\n",
            stats.sat_pct,
            stats.ftz_pct,
            (*(int*)&max_float_    >> 23)-127,
            (*(int*)&min_float_    >> 23)-127,
            (*(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,
            count_,
            this->name().c_str()
          );
          fclose(log);
        }
      }
    }

    Quantize<V>(stream, SMs_, e_ptr, y_ptr, x_ptr, round_scale_, trunc_mask_, max_float_, min_float_, exp_norm_, size, stoch_);

    count_ += 1;
  }