Function QuantizedGemmImpl

tensorflow/core/kernels/meta_support.cc:118–159

template <typename LeftStream, typename RightStream>
void QuantizedGemmImpl(OpKernelContext* tf_context, const quint8* a_data,
                       const quint8* b_data, qint32* c_data, int m, int n,
                       int k, int offset_a, int offset_b, int lda, int ldb,
                       int ldc) {
  typedef gemmlowp::meta::GemmParams<
      uint8_t, int32_t, LeftStream, RightStream,
      gemmlowp::meta::QuantizedStaticPreprocessedAsInt32,
      gemmlowp::meta::RowMajor>
      Params;
  Params params;

  params.m = m;
  params.n = n;
  params.k = k;

  params.lhs = reinterpret_cast<const uint8_t*>(&(a_data->value));
  params.rhs = reinterpret_cast<const uint8_t*>(&(b_data->value));
  params.result = reinterpret_cast<int32_t*>(&(c_data->value));
  params.scratch = CHECK_NOTNULL(GetScratch(tf_context));

  params.left_stream.count = k;
  params.left_stream.stride = lda;
  params.left_stream.multiplicative_sum_offset = offset_b;
  params.left_stream.additive_sum_offset = k * offset_a * offset_b;

  params.right_stream.count = k;
  params.right_stream.stride = ldb;
  params.right_stream.multiplicative_sum_offset = offset_a;
  params.right_stream.additive_sum_offset = 0;

  params.fused_kernel.kernel.count = k;
  params.fused_kernel.output_stream.stride = ldc * sizeof(int32_t);

  if (g_use_local_context) {
    LocalContext local_context(GetWorkersCount(tf_context), GetWorkersPool());
    MultiThreadGemm<LocalContext, Params>(&local_context, params);
  } else {
    auto& workers = *(tf_context->device()->tensorflow_cpu_worker_threads());
    TensorflowGemmContext context(workers.num_threads, workers.workers);
    MultiThreadGemm<TensorflowGemmContext, Params>(&context, params);
  }
}

template <typename Params, int kernel_size>
void MultiThreadTransform1D(OpKernelContext* tf_context, const Params& params) {