Repository: opencv/opencv
File: modules/dnn/src/layers/attention_layer.cpp (C++)

Code before:

static void precompRotationTable(float *data,
                                  size_t seq_len,
                                  size_t d) {
    // RoPE precomputation
    // RoPE is a positional encoding method used in transformer models.
    // It uses sine and cosine functions to encode the position of tokens in a sequence
    // initially introduced for NLP in https://arxiv.org/pdf/2104.09864

    // assume data is of shape [seq_ken,d]
    const float  logBase = std::log(10000.0f);
    const float  inv_d   = 1.0f / float(d);
    const size_t d_half = d / 2;
    for (size_t pos = 0; pos < seq_len; ++pos) {

        size_t i = 0;
        float* data_ptr = data + pos * d;

#if (CV_SIMD || CV_SIMD_SCALABLE)
        const size_t w = VTraits<v_float32>::vlanes();
        const v_float32 v_logBase = v_setall_f32(logBase);
        const v_float32 v_inv_d   = v_setall_f32(inv_d);
        const v_float32 v_neg2    = v_setall_f32(-2.0f);

        for (; i + w <= d_half; i+=w) {
            int idx_buf[CV_SIMD_WIDTH];
            for (int k = 0; k < int(w); ++k)
                idx_buf[k] = int(i + k);
            // [i, i+1, …, i+w-1]
            v_float32 v_idx = v_cvt_f32(vx_load(idx_buf));
            // [10_000^(-i/d), 10_000^(-(i+1)/d), …, 10_000^(-(i+w-1)/d)]
            v_float32 v_theta   = v_exp(v_mul(v_mul(v_neg2, v_mul(v_idx, v_inv_d)), v_logBase));
            v_theta = v_mul(v_setall_f32(float(pos)), v_theta);
            v_float32 sin_v, cos_v;
            v_sincos(v_theta, sin_v, cos_v);
            // store back with interleave
            v_store_interleave(data_ptr + 2*i, sin_v, cos_v);
        }
#endif
        // scalar tail
        for (; i < d_half; i+=1)
        {
            float theta = pos * std::exp(-2.f * i/d * logBase);
            data_ptr[2*i    ] = std::sin(theta);
            data_ptr[2*i + 1] = std::cos(theta);
        }
    }
}


// Operator spec: https://github.com/microsoft/onnxruntime/blob/v1.16.1/docs/ContribOperators.md#com.microsoft.Attention

Review comment (suggestion):
CV_SIMD_WIDTH is compile time constant. It may not work correctly with _SCALABLE branch. please use `VTraits<xxx>::max_nlanes` instead. For fixed-size SIMD it works in the same way.

Code after:

static void precompRotationTable(float *data,
                                  size_t seq_len,
                                  size_t d) {
    // RoPE precomputation
    // RoPE is a positional encoding method used in transformer models.
    // It uses sine and cosine functions to encode the position of tokens in a sequence
    // initially introduced for NLP in https://arxiv.org/pdf/2104.09864

    // assume data is of shape [seq_ken,d]
    const float  logBase = std::log(10000.0f);
    const float  inv_d   = 1.0f / float(d);
    const size_t d_half = d / 2;
    for (size_t pos = 0; pos < seq_len; ++pos) {

        size_t i = 0;
        float* data_ptr = data + pos * d;

#if (CV_SIMD || CV_SIMD_SCALABLE)
        const size_t w = VTraits<v_float32>::vlanes();
        const v_float32 v_logBase = vx_setall_f32(logBase);
        const v_float32 v_inv_d   = vx_setall_f32(inv_d);
        const v_float32 v_neg2    = vx_setall_f32(-2.0f);

        for (; i + w <= d_half; i+=w) {
            int idx_buf[VTraits<v_float32>::max_nlanes];
            for (int k = 0; k < int(w); ++k)
                idx_buf[k] = int(i + k);
            // [i, i+1, …, i+w-1]
            v_float32 v_idx = v_cvt_f32(vx_load(idx_buf));
            // [10_000^(-i/d), 10_000^(-(i+1)/d), …, 10_000^(-(i+w-1)/d)]
            v_float32 v_theta = v_exp(v_mul(v_mul(v_neg2, v_mul(v_idx, v_inv_d)), v_logBase));
            v_theta = v_mul(vx_setall_f32(float(pos)), v_theta);
            v_float32 sin_v, cos_v;
            v_sincos(v_theta, sin_v, cos_v);
            // store back with interleave
            v_store_interleave(data_ptr + 2*i, sin_v, cos_v);
        }
#endif
        // scalar tail
        for (; i < d_half; i+=1)
        {
            float theta = pos * std::exp(-2.f * i * inv_d * logBase);
            data_ptr[2*i    ] = std::sin(theta);
            data_ptr[2*i + 1] = std::cos(theta);
        }
    }
}


// Operator spec: https://github.com/microsoft/onnxruntime/blob/v1.16.1/docs/ContribOperators.md#com.microsoft.Attention