1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

void gradient_descent(float *points, uint32_t N, uint32_t M, float eta, const PolyParams* params) {
    // 预先计算常量向量
    __m256 a4 = _mm256_set1_ps(4.0f * params->a);
    __m256 b3 = _mm256_set1_ps(3.0f * params->b);
    __m256 c2 = _mm256_set1_ps(2.0 * params->c);
    __m256 d_vec = _mm256_set1_ps(params->d);
    __m256 eta_vec = _mm256_set1_ps(eta);

    for (uint32_t j = 0; j < M; ++j) {  // AVX2 一次处理 8 个 float (256b)
        uint32_t i = 0;
        for (; i + 8 <= N; i += 8) {
            __m256 x = _mm256_loadu_ps(&points[i]);

            // compute x^2, x^3
            __m256 x2 = _mm256_mul_ps(x, x);
            __m256 x3 = _mm256_mul_ps(x2, x);

            // compute gradient
            __m256 grad = _mm256_mul_ps(a4, x3);
            grad = _mm256_add_ps(grad, _mm256_mul_ps(b3, x2));
            grad = _mm256_add_ps(grad, _mm256_mul_ps(c2, x));
            grad = _mm256_add_ps(grad, d_vec);

            // update x -= eta * grad
            __m256 delta = _mm256_mul_ps(eta_vec, grad);
            x = _mm256_sub_ps(x, delta);

            // restore to points
            _mm256_storeu_ps(&points[i], x);
        }

        for (; i < N; ++i) {  // 剩余点
            float x = points[i];
            float grad = poly_gradient(x, params);
            x -= eta * grad;
            points[i] = x;
        }
    }
}