Method Run

tensorflow/lite/experimental/ruy/trmul.cc:63–115

        local_packed{nullptr, nullptr} {}

  void Run() override {
    TraceRecordThreadStart(thread_id, trace);

    for (Side side : {Side::kLhs, Side::kRhs}) {
      if (!params->is_prepacked[side]) {
        const int size = NumBlocksPerSide(side, block_map);
        local_allocator->Allocate(size, &local_packed[side]);
        memset(local_packed[side], 0, size * sizeof(bool));
      }
    }

    const int num_blocks = NumBlocks(block_map);

    const Tuning tuning = tuning_resolver->Resolve();

    TraceRecordThreadLoopStart(thread_id, trace);

    SidePair<int> block;
    SidePair<int> start;
    SidePair<int> end;

    // Each thread starts by initially reserving the block whose id
    // is the thread id.
    int block_id = thread_id;
    TraceRecordBlockReserved(thread_id, block_id, trace);

    while (block_id < num_blocks) {
      // Reserve the next block to handle. In order to hide the latency
      // (typically comparable to an access to the level of data cache that
      // is shared among CPU cores, e.g. 60 cycles on an ARM CPU as of 2019)
      // of this atomic operation, we structure this code so as to avoid
      // immediately depending on the `next_n` result.
      const int next_block_id =
          atomic_block_id->fetch_add(1, std::memory_order_relaxed);
      TraceRecordBlockReserved(thread_id, next_block_id, trace);
      // Get coordinates of the current block to handle, in "block space".
      GetBlockByIndex(block_map, block_id, &block);
      // Get coordinates of the current block to handle, in matrix space.
      GetBlockMatrixCoords(block_map, block, &start, &end);
      // Maybe pack the current LHS/RHS block, if not already packed.
      EnsurePacked(block, start, end, tuning);
      // Actually do matrix multiplication work
      params->RunKernel(tuning, start, end);
      TraceRecordBlockFinished(thread_id, block_id, trace);
      // Move on to the next block as obtained by the atomic increment
      // at the start of this while loop iteration.
      block_id = next_block_id;
    }

    local_allocator->FreeAll();

    TraceRecordThreadEnd(thread_id, trace);
  }

 private:
  // Tries to pack a block, without blocking.