121 lines
2.9 KiB
ArmAsm
121 lines
2.9 KiB
ArmAsm
/*++
|
|
|
|
Copyright (c) Microsoft Corporation. All rights reserved.
|
|
|
|
Licensed under the MIT License.
|
|
|
|
Module Name:
|
|
|
|
SgemmTransposePackB16x4Avx.s
|
|
|
|
Abstract:
|
|
|
|
This module implements routines for packing buffers for the single precision
|
|
matrix/matrix multiply operation (SGEMM).
|
|
|
|
This implementation uses AVX instructions.
|
|
|
|
--*/
|
|
|
|
#include "asmmacro.h"
|
|
|
|
.intel_syntax noprefix
|
|
|
|
.text
|
|
|
|
/*++
|
|
|
|
Macro Description:
|
|
|
|
4 columns of 8 rows from the source matrix are transposed to 8 columns of 4
|
|
rows in the destination packed buffer.
|
|
|
|
Arguments:
|
|
|
|
StoreOffset - Supplies the relative byte offset into the destination packed
|
|
buffer.
|
|
|
|
Implicit Arguments:
|
|
|
|
rdi - Supplies the address of the destination packed buffer.
|
|
|
|
rsi - Supplies the address of the source matrix.
|
|
|
|
rdx - Supplies the number of elements per row of the source matrix.
|
|
|
|
--*/
|
|
|
|
.macro TransposePackB8x4BlockAvx StoreOffset
|
|
|
|
//
|
|
// Load 4 columns from 8 rows of the source matrix into the lower and upper
|
|
// halves of 4 YMM registers.
|
|
//
|
|
|
|
lea rax,[rsi+rdx*2]
|
|
vmovups xmm0,XMMWORD PTR [rsi]
|
|
vmovups xmm1,XMMWORD PTR [rsi+rdx]
|
|
lea rsi,[rax+rdx*2]
|
|
vmovups xmm2,XMMWORD PTR [rax]
|
|
vmovups xmm3,XMMWORD PTR [rax+rdx]
|
|
lea rax,[rsi+rdx*2]
|
|
vinsertf128 ymm0,ymm0,XMMWORD PTR [rsi],1
|
|
vinsertf128 ymm1,ymm1,XMMWORD PTR [rsi+rdx],1
|
|
vinsertf128 ymm2,ymm2,XMMWORD PTR [rax],1
|
|
vinsertf128 ymm3,ymm3,XMMWORD PTR [rax+rdx],1
|
|
|
|
//
|
|
// Transpose the lower and upper halves of the 4 YMM registers as two 4x4
|
|
// matrices and store the output to the destination packed buffer.
|
|
//
|
|
|
|
vunpcklps ymm4,ymm0,ymm1
|
|
vunpckhps ymm5,ymm0,ymm1
|
|
vunpcklps ymm0,ymm2,ymm3
|
|
vunpckhps ymm1,ymm2,ymm3
|
|
vunpcklpd ymm2,ymm4,ymm0
|
|
vunpckhpd ymm3,ymm4,ymm0
|
|
vmovaps YMMWORD PTR [rdi+16*4*0+\StoreOffset\()],ymm2
|
|
vmovaps YMMWORD PTR [rdi+16*4*1+\StoreOffset\()],ymm3
|
|
vunpcklpd ymm0,ymm5,ymm1
|
|
vunpckhpd ymm4,ymm5,ymm1
|
|
vmovaps YMMWORD PTR [rdi+16*4*2+\StoreOffset\()],ymm0
|
|
vmovaps YMMWORD PTR [rdi+16*4*3+\StoreOffset\()],ymm4
|
|
|
|
.endm
|
|
|
|
/*++
|
|
|
|
Routine Description:
|
|
|
|
This routine transposes elements from the source matrix to the destination
|
|
packed buffer.
|
|
|
|
4 columns of 16 rows from the source matrix are transposed to 16 columns of 4
|
|
rows in the destination packed buffer.
|
|
|
|
Arguments:
|
|
|
|
D (rdi) - Supplies the address of the destination packed buffer.
|
|
|
|
B (rsi) - Supplies the address of the source matrix.
|
|
|
|
ldb (rdx) - Supplies the number of elements per row of the source matrix.
|
|
|
|
Return Value:
|
|
|
|
None.
|
|
|
|
--*/
|
|
|
|
FUNCTION_ENTRY MlasSgemmTransposePackB16x4Avx
|
|
|
|
shl rdx,2 # convert ldb to bytes
|
|
TransposePackB8x4BlockAvx 0*4
|
|
lea rsi,[rax+rdx*2]
|
|
TransposePackB8x4BlockAvx 8*4
|
|
vzeroupper
|
|
ret
|
|
|
|
.end
|