;++ ; ; Copyright (c) Microsoft Corporation. All rights reserved. ; ; Licensed under the MIT License. ; ; Module Name: ; ; SconvKernelAvx512F.asm ; ; Abstract: ; ; This module implements the kernels for the single precision convolution ; operation. ; ; This implementation uses AVX512F instructions. ; ;-- .xlist INCLUDE mlasi.inc INCLUDE SconvKernelCommon.inc .list ; ; Macro Description: ; ; This macro generates code to clear the block accumulators. ; ; Arguments: ; ; FilterCount - Supplies the number of rows from the filter to process. ; ; OutputCount - Supplies the number of output blocks to produce. ; ; Implicit Arguments: ; ; zmm0-zmm23 - Supplies the block accumulators. ; ClearBlock MACRO FilterCount, OutputCount EmitIfCount2GE FilterCount, 1, OutputCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 2, EmitIfCount2GE FilterCount, 1, OutputCount, 3, EmitIfCount2GE FilterCount, 1, OutputCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 5, EmitIfCount2GE FilterCount, 1, OutputCount, 6, EmitIfCount2GE FilterCount, 2, OutputCount, 1, EmitIfCount2GE FilterCount, 2, OutputCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCount2GE FilterCount, 3, OutputCount, 1, EmitIfCount2GE FilterCount, 3, OutputCount, 2, EmitIfCount2GE FilterCount, 3, OutputCount, 3, EmitIfCount2GE FilterCount, 3, OutputCount, 4, EmitIfCount2GE FilterCount, 3, OutputCount, 5, EmitIfCount2GE FilterCount, 3, OutputCount, 6, EmitIfCount2GE FilterCount, 4, OutputCount, 1, EmitIfCount2GE FilterCount, 4, OutputCount, 2, EmitIfCount2GE FilterCount, 4, OutputCount, 3, EmitIfCount2GE FilterCount, 4, OutputCount, 4, EmitIfCount2GE FilterCount, 4, OutputCount, 5, EmitIfCount2GE FilterCount, 4, OutputCount, 6, ENDM ; ; Macro Description: ; ; This macro multiplies and accumulates for FilterCount by OutputCount block ; of the output buffer. ; ; Arguments: ; ; KernelType - Supplies the type of kernel to be generated. ; ; FilterCount - Supplies the number of rows from the filter to process. ; ; OutputCount - Supplies the number of output blocks to produce. ; ; VectorOffset - Supplies the byte offset from the filter buffer to fetch ; elements. ; ; BroadcastOffset - Supplies the byte offset from the input buffer to fetch ; elements. ; ; Implicit Arguments: ; ; rcx - Supplies the address of the input buffer. ; ; rdx - Supplies the address of the filter buffer. ; ; rsi - Supplies the FilterStride parameter (see function description). ; ; rbx - Supplies the address of the filter buffer plus 2 * FilterStride. ; ; r9 - Supplies the StrideWidth parameter (see function description). ; ; r14 - Supplies the address of the input buffer plus 3 * StrideWidth. ; ; zmm0-zmm23 - Supplies the block accumulators. ; ComputeBlock MACRO KernelType, FilterCount, OutputCount, VectorOffset, BroadcastOffset IFIDNI , vmovups zmm24,ZMMWORD PTR [rdx+VectorOffset] EmitIfCountGE OutputCount, 1, EmitIfCountGE OutputCount, 2, EmitIfCountGE OutputCount, 3, EmitIfCountGE OutputCount, 4, EmitIfCountGE OutputCount, 5, EmitIfCountGE OutputCount, 6, ELSE IF FilterCount EQ 1 vmovups zmm24,ZMMWORD PTR [rdx+VectorOffset] EmitIfCountGE OutputCount, 1, EmitIfCountGE OutputCount, 2, EmitIfCountGE OutputCount, 3, EmitIfCountGE OutputCount, 4, EmitIfCountGE OutputCount, 5, EmitIfCountGE OutputCount, 6, ELSE EmitIfCountGE OutputCount, 1, EmitIfCountGE OutputCount, 2, EmitIfCountGE OutputCount, 3, EmitIfCountGE OutputCount, 4, EmitIfCountGE OutputCount, 5, EmitIfCountGE OutputCount, 6, IF OutputCount EQ 1 EmitIfCountGE FilterCount, 1, EmitIfCountGE FilterCount, 2, EmitIfCountGE FilterCount, 3, EmitIfCountGE FilterCount, 4, ELSE EmitIfCountGE FilterCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 2, EmitIfCount2GE FilterCount, 1, OutputCount, 3, EmitIfCount2GE FilterCount, 1, OutputCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 5, EmitIfCount2GE FilterCount, 1, OutputCount, 6, EmitIfCountGE FilterCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 1, EmitIfCount2GE FilterCount, 2, OutputCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCountGE FilterCount, 3, EmitIfCount2GE FilterCount, 3, OutputCount, 1, EmitIfCount2GE FilterCount, 3, OutputCount, 2, EmitIfCount2GE FilterCount, 3, OutputCount, 3, EmitIfCount2GE FilterCount, 3, OutputCount, 4, EmitIfCount2GE FilterCount, 3, OutputCount, 5, EmitIfCount2GE FilterCount, 3, OutputCount, 6, EmitIfCountGE FilterCount, 4, EmitIfCount2GE FilterCount, 4, OutputCount, 1, EmitIfCount2GE FilterCount, 4, OutputCount, 2, EmitIfCount2GE FilterCount, 4, OutputCount, 3, EmitIfCount2GE FilterCount, 4, OutputCount, 4, EmitIfCount2GE FilterCount, 4, OutputCount, 5, EmitIfCount2GE FilterCount, 4, OutputCount, 6, ENDIF ENDIF ENDIF ENDM ; ; Macro Description: ; ; This macro generates code to compute the convolution for a specified number ; of filter rows. ; ; Arguments: ; ; KernelFrame - Supplies the symbol name to access the convolution kernel ; stack. ; ; KernelType - Supplies the type of kernel to be generated. ; ; FilterCount - Supplies the number of rows from the filter to process. ; ; Implicit Arguments: ; ; rdi - Supplies the address of the input buffer. ; ; rsi - Supplies the FilterStride parameter (see function description) when ; KernelType!=Depthwise. Supplies the address of the filter buffer when ; KernelType=Depthwise. ; ; rbp - Supplies the DilationWidth parameter (see function description). ; ; r8 - Supplies the address of the output buffer. ; ; r9 - Supplies the StrideWidth parameter (see function description). ; ; r15 - Supplies the InputStride parameter (see function description). ; ProcessFilterCountN MACRO KernelFrame, KernelType, FilterCount LOCAL ProcessOutputCount LOCAL ProcessNextOutputCountBy6 LOCAL ProcessRemainingOutputCount LOCAL ProcessRemainingOutputCountLessThan3 LOCAL ProcessRemainingOutputCount1 LOCAL ProcessOutputCountRightPadAndRemaining ; ; Process the output blocks that include left padding. ; mov r10,KernelFrame.OutputCountLeftPad[rsp] test r10,r10 jz ProcessOutputCount call MlasConv&KernelType&FloatSingleAvx512FFilter&FilterCount ; ; Process the output blocks that do not include any padding. ; ProcessOutputCount: mov r10,KernelFrame.OutputCount[rsp] sub r10,6 jb ProcessRemainingOutputCount ProcessNextOutputCountBy6: ProcessOutputCountN Avx512F, KernelFrame, KernelType, 16, FilterCount, 6 lea rax,[r9*2+r9] lea rdi,[rdi+rax*2] ; advance input by 6 elements sub r10,6 jae ProcessNextOutputCountBy6 ProcessRemainingOutputCount: add r10,6 ; correct for over-subtract above jz ProcessOutputCountRightPadAndRemaining cmp r10,3 jb ProcessRemainingOutputCountLessThan3 ProcessOutputCountN Avx512F, KernelFrame, KernelType, 16, FilterCount, 3 lea rax,[r9*2+r9] add rdi,rax ; advance input by 3 elements sub r10,3 jz ProcessOutputCountRightPadAndRemaining ProcessRemainingOutputCountLessThan3: cmp r10,1 je ProcessOutputCountRightPadAndRemaining ProcessOutputCountN Avx512F, KernelFrame, KernelType, 16, FilterCount, 2 lea rdi,[rdi+r9*2] ; advance input by 2 elements sub r10,2 ; ; Process the output blocks that include right padding plus any remaining output ; blocks from above. ; ProcessOutputCountRightPadAndRemaining: add r10,KernelFrame.OutputCountRightPad[rsp] jz ExitKernel call MlasConv&KernelType&FloatSingleAvx512FFilter&FilterCount ENDM ; ; Macro Description: ; ; This macro generates code to compute the convolution for a specified number ; of filter rows for a pointwise convolution. ; ; Arguments: ; ; FilterCount - Supplies the number of rows from the filter to process. ; ; Implicit Arguments: ; ; rdi - Supplies the address of the input buffer. ; ; rsi - Supplies the FilterStride parameter (see function description). ; ; rbp - Supplies the InputStride parameter (see function description). ; ; r8 - Supplies the address of the output buffer. ; ; r9 - Supplies the StrideWidth parameter (see function description). ; ; r10 - Supplies the OutputCount parameter (see function description). ; ; r12 - Supplies the address of the filter buffer. ; ProcessPointwiseFilterCountN MACRO FilterCount LOCAL ProcessNextOutputCountBy6 LOCAL ProcessRemainingOutputCount LOCAL ProcessRemainingOutputCountLessThan3 LOCAL ProcessRemainingOutputCount1 sub r10,6 jb ProcessRemainingOutputCount ProcessNextOutputCountBy6: ProcessPointwiseOutputCountN Avx512F, 16, FilterCount, 6 lea rax,[r9*2+r9] lea rdi,[rdi+rax*2] ; advance input by 6 elements sub r10,6 jae ProcessNextOutputCountBy6 ProcessRemainingOutputCount: add r10,6 ; correct for over-subtract above jz ExitKernel cmp r10,3 jb ProcessRemainingOutputCountLessThan3 ProcessPointwiseOutputCountN Avx512F, 16, FilterCount, 3 lea rax,[r9*2+r9] add rdi,rax ; advance input by 3 elements sub r10,3 jz ExitKernel ProcessRemainingOutputCountLessThan3: cmp r10,2 jb ProcessRemainingOutputCount1 ProcessPointwiseOutputCountN Avx512F, 16, FilterCount, 2 jmp ExitKernel ProcessRemainingOutputCount1: ProcessPointwiseOutputCountN Avx512F, 16, FilterCount, 1 ENDM ; ; Generate the convolution kernels. ; ; N.B. BiasFilter is not used here as the AVX-512 EVEX instruction encoding ; efficiently compresses aligned relative byte offsets. ; SconvKernelFunction Nchw, 16, Avx512F SconvKernelFunction Nchwc, 16, Avx512F SconvKernelDepthwiseFunction 16, Avx512F SconvKernelPointwiseFunction Avx512F ; ; Macro Description: ; ; This macro generates code to process an output block after the inner ; convolution kernel has executed and then stores the output block to the ; output buffer. ; ; Arguments: ; ; FilterCount - Supplies the number of rows from the filter to process. ; ; OutputCount - Supplies the number of output blocks to produce. ; IRP FilterCount, <1, 2, 3, 4> IRP OutputCount, <1, 2, 3, 6> LEAF_ENTRY MlasConvPostProcessFloatAvx512FFilter&FilterCount&Output&OutputCount, _TEXT IF FilterCount GT 2 lea rbx,[r8+rax*2] ; compute output plus 2 rows ENDIF ; ; Test if the existing contents of the output buffer should be accumulated ; with the output block. ; test dl,MLAS_CONV_KERNEL_FLAG_ACCUMULATE_OUTPUT jz SkipAccumulateOutput EmitIfCount2GE FilterCount, 1, OutputCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 2, EmitIfCount2GE FilterCount, 1, OutputCount, 3, EmitIfCount2GE FilterCount, 1, OutputCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 5, EmitIfCount2GE FilterCount, 1, OutputCount, 6, EmitIfCount2GE FilterCount, 2, OutputCount, 1, EmitIfCount2GE FilterCount, 2, OutputCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCount2GE FilterCount, 3, OutputCount, 1, EmitIfCount2GE FilterCount, 3, OutputCount, 2, EmitIfCount2GE FilterCount, 3, OutputCount, 3, EmitIfCount2GE FilterCount, 3, OutputCount, 4, EmitIfCount2GE FilterCount, 3, OutputCount, 5, EmitIfCount2GE FilterCount, 3, OutputCount, 6, EmitIfCount2GE FilterCount, 4, OutputCount, 1, EmitIfCount2GE FilterCount, 4, OutputCount, 2, EmitIfCount2GE FilterCount, 4, OutputCount, 3, EmitIfCount2GE FilterCount, 4, OutputCount, 4, EmitIfCount2GE FilterCount, 4, OutputCount, 5, EmitIfCount2GE FilterCount, 4, OutputCount, 6, SkipAccumulateOutput: ; ; Test if the bias buffer should be accumulated with the output block. ; test dl,MLAS_CONV_KERNEL_FLAG_BIAS_ADDITION jz SkipBiasAddition IF OutputCount EQ 1 EmitIfCountGE FilterCount, 1, EmitIfCountGE FilterCount, 2, EmitIfCountGE FilterCount, 3, EmitIfCountGE FilterCount, 4, ELSE EmitIfCountGE FilterCount, 1, EmitIfCountGE FilterCount, 2, EmitIfCountGE FilterCount, 3, EmitIfCountGE FilterCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 2, EmitIfCount2GE FilterCount, 1, OutputCount, 3, EmitIfCount2GE FilterCount, 1, OutputCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 5, EmitIfCount2GE FilterCount, 1, OutputCount, 6, EmitIfCount2GE FilterCount, 2, OutputCount, 1, EmitIfCount2GE FilterCount, 2, OutputCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCount2GE FilterCount, 3, OutputCount, 1, EmitIfCount2GE FilterCount, 3, OutputCount, 2, EmitIfCount2GE FilterCount, 3, OutputCount, 3, EmitIfCount2GE FilterCount, 3, OutputCount, 4, EmitIfCount2GE FilterCount, 3, OutputCount, 5, EmitIfCount2GE FilterCount, 3, OutputCount, 6, EmitIfCount2GE FilterCount, 4, OutputCount, 1, EmitIfCount2GE FilterCount, 4, OutputCount, 2, EmitIfCount2GE FilterCount, 4, OutputCount, 3, EmitIfCount2GE FilterCount, 4, OutputCount, 4, EmitIfCount2GE FilterCount, 4, OutputCount, 5, EmitIfCount2GE FilterCount, 4, OutputCount, 6, ENDIF SkipBiasAddition: ; ; Test for fused ReLU activation. ; test dl,MLAS_CONV_KERNEL_FLAG_RELU_ACTIVATION jz SkipReluActivation vpxord zmm24,zmm24,zmm24 EmitIfCount2GE FilterCount, 1, OutputCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 2, EmitIfCount2GE FilterCount, 1, OutputCount, 3, EmitIfCount2GE FilterCount, 1, OutputCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 5, EmitIfCount2GE FilterCount, 1, OutputCount, 6, EmitIfCount2GE FilterCount, 2, OutputCount, 1, EmitIfCount2GE FilterCount, 2, OutputCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCount2GE FilterCount, 3, OutputCount, 1, EmitIfCount2GE FilterCount, 3, OutputCount, 2, EmitIfCount2GE FilterCount, 3, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCount2GE FilterCount, 4, OutputCount, 1, EmitIfCount2GE FilterCount, 4, OutputCount, 2, EmitIfCount2GE FilterCount, 4, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, SkipReluActivation: ; ; Store the output block in the output buffer. ; EmitIfCount2GE FilterCount, 1, OutputCount, 1, EmitIfCount2GE FilterCount, 1, OutputCount, 2, EmitIfCount2GE FilterCount, 1, OutputCount, 3, EmitIfCount2GE FilterCount, 1, OutputCount, 4, EmitIfCount2GE FilterCount, 1, OutputCount, 5, EmitIfCount2GE FilterCount, 1, OutputCount, 6, EmitIfCount2GE FilterCount, 2, OutputCount, 1, EmitIfCount2GE FilterCount, 2, OutputCount, 2, EmitIfCount2GE FilterCount, 2, OutputCount, 3, EmitIfCount2GE FilterCount, 2, OutputCount, 4, EmitIfCount2GE FilterCount, 2, OutputCount, 5, EmitIfCount2GE FilterCount, 2, OutputCount, 6, EmitIfCount2GE FilterCount, 3, OutputCount, 1, EmitIfCount2GE FilterCount, 3, OutputCount, 2, EmitIfCount2GE FilterCount, 3, OutputCount, 3, EmitIfCount2GE FilterCount, 3, OutputCount, 4, EmitIfCount2GE FilterCount, 3, OutputCount, 5, EmitIfCount2GE FilterCount, 3, OutputCount, 6, EmitIfCount2GE FilterCount, 4, OutputCount, 1, EmitIfCount2GE FilterCount, 4, OutputCount, 2, EmitIfCount2GE FilterCount, 4, OutputCount, 3, EmitIfCount2GE FilterCount, 4, OutputCount, 4, EmitIfCount2GE FilterCount, 4, OutputCount, 5, EmitIfCount2GE FilterCount, 4, OutputCount, 6, add_immed r8,OutputCount*16*4 ; advance output by N nchw16c blocks ret LEAF_END MlasConvPostProcessFloatAvx512FFilter&FilterCount&Output&OutputCount, _TEXT ENDM ENDM END