Function compute_embedding_node

cactus/graph/graph_ops_tensor.cpp:202–314

}

void compute_embedding_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes, const std::unordered_map<size_t, size_t>& node_index_map) {
    const auto& embeddings_buffer = get_input(node, 0, nodes, node_index_map);
    const auto& indices_buffer = get_input(node, 1, nodes, node_index_map);

    size_t hidden_dim = embeddings_buffer.shape[1];
    size_t num_indices = indices_buffer.total_size;
    size_t vocab_size = embeddings_buffer.original_N > 0
                      ? embeddings_buffer.original_N
                      : embeddings_buffer.shape[0];

    std::vector<float> indices_float;
    const float* indices_ptr;
    if (indices_buffer.precision == Precision::FP32) {
        indices_ptr = indices_buffer.data_as<float>();
    } else {
        indices_float.resize(num_indices);
        const int8_t* int_indices = indices_buffer.data_as<int8_t>();
        for (size_t i = 0; i < num_indices; i++) {
            indices_float[i] = static_cast<float>(int_indices[i]);
        }
        indices_ptr = indices_float.data();
    }

    __fp16* output = node.output_buffer.data_as<__fp16>();

    Precision emb_prec = embeddings_buffer.precision;
    if (PrecisionTraits::is_integer(emb_prec) && embeddings_buffer.group_size > 0) {
        const int8_t* embeddings = embeddings_buffer.data_as<int8_t>();
        const __fp16* scales = embeddings_buffer.scales_as_fp16();
        size_t group_size = embeddings_buffer.group_size;
        size_t num_groups = embeddings_buffer.num_groups;

        static const uint8_t gather_indices[4][16] = {
            {0, 1, 2, 3, 16, 17, 18, 19, 32, 33, 34, 35, 48, 49, 50, 51},  // lane 0
            {4, 5, 6, 7, 20, 21, 22, 23, 36, 37, 38, 39, 52, 53, 54, 55},  // lane 1
            {8, 9, 10, 11, 24, 25, 26, 27, 40, 41, 42, 43, 56, 57, 58, 59}, // lane 2
            {12, 13, 14, 15, 28, 29, 30, 31, 44, 45, 46, 47, 60, 61, 62, 63} // lane 3
        };

        auto load_table = [emb_prec](const int8_t* base) -> int8x16x4_t {
            if (emb_prec == Precision::INT4) {
                const uint8_t* ubase = reinterpret_cast<const uint8_t*>(base);
                int8x16_t low_a, high_a, low_b, high_b;
                unpack_int4_as_int8x16x2(ubase, high_a, low_a);
                unpack_int4_as_int8x16x2(ubase + 16, high_b, low_b);
                return {low_a, high_a, low_b, high_b};
            }
            return vld1q_s8_x4(base);
        };

        for (size_t i = 0; i < num_indices; i++) {
            size_t idx = static_cast<size_t>(indices_ptr[i]);
            if (idx >= vocab_size) {
                throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size));
            }

            size_t block = idx / 4;
            size_t lane = idx % 4;