Function TrMul

tensorflow/lite/experimental/ruy/trmul.cc:266–379

}  // namespace

void TrMul(TrMulParams* params, Context* context) {
  gemmlowp::ScopedProfilingLabel label("TrMul");

  PMatrix& packed_lhs = params->packed[Side::kLhs];
  PMatrix& packed_rhs = params->packed[Side::kRhs];
  DMatrix& lhs = params->src[Side::kLhs];
  DMatrix& rhs = params->src[Side::kRhs];

  const int rows = lhs.layout.cols;
  const int cols = rhs.layout.cols;
  const int depth = lhs.layout.rows;

  int thread_count = GetThreadCount(context, rows, cols, depth);
  const auto loop_structure =
      GetLoopStructure(thread_count, rows, cols, depth,
                       params->cache_friendly_traversal_threshold);
  Allocator* allocator = context->GetMainAllocator();

  // Allocate packed matrices
  for (Side side : {Side::kLhs, Side::kRhs}) {
    if (!params->is_prepacked[side]) {
      AllocatePMatrix(allocator, &params->packed[side]);
    }
  }

  // Case of running this TrMul as a simple loop.
  // This is a good place to start reading this function: all the rest
  // of this function is just an optimized, but functionally equivalent,
  // version of that.
  if (loop_structure == LoopStructure::kSimple) {
    gemmlowp::ScopedProfilingLabel label_simple("TrMulImpl, simple loop");
    Tuning tuning = context->GetMainThreadTuning();

    const SidePair<int> origin{0, 0};
    const SidePair<int> rounded_dims{packed_lhs.layout.cols,
                                     packed_rhs.layout.cols};
    for (Side side : {Side::kLhs, Side::kRhs}) {
      if (!params->is_prepacked[side]) {
        params->RunPack(side, tuning, origin[side], rounded_dims[side]);
      }
    }
    params->RunKernel(tuning, origin, rounded_dims);

    allocator->FreeAll();
    return;
  }

  gemmlowp::ScopedProfilingLabel label_general("TrMulImpl, general case");

  auto* trace = NewTraceOrNull(&context->tracing, rows, depth, cols);
  TraceRecordStart(trace);

  // Initialize block map.
  BlockMap block_map;
  MakeBlockMap(packed_lhs.layout.cols, packed_rhs.layout.cols, depth,
               packed_lhs.layout.kernel.cols, packed_rhs.layout.kernel.cols,
               packed_lhs.data_type.size, packed_rhs.data_type.size,
               params->cache_friendly_traversal_threshold, &block_map);