Simpler way to efficiently copy strided pixel data using RISC-V vector assembly (V)?

Problem statement

void copy(char *dst, const char *src, const ptrdiff_t stride, const int w, int h)
{
    do {
        memcpy(dst, src, w);
        dst += stride;
        src += stride;
    } while (--h);
}

I’m struggling to find a way to efficiently implement this using RISC-V scalable vectors in a way that does not make unnecessary assumptions about the size of VLEN, and which scales to both large and small blocks.

Here are approaches I tried:

Approach 1: Naive row-wise copy

.global copy_naive
copy_naive:
    mv t3, a3       # save width
    sub a2, a2, a3  # stride -= width
0:
    vsetvli t0, a3, e8, m8, ta, ma
    vle8.v v0, (a1) # load row from src
    add a1, a1, t0  # bump src ptr
    vse8.v v0, (a0) # store row to dst
    add a0, a0, t0  # bump dst ptr
    sub a3, a3, t0  # subtract count
    bnez a3, 0b     # any left in row?
    add a0, a0, a2  # dst += stride
    add a1, a1, a2  # src += stride
    mv a3, t3       # restore width
    addi a4, a4, -1 # decrement h
    bnez a4, 0b
    ret

This obviously works, and directly translates to the C code given, but it is very inefficient for any w <= VLMAX. (Consider the extreme case of w = 1)

Approach 2: Strided load/stores for very small widths

For sufficiently small, power-of-two values of w, we can make use of the strided integer load/stores, and switch to processing data in discrete columns:

.global copy_strided
copy_strided:
    li t0, 1
    beq a3, t0, 1f
    li t0, 2
    beq a3, t0, 2f
    li t0, 4
    beq a3, t0, 4f
    li t0, 8
    beq a3, t0, 8f
    j copy_naive
1:
    vsetvli t0, a4, e8, m8, ta, ma
    vlse8.v v0, (a1), a2
    add a1, a1, t0
    vsse8.v v0, (a0), a2
    add a0, a0, t0
    sub a4, a4, t0
    bnez a4, 1b
    ret
2:
    vsetvli t0, a4, e16, m8, ta, ma
    slliw t1, t0, 1
    vlse16.v v0, (a1), a2
    add a1, a1, t1
    vsse16.v v0, (a0), a2
    add a0, a0, t1
    sub a4, a4, t0
    bnez a4, 2b
    ret
4:
    vsetvli t0, a4, e32, m8, ta, ma
    slliw t1, t0, 2
    vlse32.v v0, (a1), a2
    add a1, a1, t1
    vsse32.v v0, (a0), a2
    add a0, a0, t1
    sub a4, a4, t0
    bnez a4, 4b
    ret
8:
    vsetvli t0, a4, e64, m8, ta, ma
    slliw t1, t0, 3
    vlse64.v v0, (a1), a2
    add a1, a1, t1
    vsse64.v v0, (a0), a2
    add a0, a0, t1
    sub a4, a4, t0
    bnez a4, 8b
    ret

This works very efficiently regardless of VLEN, but only for power-of-two widths up to e64 (w=8). After that, the approach breaks down. (It also doesn’t handle odd sizes of w, but I’m not too worried about those in my actual application)

Approach 3: Segmented strided loads/stores

We can extend the approach 2 using segmented loads/stores, by adding yet more special cases:

.global copy_segmented
copy_segmented:
    li t0, 16
    beq a3, t0, 16f
    li t0, 32
    beq a3, t0, 32f
    li t0, 64
    beq a3, t0, 64f
    j copy_strided
16:
    vsetvli t0, a4, e64, m4, ta, ma
    slliw t1, t0, 4
    vlsseg2e64.v v0, (a1), a2
    add a1, a1, t1
    vssseg2e64.v v0, (a0), a2
    add a0, a0, t1
    sub a4, a4, t0
    bnez a4, 16b
    ret
32:
    vsetvli t0, a4, e64, m2, ta, ma
    slliw t1, t0, 5
    vlsseg4e64.v v0, (a1), a2
    add a1, a1, t1
    vssseg4e64.v v0, (a0), a2
    add a0, a0, t1
    sub a4, a4, t0
    bnez a4, 32b
    ret
64:
    vsetvli t0, a4, e64, m1, ta, ma
    slliw t1, t0, 6
    vlsseg8e64.v v0, (a1), a2
    add a1, a1, t1
    vssseg8e64.v v0, (a0), a2
    add a0, a0, t1
    sub a4, a4, t0
    bnez a4, 64b
    ret

This continues the pattern, though in a bit more clumsy way: we now have multiple columns to worry about, which makes our life harder if we want to do anything other than simply copy data around. For example, if we wanted to add 1 to every value being copied, we now need to include a vadd.vi instruction for every column – potentially up to 8 times in the w=64 case. In general, all business logic needs to be duplicated up to 8 times.

Approach 4: Indexed gather/store

One approach to load/store an arbitrary amount of strided data in a general way would be to use an indexed load/store. In particular, this works well for multiples of 8:

.global copy_indexed
copy_indexed:
    andi t0, a3, 7      # works only for multiples of 8
    bnez t0, copy_naive
    ctzw t0, a3         # log2(width)
    addi t0, t0, -3     # log2(width/64)
    sllw a4, a4, t0     # total number of dwords to process
    vsetvli t1, a4, e64, m8, ta, ma
    srlw t1, t1, t0     # number of rows per iteration
    vid.v v0
    vsrl.vx v8, v0, t0  # y coordinates
    vmul.vx v8, v8, a2  # stride offsets
    srlw t0, a3, 3
    addi t0, t0, -1     # x coordinate mask
    vand.vx v0, v0, t0  # x coordinates
    vadd.vv v8, v0, v8  # address of every dword
    mulw a2, a2, t1     # ptr bump per iteration
0:
    vsetvli t0, a4, e64, m8, ta, ma
    vluxei64.v v0, (a1), v8
    add a1, a1, a2
    vsuxei64.v v0, (a0), v8
    add a0, a0, a2
    subw a4, a4, t0
    bnez a4, 0b
    ret

This has the advantage of naturally scaling to arbitrary vector lengths and block sizes, but requires very expensive setup code, and also only works for powers of two.

Approach 5: Some combination of the above?

Depending on runtime tests for w, VLEN etc., we can pick an appropriate primitive from the available approaches. However, this does not really feel elegant nor in the “spirit” of scalable vector extensions, which are supposed to let you write code in a way that’s agnostic to the particular vector size.

Final thoughts

This type of operation comes up extremely often in image processing, and so I’m a bit surprised that there’s no native “stride-every-N” addressing mode. In particular, it makes doing efficient image processing in a way that scales to arbitrary vector sizes very frustrating.

While my actual use case is mercifully constrained to powers of two, having so many special cases (and needing to duplicate a bunch of code as a result) makes writing functions that deal with strided pixel blocks extremely cumbersome to say the least. While I picked a toy memcpy function for my example, imagine scaling this up to hundreds of different image processing primitives.

It seems to me that there should exist a more elegant solution here. What am I missing?

New contributor

Filed under: Kiến thức lập trình - @ 01:44

Thẻ: assemblyoptimizationsimdriscv

Thiết kế website giá rẻ

Danh mục