Function main

subprojects/llama.cpp/tests/test-rope.cpp:127–263

}

int main(int /*argc*/, const char ** /*argv*/) {
    struct ggml_init_params params = {
        /* .mem_size   = */ 128*1024*1024,
        /* .mem_buffer = */ NULL,
        /* .no_alloc   = */ false,
    };

    std::vector<uint8_t> work_buffer;

    struct ggml_context * ctx0 = ggml_init(params);

    struct ggml_tensor * x;

    // rope f32
    for (int m = 0; m < 5; ++m) {
        const int ndims = 4;

        const int64_t n_rot = 128;
        const int64_t ne[4] = { 2*n_rot, 32, 73, 1 };

        const int n_past_0 = 100;
        const int n_past_2 = 33;

        struct ggml_tensor * r0;
        struct ggml_tensor * r1;
        struct ggml_tensor * r2;
        x = get_random_tensor_f32(ctx0, ndims, ne, -1.0f, 1.0f);
        int mode = -1;

        if (m < 2) {
            struct ggml_tensor * p0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2]);
            struct ggml_tensor * p1 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2]);
            struct ggml_tensor * p2 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2]);

            for (int i = 0; i < ne[2]; ++i) {
                ((int32_t *) p0->data)[i] = n_past_0 + i;
                ((int32_t *) p1->data)[i] = n_past_2 - n_past_0;
                ((int32_t *) p2->data)[i] = n_past_2 + i;
            }
            // test mode 0, 2  (standard, GPT-NeoX)
            mode = m == 0 ? GGML_ROPE_TYPE_NORMAL : GGML_ROPE_TYPE_NEOX;

            // 100, 101, 102, ..., 172
            r0 = ggml_rope(ctx0, x,  p0, n_rot, mode);
            // -67, -67, -67, ..., -67
            r1 = ggml_rope(ctx0, r0, p1, n_rot, mode); // "context swap", i.e. forget n_past_0 - n_past_2 tokens

            //  33,  34,  35, ..., 105
            r2 = ggml_rope(ctx0, x,  p2, n_rot, mode);
        } else {
            // testing multi-dimension rope position embedding mode
            struct ggml_tensor * p0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2] * 4);
            struct ggml_tensor * p1 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2] * 4);
            struct ggml_tensor * p2 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ne[2] * 4);

            int sections[4] = {16, 24, 24, 0};

            mode = (m == 2) ? GGML_ROPE_TYPE_MROPE : (m == 3) ? GGML_ROPE_TYPE_VISION : GGML_ROPE_TYPE_IMROPE;