Skip to content

[ARM:Perf] Add fast path for SME2 FP16 Pack with small channel count #4578

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lanling-47 wants to merge 1 commit into
base: master
Choose a base branch
from
+33 −15
Open

[ARM:Perf] Add fast path for SME2 FP16 Pack with small channel count #4578

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 33 additions & 15 deletions source/backend/arm82/Arm82Functions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1226,23 +1226,41 @@ static void Sme2MNNPackC4ForMatMul_A_FP16(float* destOrigin, float const** sourc
const int eHandled = eMain * eTile;
const int lHandled = lMain * lTile;

// Process remaining rows
for (int y = eHandled; y < e; ++y) {
int yR = y % eDest;
for (int x = 0; x < l; ++x) {
int xR = x % pack;
int xC = x / pack;
destBase[(x / lP) * dstColBlockStride + yR * lP + (x % lP)] = sourceBase[xC * srcColBlockStride + y * srcRowStride + xR];
if (lMain == 0 && l <= pack) {
// Fast path: all channels fit in one NC8HW8 slice (e.g., ic=3 first conv).
// Avoids per-element division/modulo by using simplified addressing.
auto dest16 = (FLOAT16*)destBase;
for (int y = 0; y < e; ++y) {
auto src = sourceBase + y * srcRowStride;
int x = 0;
for (; x + 1 < l; x += 2) {
auto dstAddr = dest16 + (x / lP) * dstColBlockStride + y * lP;
*(uint32_t*)(dstAddr) = *(const uint32_t*)(src + x);
}
if (x < l) {
auto dstAddr = dest16 + (x / lP) * dstColBlockStride + y * lP + (x % lP);
*dstAddr = src[x];
}
}
} else {
// Process remaining rows
for (int y = eHandled; y < e; ++y) {
int yR = y % eDest;
for (int x = 0; x < l; ++x) {
int xR = x % pack;
int xC = x / pack;
destBase[(x / lP) * dstColBlockStride + yR * lP + (x % lP)] = sourceBase[xC * srcColBlockStride + y * srcRowStride + xR];
}
}
}

// Process remaining columns for the already handled rows
for (int y = 0; y < eHandled; ++y) {
int yR = y % eDest;
for (int x = lHandled; x < l; ++x) {
int xR = x % pack;
int xC = x / pack;
destBase[(x / lP) * dstColBlockStride + yR * lP + (x % lP)] = sourceBase[xC * srcColBlockStride + y * srcRowStride + xR];
// Process remaining columns for the already handled rows
for (int y = 0; y < eHandled; ++y) {
int yR = y % eDest;
for (int x = lHandled; x < l; ++x) {
int xR = x % pack;
int xC = x / pack;
destBase[(x / lP) * dstColBlockStride + yR * lP + (x % lP)] = sourceBase[xC * srcColBlockStride + y * srcRowStride + xR];
}
}
}
}
Expand Down