Julia中实现@parallel_map宏：高效多核循环避免Any类型

核心思路 #

通过宏在编译期解析循环结构，先推断循环体的返回值类型，再创建对应类型的数组，最后用多线程填充，从根源上避免Vector{Any}带来的性能损耗。

基础实现代码 #

macro parallel_map(for_expr)
    # 校验输入必须是for循环表达式
    @assert for_expr.head == :for "Input must be a for loop"
    loop = for_expr.args[1]
    body = for_expr.args[2]
    # 解析循环变量与迭代器
    @assert loop.head == :(=) "Invalid for loop syntax"
    var = loop.args[1]
    iter = loop.args[2]

    quote
        iter_inst = $(esc(iter))
        len = length(iter_inst)
        # 通过单次执行循环体推断返回值类型
        elem_type = typeof($(esc(body)))
        # 初始化指定类型的数组
        results = Vector{elem_type}(undef, len)
        # 多线程填充数组
        Threads.@threads for (idx, $(esc(var))) in enumerate(iter_inst)
            results[idx] = $(esc(body))
        end
        results
    end
end

扩展：支持手动指定类型 #

如果迭代器可能为空，或者需要强制指定返回类型，可以扩展宏的参数：

macro parallel_map(T, for_expr)
    @assert for_expr.head == :for "Input must be a for loop"
    loop = for_expr.args[1]
    body = for_expr.args[2]
    var = loop.args[1]
    iter = loop.args[2]

    quote
        iter_inst = $(esc(iter))
        len = length(iter_inst)
        # 使用用户指定的类型初始化数组
        results = Vector{$(esc(T))}(undef, len)
        Threads.@threads for (idx, $(esc(var))) in enumerate(iter_inst)
            results[idx] = $(esc(body))
        end
        results
    end
end

使用示例：

# 自动推断类型
results = @parallel_map for i in 1:10
    i^2
end

# 手动指定类型
results = @parallel_map Int64 for i in 1:10
    i^2
end

效果验证 #

执行上述代码后，typeof(results)会返回Vector{Int64}而非Vector{Any}，数组类型稳定，不会出现类型不稳定导致的性能下降。

内容的提问来源于stack exchange，提问作者Alex Craft