[AVX2] SAD pattern detection is too strict

[adworacz]

Reference code: Zig Godbolt

Some opportunities for producing optimized sum of absolute differences (SAD) calculations are being missed. It looks like prior support for this was overly restrictive.

Essentially, the absolute difference is being calculated, when it should just be handled by the dedicated SAD instruction.

Here's the code inline:

const block_width = 8;
const T = u8;
const VT = @Vector(block_width, T);

export fn sad(noalias srcp: [*]const u8, noalias refp: [*]const u8, height: usize, stride: usize) u32 {
    const src = srcp[0..height * stride];
    const ref = refp[0..height * stride];
    
    var sum: u32 = 0;

    const s: VT = src[0*stride..][0..block_width].*;
    const r: VT = ref[0*stride..][0..block_width].*;

    // Should work, but doesn't.
    const absdiff = @max(s,r) - @min(s,r);
    sum += @reduce(.Add, absdiff);
        
    // Should work, but doesn't
    //const VTI = @Vector(block_width, i16);
    //sum += @reduce(.Add, @abs(@as(VTI, s) - @as(VTI, r)));

    // Does work
    //const VTI = @Vector(block_width, i32);
    //sum += @reduce(.Add, @abs(@as(VTI, s) - @as(VTI, r)));
    
    return sum;
}

Which produces:

sad:
        push    rbp
        mov     rbp, rsp
        vmovq   xmm0, qword ptr [rdi]
        vmovq   xmm1, qword ptr [rsi]
        vpminub xmm2, xmm0, xmm1
        vpmaxub xmm0, xmm0, xmm1
        vpxor   xmm1, xmm1, xmm1
        vpsubb  xmm0, xmm0, xmm2
        vpsadbw xmm0, xmm0, xmm1
        vpextrb eax, xmm0, 0
        pop     rbp
        ret

But it should be:

sad:
        push    rbp
        mov     rbp, rsp
        vmovq   xmm0, qword ptr [rdi]
        vmovq   xmm1, qword ptr [rsi]
        vpsadbw xmm0, xmm0, xmm1
        vmovd   eax, xmm0
        pop     rbp
        ret

[llvmbot]

@llvm/issue-subscribers-backend-x86

Author: Austin Dworaczyk Wiltshire (adworacz)

Reference code: [Zig Godbolt](https://zig.godbolt.org/z/c5ezdxojG)

Some opportunities for producing optimized sum of absolute differences (SAD) calculations are being missed. It looks like prior support for this was overly restrictive.

Essentially, the absolute difference is being calculated, when it should just be handled by the dedicated SAD instruction.

Here's the code inline:

const block_width = 8;
const T = u8;
const VT = @<!-- -->Vector(block_width, T);

export fn sad(noalias srcp: [*]const u8, noalias refp: [*]const u8, height: usize, stride: usize) u32 {
    const src = srcp[0..height * stride];
    const ref = refp[0..height * stride];
    
    var sum: u32 = 0;

    const s: VT = src[0*stride..][0..block_width].*;
    const r: VT = ref[0*stride..][0..block_width].*;

    // Should work, but doesn't.
    const absdiff = @<!-- -->max(s,r) - @<!-- -->min(s,r);
    sum += @<!-- -->reduce(.Add, absdiff);
        
    // Should work, but doesn't
    //const VTI = @<!-- -->Vector(block_width, i16);
    //sum += @<!-- -->reduce(.Add, @<!-- -->abs(@<!-- -->as(VTI, s) - @<!-- -->as(VTI, r)));

    // Does work
    //const VTI = @<!-- -->Vector(block_width, i32);
    //sum += @<!-- -->reduce(.Add, @<!-- -->abs(@<!-- -->as(VTI, s) - @<!-- -->as(VTI, r)));
    
    return sum;
}

Which produces:

sad:
        push    rbp
        mov     rbp, rsp
        vmovq   xmm0, qword ptr [rdi]
        vmovq   xmm1, qword ptr [rsi]
        vpminub xmm2, xmm0, xmm1
        vpmaxub xmm0, xmm0, xmm1
        vpxor   xmm1, xmm1, xmm1
        vpsubb  xmm0, xmm0, xmm2
        vpsadbw xmm0, xmm0, xmm1
        vpextrb eax, xmm0, 0
        pop     rbp
        ret

But it should be:

sad:
        push    rbp
        mov     rbp, rsp
        vmovq   xmm0, qword ptr [rdi]
        vmovq   xmm1, qword ptr [rsi]
        vpsadbw xmm0, xmm0, xmm1
        vmovd   eax, xmm0
        pop     rbp
        ret

[RKSimon]

define i32 @sad(ptr noalias nocapture nonnull readonly align 1 %0, ptr noalias nocapture nonnull readonly align 1 %1, i64 %2, i64 %3) {
Entry:
  %4 = load <8 x i8>, ptr %0, align 1
  %5 = load <8 x i8>, ptr %1, align 1
  %6 = tail call <8 x i8> @llvm.umax.v8i8(<8 x i8> %4, <8 x i8> %5)
  %7 = tail call <8 x i8> @llvm.umin.v8i8(<8 x i8> %4, <8 x i8> %5)
  %8 = sub nuw <8 x i8> %6, %7
  %rdx.shuf = shufflevector <8 x i8> %8, <8 x i8> poison, <8 x i32> <i32 4, i32 5, i32 6, i32 7, i32 poison, i32 poison, i32 poison, i32 poison>
  %bin.rdx = add <8 x i8> %8, %rdx.shuf
  %rdx.shuf15 = shufflevector <8 x i8> %bin.rdx, <8 x i8> poison, <8 x i32> <i32 2, i32 3, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison>
  %bin.rdx16 = add <8 x i8> %bin.rdx, %rdx.shuf15
  %rdx.shuf17 = shufflevector <8 x i8> %bin.rdx16, <8 x i8> poison, <8 x i32> <i32 1, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison>
  %bin.rdx18 = add <8 x i8> %bin.rdx16, %rdx.shuf17
  %9 = extractelement <8 x i8> %bin.rdx18, i32 0
  %10 = zext i8 %9 to i32
  ret i32 %10
}

declare <8 x i8> @llvm.umax.v8i8(<8 x i8>, <8 x i8>) #1
declare <8 x i8> @llvm.umin.v8i8(<8 x i8>, <8 x i8>) #1
declare i8 @llvm.vector.reduce.add.v8i8(<8 x i8>) #1