How to Lower to VNNI Instructions?

[FabianSchuetze]

I am trying to use Intels VNNI instructions. Eventually, I would like to use them for convolutions but start with a simple example based on PR #5807.

The following code compiles but does not produce the desired assembly.

class Reduction : public Halide::Generator<Reduction> {
public:
    Input<Buffer<uint8_t, 2>> input{"input"};
    Input<Buffer<int8_t, 1>> filter{"filter"};
    Output<Buffer<int32_t, 2>> out{"out"};
    void generate() {
        Var x("x"), y("y");

        RDom r(0, 4);

        out(x, y) = 0;

        out(x, y) += ConciseCasts::i32(ConciseCasts::i8(filter(r)) * ConciseCasts::u8(input(x + r, y)));
        // out(x, y) = Halide::saturating_sum(ConciseCasts::i32(0), ConciseCasts::i16(ConciseCasts::i8(filter(r)) * ConciseCasts::u8(input(x + r, y))));

        bool override_associativity_test = true;
        int vector_width = 4;
        Var xo, xi;
        out.update()
            .split(x, xo, xi, vector_width)
            .atomic(override_associativity_test)
            .vectorize(r)
            .vectorize(xi);
    }
};

I compile it with the target option target=host-avx512_sapphirerapids. Unfortunately, the code doesn't seem to be using VNNI instructions. The generated IR is:

for (out.s1.x.v1, 0, t40){
let t19 = ((out.s1.x.v1 * 4) + t44)
out[ramp(t19, 1, 4)] = ((int32x4)vector_reduce_add(int32x16((shuffle(filter[ramp(t42, 1, 4)], 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3) * int8x16(shuffle(input[ramp(((out.s1.x.v1 * 4) + t45), 1, 7)], 0, 1, 2, 3, 1, 2, 3, 4, 2, 3, 4, 5, 3, 4, 5, 6))))) + out[ramp(t19, 1, 4)])
}
}

And the relevant assembly:

movl	(%rdx,%rdi,4), %r8d
	vmovd	%r8d, %xmm4
	vpinsrw	$2, 4(%rdx,%rdi,4), %xmm4, %xmm4
	movl	%r8d, %r9d
	shrl	$8, %r9d
	vpinsrb	$1, %r9d, %xmm4, %xmm5
	movl	%r8d, %r9d
	shrl	$16, %r9d
	vpinsrb	$2, %r9d, %xmm5, %xmm5
	shrl	$24, %r8d
	vpinsrb	$3, %r8d, %xmm5, %xmm5
	vpblendw	$4, %xmm4, %xmm5, %xmm4         # xmm4 = xmm5[0,1],xmm4[2],xmm5[3,4,5,6,7]
	vpinsrb	$6, 6(%rdx,%rdi,4), %xmm4, %xmm4
	vpshufb	%xmm0, %xmm4, %xmm4
	vpmovzxbw	%xmm4, %ymm4            # ymm4 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero,xmm4[8],zero,xmm4[9],zero,xmm4[10],zero,xmm4[11],zero,xmm4[12],zero,xmm4[13],zero,xmm4[14],zero,xmm4[15],zero
	vpmullw	%ymm1, %ymm4, %ymm4
	vpmovwb	%ymm4, %xmm4
	vpmaddubsw	%xmm4, %xmm2, %xmm4
	vmovdqu	(%rsi), %xmm5
	{vex}	vpdpwssd	%xmm3, %xmm4, %xmm5
	vmovdqu	%xmm5, (%rsi)
	incq	%rdi
	addq	$16, %rsi
	cmpq	%rdi, %r15
	jne	.LBB250_32

Instead, when I try to follow the example of the PR the compilation aborts. The code:

out(x, y) = Halide::saturating_sum(ConciseCasts::i32(0), ConciseCasts::i16(ConciseCasts::i8(filter(r)) * ConciseCasts::u8(input(x + r, y))));

leads to a termination of the compilation because saturating_sum cannot matched.

I know that there's also the unit-test in simd_op_check_x86:

RDom r(0, 4);
check("vpdpbusds*zmm", 16, saturating_sum(i32(in_u8(4 * x + r)) * in_i8(4 * x + r + 32)));
check("vpdpbusds*zmm", 16, saturating_sum(i32(in_i8(4 * x + r)) * in_u8(4 * x + r + 32)));

But I unfortunately can't figure out to which generator and schedule this maps to.

What should I do to produce VNNI instructions? And, are there any suggestions for generating VNNI instructions for a convolution?