Method load_embedding

include/core/solver.h:1349–1386

* @brief Load an embedding matrix and its moment matrices to GPU cache * @param id id of embedding matrix * * The actual operation depends on the protocol. It is safe to call this function multiple times. */

     * The actual operation depends on the protocol. It is safe to call this function multiple times.
     */
    void load_embedding(int id) {
        Protocol protocol = protocols[id];
        if ((protocol & kSharedWithPredecessor) && ((protocol & kGlobal) || head_partition_id == tail_partition_id)) {
            embeddings[id] = embeddings[id - 1];
            moments[id] = moments[id - 1];
            return;
        }
        if (id > 0 && embeddings[id] == embeddings[id - 1]) {
            embeddings[id] = std::make_shared<Memory<Vector, Index>>(device_id, 0, work_stream);
            gradients[id] = std::make_shared<Memory<Vector, Index>>(device_id, 0, work_stream);
            moments[id] = std::make_shared<std::vector<Memory<Vector, Index>>>();
            for (int i = 0; i < num_moment; i++)
                moments[id]->push_back(Memory<Vector, Index>(device_id, 0, work_stream));
        }
        auto &global_embedding = *(solver->embeddings[id]);
        auto &global_moment = *(solver->moments[id]);
        auto &embedding = *embeddings[id];
        auto &moment = *moments[id];
        auto &gradient = *gradients[id];
        std::vector<Index> none, *mapping = &none;

        if (protocol & kHeadPartition)
            mapping = &head_global_ids;
        if (protocol & kTailPartition)
            mapping = &tail_global_ids;

        embedding.gather(global_embedding, *mapping);
        embedding.to_device_async();

        // only load partitioned moments, or global moments if uninitialized
        if (solver->is_train)
            if (!(protocol & kGlobal) || (num_moment && moment[0].count == 0)) {
                for (int i = 0; i < num_moment; i++) {
                    moment[i].gather(global_moment[i], *mapping);
                    moment[i].to_device_async();
                }
            }
    }

    /**
     * @brief Write back an embedding matrix and its moment matrices from GPU cache