at (@) 符號在 numpy 的矢量化 C 內核中起什么作用？

Question

Numpy 有一個驚人的 CPU 調度機制，允許它計算出 CPU 在運行時使用哪個指令集。 它使用該信息然后為請求的調用“插入”/選擇優化的 kernel，即支持 AVX 的機器運行 AVX 優化循環，只有 SSE2 的機器使用 SSE2，等等。

我試圖弄清楚它是如何工作的，所以我正在閱讀 numpy 的源代碼（即 python、C 和 C++），我遇到了一些我無法理解的事情：

// copied from numpy/core/src/umath/loops_minmax.dispatch.c.src

// contiguous input.
static inline void
simd_reduce_c_@intrin@_@sfx@(const npyv_lanetype_@sfx@ *ip, npyv_lanetype_@sfx@ *op1, npy_intp len)
{
    if (len < 1) {
        return;
    }
    const int vstep = npyv_nlanes_@sfx@;
    const int wstep = vstep*8;
    npyv_@sfx@ acc = npyv_setall_@sfx@(op1[0]);
    for (; len >= wstep; len -= wstep, ip += wstep) {
    #ifdef NPY_HAVE_SSE2
        NPY_PREFETCH(ip + wstep, 0, 3);
    #endif
        npyv_@sfx@ v0 = npyv_load_@sfx@(ip + vstep * 0);
        npyv_@sfx@ v1 = npyv_load_@sfx@(ip + vstep * 1);
        npyv_@sfx@ v2 = npyv_load_@sfx@(ip + vstep * 2);
        npyv_@sfx@ v3 = npyv_load_@sfx@(ip + vstep * 3);

        npyv_@sfx@ v4 = npyv_load_@sfx@(ip + vstep * 4);
        npyv_@sfx@ v5 = npyv_load_@sfx@(ip + vstep * 5);
        npyv_@sfx@ v6 = npyv_load_@sfx@(ip + vstep * 6);
        npyv_@sfx@ v7 = npyv_load_@sfx@(ip + vstep * 7);

        npyv_@sfx@ r01 = V_INTRIN(v0, v1);
        npyv_@sfx@ r23 = V_INTRIN(v2, v3);
        npyv_@sfx@ r45 = V_INTRIN(v4, v5);
        npyv_@sfx@ r67 = V_INTRIN(v6, v7);
        acc = V_INTRIN(acc, V_INTRIN(V_INTRIN(r01, r23), V_INTRIN(r45, r67)));
    }
    for (; len >= vstep; len -= vstep, ip += vstep) {
        acc = V_INTRIN(acc, npyv_load_@sfx@(ip));
    }
    npyv_lanetype_@sfx@ r = V_REDUCE_INTRIN(acc);
    // Scalar - finish up any remaining iterations
    for (; len > 0; --len, ++ip) {
        const npyv_lanetype_@sfx@ in2 = *ip;
        r = SCALAR_OP(r, in2);
    }
    op1[0] = r;
}

@符號在這里做什么？ 據我所知，這不是 C 中的有效字符，所以我有點迷茫。 我假設它是模板魔法，用矢量擴展的后綴替換，例如@sfx@ ，但這是如何工作的？

Answer 1

正如評論已經懷疑的那樣，這確實是一種替代/模板機制。 然而，這一次，它是一個位於 numpy 的distutils模塊（負責為 python<3.12 構建 numpy 的模塊）中的“自制軟件”。

在大多數情況下，它允許循環和變量替換。 重復塊使用/** begin repeat... */和/*end repeat*/聲明。 使用* #var = 1,2,3, ...定義循環的塊注釋中聲明了一個變量。 還支持嵌套循環，並在repeat之后使用數字標識，例如/**repeat1...表示第一個嵌套循環。 在循環內，短語@var@然后被替換為相應的變量。

示例源文件如下

/**begin repeat
* #a = 1,2,3#
* #b = 1,2,3#
*/
/**begin repeat1
* #c = ted, jim#
*/
@a@, @b@, @c@
/**end repeat1**/
/**end repeat**/

使用以下命令處理

from numpy.distutils.conv_template import process_file
from pathlib import Path

generated = process_file(

# write somewhere
Path("test.src").write_text(generated)

產生：

#line 1 "test.c.src"

/*
 *****************************************************************************
 **       This file was autogenerated from a template  DO NOT EDIT!!!!      **
 **       Changes should be made to the original source (.src) file         **
 *****************************************************************************
 */

#line 1
#line 5
#line 8
1, 1, ted

#line 8
1, 1, jim


#line 5
#line 8
2, 2, ted

#line 8
2, 2, jim


#line 5
#line 8
3, 3, ted

#line 8
3, 3, jim

---

回到我在問題中分享的片段，解開代碼的相關部分是：

/**begin repeat
 * #sfx = s8, u8, s16, u16, s32, u32, s64, u64, f32, f64#
 * #simd_chk = NPY_SIMD*8, NPY_SIMD_F32, NPY_SIMD_F64#
 * #is_fp = 0*8, 1, 1#
 * #scalar_sfx = i*8, f, d#
 */
/**begin repeat1
 * # intrin = max, min, maxp, minp#
 * # fp_only = 0, 0, 1, 1#
 */

在同一個文件中更高一點。

從本質上講，這意味着該模板將為數據類型 ( @sfx@ ) 和縮減運算符 ( @intrin@ ) 的各種組合生成高達 40 個縮減循環。 （注意：結果在預處理器階段被進一步替換，其中像npyv_load_f64這樣的宏被替換為執行請求操作的片段，具體取決於將被編譯的指令集。）