Commit cee92951 authored by Jing Zhang's avatar Jing Zhang
Browse files

changing layout

parent 92c615b6
...@@ -106,13 +106,50 @@ int main(int argc, char* argv[]) ...@@ -106,13 +106,50 @@ int main(int argc, char* argv[])
int StrideA = K; int StrideA = K;
int StrideB = K; int StrideB = K;
int StrideM0 = S * NumHead * HeadDim; if(i % 2 == 0)
int StrideM1 = HeadDim; {
int StrideN0 = S * HeadDim;
int StrideN1 = 1;
gemm_descs.push_back( int StrideM0 = S * NumHead * HeadDim;
{M, N, K, StrideA, StrideB, M0, M1, N0, N1, StrideM0, StrideM1, StrideN0, StrideN1}); int StrideM1 = 1;
int StrideN0 = S * HeadDim;
int StrideN1 = S;
gemm_descs.push_back({M,
N,
K,
StrideA,
StrideB,
M0,
M1,
N0,
N1,
StrideM0,
StrideM1,
StrideN0,
StrideN1});
}
else
{
int StrideM0 = S * NumHead * HeadDim;
int StrideM1 = HeadDim;
int StrideN0 = S * HeadDim;
int StrideN1 = 1;
gemm_descs.push_back({M,
N,
K,
StrideA,
StrideB,
M0,
M1,
N0,
N1,
StrideM0,
StrideM1,
StrideN0,
StrideN1});
}
} }
auto f_host_tensor_descriptor = auto f_host_tensor_descriptor =
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment