< Previous post | Next post >

[2025-10-23][Daily]

Using inline assembly to implement a "valid" size convolution without padding

Results with O0, not really interesting though: Padded version 43-110 ms, Padding free version 50-170 ms

However, what really matters for the padding free version, is compilation with O3 optimizations, because it completly removes conditionnal branches in the produced assembly for the sub_sat function which is extremely useful in our for loops. This works with saturation arithmetic that can be implemented very efficiently as such:

size_t add_sat(size_t a, size_t b) {
    size_t c = a + b;
    if (c < a) /* Can only happen due to overflow */
        c = -1;
    return c;
}

which produces really nice assembly:

add_sat:
        mov     rax, rdi  # c, a
        mov     rdx, -1   # tmp92,
        add     rax, rsi  # c, b
        cmovc   rax, rdx  # c,, c, tmp92
        ret

This works because we can reuse some flags produce by add that indicates if c overflowed, so the if can be simplified into using cmovc, see: https://www.felixcloutier.com/x86/cmovcc. However, sub function isn't as nicely compiled:

size_t sub_sat(size_t a, size_t b)
{
    size_t c = a - b;
    if (c > a)
        c = 0;
    return c;
}
sub_sat:
        mov     rax, rdi  # c, a
        mov     edx, 0    # tmp92,
        sub     rax, rsi  # c, b
        cmp     rdi, rax  # a, c
        cmovb   rax, rdx  # c,, c, tmp92
        ret

Here it need an additional cmp instruction because the compiler is not able to guess to c > a is equivalent to using the unsigned carry flag. This may be sufficient, but if we really want, we could write inline assembly to obtain even better results:

size_t sub_sat_asm(size_t a, size_t b)
{
    size_t tmp = 0;
               // { AT&T    | Intel }  syntax alternatives.  The other versions without this  will break with -masm=intel
    asm ("sub     { %[b],%[a]   | %[a],  %[b] }\n\t"
         "cmovnc  { %[a],%[dst] | %[dst],%[a] }"
         : [dst] "+r" (tmp), [a] "+&r" (a)
         : [b] "g" (b)
         );
    return tmp;
}
sub_sat_asm:
        xor     eax, eax  # tmp
        sub     rdi, rsi  # a, b
        cmovnc  rax, rdi  # tmp, a
        ret

In O3 we obtain pretty close results in terms of pure codelet execution, however we gain clear advantages: no need to perform copy and add padding, reduced number of task and synchronization, better parallelization, less memory usage.

Results with O3: Padded version ~20 ms, Padding free version ~10 ms

Obviously the gains are not only runtime wise, but also for memory consumption, as we don't need to 1. store the zero padding and 2. pontentially duplicate the tensor so that we have a padded and unpadded version. However it is clear that we could pre-compute ranges at model initialization and have clean loops, without inline assemebly.

Here is the full implementation if you are interested:

// Saturating substraction operation. If the result was supposed to underflow, return 0 instead.
// Used for convolution backward without padding.
#if defined (__i386__) || defined (__x86_64__)
static inline size_t sub_sat(size_t a, size_t b)
{  
    size_t tmp = 0;

    // { AT&T    | Intel }  syntax alternatives.
    // The other versions without this  will break with -masm=intel
    __asm__("sub     { %[b],%[a]   | %[a],  %[b] }\n\t"
            "cmovnc  { %[a],%[dst] | %[dst],%[a] }"
            : [dst] "+r" (tmp), [a] "+&r" (a)
            : [b] "g" (b)
    );
    return tmp;
}
#else
// sub_sat equivalent without using inlined assembly. Not perfectly optimized though.
static inline size_t sub_sat(size_t a, size_t b)
{
    size_t c = a - b;
    // This if can only happen if `c` underflow'd, thus the compiler can reuse overflow flag and
    // use a conditional move to execute c = 0;
    if (c > a)
        c = 0;
    return c;
}
#endif

/*
 * This function implements a "full" convolution with padding free input. This means that the output
 * is larger than the input, but we don't need to use zero padding and compute useless operations on
 * the padding.
 * We do that by computing start/end indexes of each kernel window so that we ignore out-of-bound
 * kernel values.
 * It uses saturating arithmetic trick to prevent conditionnal branches to appear in for loops.
 *
 *      kernel size
 *     ┌───────────┐
 *     ▼           ▼
 *        Actual range we want
 *             ┌───┐
 *             ▼   ▼
 *      -2  -1   0   1   2   3     
 *     ┌ ─ ┬ ─ ┬ ─ ┬ ─ ┬ ─ ┬ ─ ┐
 *  -2   0   0   0   0   0   0                            0   1   2   3
 *     ├ ─ ┼ ─ ┼ ─ ┼ ─ ┼ ─ ┼ ─ ┤        0   1   2       ┌───┬───┬───┬───┐
 *  -1   0   0   0   0   0   0        ┌───┬───┬───┐   0 │   │   │   │   │
 *     ├ ─ ┼ ─ ┼───┼───┼ ─ ┼ ─ ┤    0 │   │   │   │     ├───┼───┼───┼───┤
 *   0   0   0 │   │   │ 0   0        ├───┼───┼───┤   1 │   │   │   │   │
 *     ├ ─ ┼ ─ ┼───┼───┼ ─ ┼ ─ ┤    1 │   │   │   │     ├───┼───┼───┼───┤
 *   1   0   0 │   │   │ 0   0        ├───┼───┼───┤   2 │   │   │   │   │
 *     ├ ─ ┼ ─ ┼───┼───┼ ─ ┼ ─ ┤    2 │   │   │   │     ├───┼───┼───┼───┤
 *   2   0    Input buffer   0        └───┴───┴───┘   3 │   │   │   │   │
 *     ├ ─ ┼ ─ ┼ ─ ┼ ─ ┼ ─ ┼ ─ ┤         Kernel         └───┴───┴───┴───┘
 *   3   0   0   0   0   0   0       (omitted z dim)         Output
 *     └ ─ ┴ ─ ┴ ─ ┴ ─ ┴ ─ ┴ ─ ┘                         (omitted z dim)
 *            Fake Padding
 */
void convolution_2d_backward_input_padding_free(void* buffers[3], void* cl_arg)
{
    // Input matrix, here the gradients output of the layer just after the convolution
    auto in = STARPU_MATRIX_GET(buffers[0]);
    auto in_p = (dahl_fp const*)in.ptr;

    // Kernel block, here the filters (weights) associated to the convolution
    auto ker = STARPU_BLOCK_GET(buffers[1]);
    auto ker_p = (dahl_fp const*)ker.ptr;

    // Output block, here the loss derivative of the input
    auto out = STARPU_BLOCK_GET(buffers[2]);
    auto out_p = (dahl_fp*)out.ptr;

    assert(out.nx == in.nx + ker.nx - 1);
    assert(out.ny == in.ny + ker.ny - 1);
    assert(out.nz == ker.nz);

    // Compute padding size required to produce the correct output shape.
    size_t pad_nx = ker.nx - 1;
    size_t pad_ny = ker.ny - 1;

    // loop through i,j,k on axes x,y,z of the output block
    for (size_t k = 0; k < out.nz; k++)
    {
        for (size_t j = 0; j < out.ny; j++)
        {
            for (size_t i = 0; i < out.nx; i++)
            {
                dahl_fp cell_res = 0.0F;

                // [`y_start`, `y_end`[ corresponds to the window where we pull values from the                 
                // input matrix, `y_ker` simulates how shifted the kernel is, its values gets    
                // incremented each loop to simulate shifting. Same principle for x dimension.   
                
                // Here, we will always shift our kernel up from `pad_ny` values relative to j.
                // Using sub_sat let us handle cases where the kernel has values outside the matrix,
                // underflow will saturate to value 0.
                size_t y_start = sub_sat(j, pad_ny);

                // We can use the same trick for the end index, but here we "reverse" indexes, so we
                // can use the underflow mechanic. Minus 1 because this is the remaining part of the
                // kernel (pad_ny being included in y_start).
                size_t y_end = sub_sat(in.ny - 1, j); 

                // Reverse indexes once again to retrieve the actual index of the end. When we reach
                // the bottom of the matrix, the kernel will be outside bounds so `y_end` values
                // will be saturated to 0, thus `y_end` will equal `in_ny`.
                y_end = in.ny - y_end;

                // Here we compute where the kernel should start by substracting the end index.
                // If it overflows, it means there's no more need to shift the kernel, all start at
                // 0
                size_t y_ker = sub_sat(ker.ny, y_end);

                // Loop through l,m on axes x,y of computed window
                for (size_t m = y_start; m < y_end; m++)
                {
                    size_t x_start = sub_sat(i, pad_nx);
                    size_t x_end = sub_sat(in.nx - 1, i); 
                    x_end = in.nx - x_end;
                    size_t x_ker = sub_sat(ker.nx, x_end);

                    for (size_t l = x_start; l < x_end; l++)
                    {
                        // Reverse indexes x_ker and y_ker so we don't actually have to rotate(180)
                        // the kernel matrix. However we still use k for axis z because we write
                        // each result for the current `kernel` channel into the same `out` channel.
                        dahl_fp kernel_value = ker_p[
                           (k * ker.ldz) + ((ker.ny - 1 - y_ker) * ker.ldy) + (ker.nx - 1 - x_ker)];
                        dahl_fp in_value = in_p[(m * in.ld) + l];

                        cell_res += in_value * kernel_value;
                        x_ker++;
                    }

                    y_ker++;
                }

                // Set the corresponding value for index i,j,k
                out_p[(k * out.ldz) + (j * out.ldy) + i] = cell_res;
            }
        }
    }
}

< Previous post | Next post >