halide - 为什么我的表现不好？ (菜鸟调度)

Question

我主要是一名非常高级的程序员，因此思考 CPU 局部性等问题对我来说是非常新鲜的。

我正在研究一个基本的双线性去马赛克(用于 RGGB 传感器数据)，并且我的算法是正确的(根据结果判断)，但它的性能没有我希望的那么好(~210Mpix/s)。

这是我的代码(输入是具有单 channel RGGB 的 4640x3472 图像):

def get_bilinear_debayer(input_raw, print_nest=False):
    x, y, c = Var(), Var(), Var()

    # Clamp and move to 32 bit for lots of space for averaging.
    input = Func()
    input[x,y] = cast(
        UInt(32),
        input_raw[
            clamp(x,0,input_raw.width()-1),
            clamp(y,0,input_raw.height()-1)]
    )

    # Interpolate vertically
    vertical = Func()
    vertical[x,y] = (input[x,y-1] + input[x,y+1])/2

    # Interpolate horizontally
    horizontal = Func()
    horizontal[x,y] = (input[x-1,y] + input[x+1,y])/2

    # Interpolate on diagonals
    diagonal_average = Func()
    diagonal_average[x, y] = (
        input[x+1,y-1] + 
        input[x+1,y+1] +
        input[x-1,y-1] +
        input[x-1,y+1])/4

    # Interpolate on adjacents
    adjacent_average = Func()
    adjacent_average[x, y] = (horizontal[x,y] + vertical[x,y])/2

    red, green, blue = Func(), Func(), Func()

    # Calculate the red channel
    red[x, y, c] = select(
        # Red photosite
        c == 0, input[x, y],
        # Green photosite
        c == 1, select(x%2 == 0, vertical[x,y],
                                 horizontal[x,y]),
        # Blue photosite
        diagonal_average[x,y]
    )

    # Calculate the blue channel
    blue[x, y, c] = select(
        # Blue photosite
        c == 2, input[x, y],
        # Green photosite
        c == 1, select(x%2 == 1, vertical[x,y],
                                 horizontal[x,y]),
        # Red photosite
        diagonal_average[x,y]
    )

    # Calculate the green channel
    green[x, y, c] = select(
        # Green photosite
        c == 1, input[x,y],
        # Red/Blue photosite
        adjacent_average[x,y]
    )

    # Switch color interpolator based on requested color.
    # Specify photosite as third argument, calculated as [x, y, z] = (0, 0, 0), (0, 1, 1), (1, 0, 1), (1, 1, 2)
    # Happily works out to a sum of x mod 2 and y mod 2.
    debayer = Func()
    debayer[x, y, c] = select(c == 0, red[x, y, x%2 + y%2],
                              c == 1, green[x, y, x%2 + y%2],
                                      blue[x, y, x%2 + y%2])


    # Scheduling
    x_outer, y_outer, x_inner, y_inner, tile_index = Var(), Var(), Var(), Var(), Var()

    bits = input_raw.get().type().bits

    output = Func()
    # Cast back to the original colour space
    output[x,y,c] = cast(UInt(bits), debayer[x,y,c])
    # Reorder so that colours are calculated in order (red runs, then green, then blue)

    output.reorder_storage(c, x, y)
    # Tile in 128x128 squares
    output.tile(x, y, x_outer, y_outer, x_inner, y_inner, 128, 128)
    # Vectorize based on colour
    output.bound(c, 0, 3)
    output.vectorize(c)
    # Fuse and parallelize
    output.fuse(x_outer, y_outer, tile_index)
    output.parallel(tile_index)

    # Debugging
    if print_nest:
        output.print_loop_nest()
        debayer.print_loop_nest()
        red.print_loop_nest()
        green.print_loop_nest()
        blue.print_loop_nest()

    return output

老实说，我不知道我在这里做什么，而且我对此太陌生，不知道在哪里或该看什么。

任何关于如何改进日程安排的建议都会有帮助。我仍在学习，但很难找到反馈。

我的时间表是我能做到的最好的，但它几乎完全是反复试验。

编辑:我通过直接在函数中进行整个相邻平均求和并对 x_inner 而不是颜色进行矢量化，额外增加了 30Mpix/s。

编辑:新时间表:

# Set input bounds.
output.bound(x, 0, (input_raw.width()/2)*2)
output.bound(y, 0, (input_raw.height()/2)*2)
output.bound(c, 0, 3)

# Reorder so that colours are calculated in order (red runs, then green, then blue)
output.reorder_storage(c, x, y)
output.reorder(c, x, y)

# Tile in 128x128 squares
output.tile(x, y, x_outer, y_outer, x_inner, y_inner, 128, 128)
output.unroll(x_inner, 2).unroll(y_inner,2)

# Vectorize based on colour
output.unroll(c)
output.vectorize(c)

# Fuse and parallelize
output.fuse(x_outer, y_outer, tile_index)
output.parallel(tile_index)

编辑:最终时间表现在击败(640MP/s)Intel Performance Primitive benchmark在 CPU twice as powerful as mine 上运行:

output = Func()

# Cast back to the original colour space
output[x,y,c] = cast(UInt(bits), debayer[x,y,c])

# Set input bounds.
output.bound(x, 0, (input_raw.width()/2)*2)
output.bound(y, 0, (input_raw.height()/2)*2)
output.bound(c, 0, 3)

# Tile in 128x128 squares
output.tile(x, y, x_outer, y_outer, x_inner, y_inner, 128, 128)
output.unroll(x_inner, 2).unroll(y_inner, 2)

# Vectorize based on colour
output.vectorize(x_inner, 16)

# Fuse and parallelize
output.fuse(x_outer, y_outer, tile_index)
output.parallel(tile_index)

target = Target()
target.arch = X86
target.os = OSX
target.bits = 64
target.set_feature(AVX)
target.set_feature(AVX2)
target.set_feature(SSE41)

output.compile_jit(target)

Answer 1

确保您使用 unroll(c) 来优化每个 channel 的选择逻辑。在 x 和 y 方向上展开 2 也会有帮助:

output.unroll(x, 2).unroll(y,2)

目标是优化偶数/奇数行和列之间的选择逻辑。为了充分利用这一点，您可能还需要告诉 Halide 最小值和范围是 2 的倍数:

output.output_buffer().set_bounds(0,
                                  (f.output_buffer().min(0) / 2) * 2,
                                  (output.output_buffer().extent(0) / 2) * 2)
output.output_buffer().set_bounds(1,
                                  (f.output_buffer().min(1) / 2) * 2,
                                  (output.output_buffer().extent(1) / 2) * 2)

尽管可能值得说明更严格的约束，例如使用 128 而不是 2 来断言图 block 大小的倍数，或者如果仅支持单个相机，则仅硬连线最小值和范围以反射(reflect)实际传感器参数。