portelli/Grid
1
0 Fork 0
mirror of https://github.com/paboyle/Grid.git synced 2026-10-24 15:53:18 +01:00
Code Releases Activity

Optimisation of A64FX kernels: FMLA instead of FCMLA, and better exterior loop

Browse Source
This commit is contained in:
portelli committed 2026-10-21 11:49:34 +09:00
1 parent 3232b438f2
commit e930240878
4 files changed
+203 -25
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Grid/qcd/action/fermion/implementation/WilsonKernelsAsmBodyA64FX.h
+8 -15
View File
@@ -95,7 +95,6 @@ Author: Nils Meyer <nils.meyer@ur.de> Regensburg University
#ifdef INTERIOR_AND_EXTERIOR
#define ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON) \
basep = st.GetPFInfo(nent,plocal); nent++; \
if ( local ) { \
LOAD_CHIMU(base); \
LOAD_TABLE(PERMUTE_DIR); \
@@ -106,23 +105,15 @@ Author: Nils Meyer <nils.meyer@ur.de> Regensburg University
} \
base = st.GetInfo(ptype,local,perm,NxtDir,ent,plocal); ent++; \
MULT_2SPIN_1(Dir); \
PREFETCH_CHIMU(base); \
PREFETCH_CHIMU_L2(basep); \
/* PREFETCH_GAUGE_L1(NxtDir); */ \
MULT_2SPIN_2; \
if (s == 0) { \
if ((Dir == 0) || (Dir == 4)) { PREFETCH_GAUGE_L2(Dir); } \
} \
RECON; \
/*
NB: picking PREFETCH_GAUGE_L2(Dir+4); here results in performance penalty
though I expected that it would improve on performance
No software prefetch: the hardware prefetcher handles the neighbour and gauge link streams better on A64FX.
*/
#define ASM_LEG_XP(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON) \
base = st.GetInfo(ptype,local,perm,Dir,ent,plocal); ent++; \
PREFETCH1_CHIMU(base); \
ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)
#define RESULT(base,basep) SAVE_RESULT(base,basep);
@@ -130,12 +121,11 @@ NB: picking PREFETCH_GAUGE_L2(Dir+4); here results in performance penalty
#endif
////////////////////////////////////////////////////////////////////////////////
// Pre comms kernel -- prefetch like normal because it is mostly right
// Pre comms kernel -- gauge link prefetch into L2 only
////////////////////////////////////////////////////////////////////////////////
#ifdef INTERIOR
#define ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON) \
basep = st.GetPFInfo(nent,plocal); nent++; \
if ( local ) { \
LOAD_CHIMU(base); \
LOAD_TABLE(PERMUTE_DIR); \
@@ -148,12 +138,15 @@ NB: picking PREFETCH_GAUGE_L2(Dir+4); here results in performance penalty
RECON; \
} \
base = st.GetInfo(ptype,local,perm,NxtDir,ent,plocal); ent++; \
PREFETCH_CHIMU(base); \
PREFETCH_CHIMU_L2(basep); \
if (s == 0) { \
if ((Dir == 0) || (Dir == 4)) { PREFETCH_GAUGE_L2(Dir); } \
} \
/* No software prefetch of the neighbours (the hardware prefetcher handles them better on A64FX);
the gauge links of the next site are prefetched into L2. */
#define ASM_LEG_XP(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON) \
base = st.GetInfo(ptype,local,perm,Dir,ent,plocal); ent++; \
PREFETCH1_CHIMU(base); \
{ ZERO_PSI; } \
ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)
Grid/qcd/action/fermion/implementation/WilsonKernelsImplementation.h
+17 -2
View File
@@ -479,6 +479,21 @@ extern "C" {
int sF = ss*Ls; \
WilsonKernels<Impl>::A(st_v,U_v,buf,sF,sU,Ls,1,in_v,out_v); \
});
#ifdef A64FX
// Exterior kernel on the surface 4d sites only (the surface list holds Ls
// consecutive 5d entries per surface 4d site)
#define ASM_CALL_EXT(A) \
const uint64_t nsurf = st.surface_list.size()/Ls; \
auto surf = &st.surface_list[0]; \
thread_for( sss, nsurf, { \
int sU = surf[sss*Ls]/Ls; \
int sF = sU*Ls; \
WilsonKernels<Impl>::A(st_v,U_v,buf,sF,sU,Ls,1,in_v,out_v); \
});
#else
#define ASM_CALL_EXT(A) ASM_CALL(A)
#endif
#define ASM_CALL_SLICE(A) \
auto grid = in.Grid() ; \
int nt = grid->LocalDimensions()[4]; \
@@ -522,7 +537,7 @@ void WilsonKernels<Impl>::DhopKernel(int Opt,StencilImpl &st, DoubledGaugeField
if (Opt == WilsonKernelsStatic::OptGeneric ) { KERNEL_CALL_EXT(GenericDhopSiteExt); return;}
if (Opt == WilsonKernelsStatic::OptHandUnroll ) { KERNEL_CALL_EXT(HandDhopSiteExt); return;}
#ifndef GRID_CUDA
if (Opt == WilsonKernelsStatic::OptInlineAsm ) { ASM_CALL(AsmDhopSiteExt); return;}
if (Opt == WilsonKernelsStatic::OptInlineAsm ) { ASM_CALL_EXT(AsmDhopSiteExt); return;}
#endif
}
GRID_ASSERT(0 && " Kernel optimisation case not covered ");
@@ -568,7 +583,7 @@ void WilsonKernels<Impl>::DhopKernel(int Opt,StencilImpl &st, DoubledGaugeField
if (Opt == WilsonKernelsStatic::OptGeneric ) { KERNEL_CALL_EXT(GenericDhopSiteDagExt); return;}
if (Opt == WilsonKernelsStatic::OptHandUnroll ) { KERNEL_CALL_EXT(HandDhopSiteDagExt); return;}
#ifndef GRID_CUDA
if (Opt == WilsonKernelsStatic::OptInlineAsm ) { ASM_CALL(AsmDhopSiteDagExt); return;}
if (Opt == WilsonKernelsStatic::OptInlineAsm ) { ASM_CALL_EXT(AsmDhopSiteDagExt); return;}
#endif
}
GRID_ASSERT(0 && " Kernel optimisation case not covered ");
Grid/simd/Fujitsu_A64FX_intrin_double.h
+89 -4
View File
@@ -39,8 +39,13 @@ Author: Nils Meyer <nils.meyer@ur.de>
#define UNLOCK_GAUGE(A)
#define MASK_REGS DECLARATIONS_A64FXd
#define SAVE_RESULT(A,B) RESULT_A64FXd(A);
#define MULT_2SPIN_1(Dir) MULT_2SPIN_1_A64FXd(Dir)
#define MULT_2SPIN_2 MULT_2SPIN_2_A64FXd
#ifdef GRID_A64FX_FCMLA
#define MULT_2SPIN_1(Dir) MULT_2SPIN_1_FCMLA_A64FXd(Dir)
#define MULT_2SPIN_2 MULT_2SPIN_2_FCMLA_A64FXd
#else
#define MULT_2SPIN_1(Dir) MULT_2SPIN_1_FMLA_A64FXd(Dir)
#define MULT_2SPIN_2 MULT_2SPIN_2_FMLA_A64FXd
#endif
#define LOAD_CHI(base) LOAD_CHI_A64FXd(base)
#define ZERO_PSI ZERO_PSI_A64FXd
#define ADD_RESULT(base,basep) LOAD_CHIMU(base); ADD_RESULT_INTERNAL_A64FXd; RESULT_A64FXd(base)
@@ -273,8 +278,88 @@ Author: Nils Meyer <nils.meyer@ur.de>
U_11 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-2)); \
U_21 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(1)); \
}
// MULT_2SPIN with FMLA instead of FCMLA (3 uops on A64FX):
// UChi = sum_b Re(U_ab) Chi_b + Im(U_ab) (i Chi_b)
// column by column. Re(U)/Im(U) are duplicated over each complex with trn1/trn2,
// i Chi_b is formed in place with fcadd once Chi_b has been used (Chi is dead after
// the multiply). U_00, U_10, U_20 hold column b, U_01 is a temporary.
// Define GRID_A64FX_FCMLA to use the original FCMLA version.
#define MULT_2SPIN_1_FMLA_A64FXd(A) \
{ \
const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-6)); \
U_10 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-3)); \
U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(0)); \
U_01 = svtrn1(U_00, U_00); \
UChi_00 = svmul_x(pg1, U_01, Chi_00); \
UChi_10 = svmul_x(pg1, U_01, Chi_10); \
U_01 = svtrn1(U_10, U_10); \
UChi_01 = svmul_x(pg1, U_01, Chi_00); \
UChi_11 = svmul_x(pg1, U_01, Chi_10); \
U_01 = svtrn1(U_20, U_20); \
UChi_02 = svmul_x(pg1, U_01, Chi_00); \
UChi_12 = svmul_x(pg1, U_01, Chi_10); \
Chi_00 = svcadd_x(pg1, zero0, Chi_00, 90); \
Chi_10 = svcadd_x(pg1, zero0, Chi_10, 90); \
U_01 = svtrn2(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_00); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_10); \
U_01 = svtrn2(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_00); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_10); \
U_01 = svtrn2(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_00); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_10); \
U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-5)); \
U_10 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-2)); \
U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(1)); \
U_01 = svtrn1(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
U_01 = svtrn1(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
U_01 = svtrn1(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
Chi_01 = svcadd_x(pg1, zero0, Chi_01, 90); \
Chi_11 = svcadd_x(pg1, zero0, Chi_11, 90); \
U_01 = svtrn2(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
U_01 = svtrn2(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
U_01 = svtrn2(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-4)); \
U_10 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-1)); \
U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(2)); \
U_01 = svtrn1(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
U_01 = svtrn1(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
U_01 = svtrn1(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
Chi_02 = svcadd_x(pg1, zero0, Chi_02, 90); \
Chi_12 = svcadd_x(pg1, zero0, Chi_12, 90); \
U_01 = svtrn2(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
U_01 = svtrn2(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
U_01 = svtrn2(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
}
#define MULT_2SPIN_2_FMLA_A64FXd
// MULT_2SPIN
#define MULT_2SPIN_1_A64FXd(A) \
#define MULT_2SPIN_1_FCMLA_A64FXd(A) \
{ \
const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-6)); \
@@ -300,7 +385,7 @@ Author: Nils Meyer <nils.meyer@ur.de>
U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(2)); \
}
// MULT_2SPIN_BACKEND
#define MULT_2SPIN_2_A64FXd \
#define MULT_2SPIN_2_FCMLA_A64FXd \
{ \
UChi_00 = svcmla_x(pg1, UChi_00, U_01, Chi_01, 0); \
UChi_10 = svcmla_x(pg1, UChi_10, U_01, Chi_11, 0); \
Grid/simd/Fujitsu_A64FX_intrin_single.h
+89 -4
View File
@@ -39,8 +39,13 @@ Author: Nils Meyer <nils.meyer@ur.de>
#define UNLOCK_GAUGE(A)
#define MASK_REGS DECLARATIONS_A64FXf
#define SAVE_RESULT(A,B) RESULT_A64FXf(A);
#define MULT_2SPIN_1(Dir) MULT_2SPIN_1_A64FXf(Dir)
#define MULT_2SPIN_2 MULT_2SPIN_2_A64FXf
#ifdef GRID_A64FX_FCMLA
#define MULT_2SPIN_1(Dir) MULT_2SPIN_1_FCMLA_A64FXf(Dir)
#define MULT_2SPIN_2 MULT_2SPIN_2_FCMLA_A64FXf
#else
#define MULT_2SPIN_1(Dir) MULT_2SPIN_1_FMLA_A64FXf(Dir)
#define MULT_2SPIN_2 MULT_2SPIN_2_FMLA_A64FXf
#endif
#define LOAD_CHI(base) LOAD_CHI_A64FXf(base)
#define ZERO_PSI ZERO_PSI_A64FXf
#define ADD_RESULT(base,basep) LOAD_CHIMU(base); ADD_RESULT_INTERNAL_A64FXf; RESULT_A64FXf(base)
@@ -273,8 +278,88 @@ Author: Nils Meyer <nils.meyer@ur.de>
U_11 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-2)); \
U_21 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(1)); \
}
// MULT_2SPIN with FMLA instead of FCMLA (3 uops on A64FX):
// UChi = sum_b Re(U_ab) Chi_b + Im(U_ab) (i Chi_b)
// column by column. Re(U)/Im(U) are duplicated over each complex with trn1/trn2,
// i Chi_b is formed in place with fcadd once Chi_b has been used (Chi is dead after
// the multiply). U_00, U_10, U_20 hold column b, U_01 is a temporary.
// Define GRID_A64FX_FCMLA to use the original FCMLA version.
#define MULT_2SPIN_1_FMLA_A64FXf(A) \
{ \
const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-6)); \
U_10 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-3)); \
U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(0)); \
U_01 = svtrn1(U_00, U_00); \
UChi_00 = svmul_x(pg1, U_01, Chi_00); \
UChi_10 = svmul_x(pg1, U_01, Chi_10); \
U_01 = svtrn1(U_10, U_10); \
UChi_01 = svmul_x(pg1, U_01, Chi_00); \
UChi_11 = svmul_x(pg1, U_01, Chi_10); \
U_01 = svtrn1(U_20, U_20); \
UChi_02 = svmul_x(pg1, U_01, Chi_00); \
UChi_12 = svmul_x(pg1, U_01, Chi_10); \
Chi_00 = svcadd_x(pg1, zero0, Chi_00, 90); \
Chi_10 = svcadd_x(pg1, zero0, Chi_10, 90); \
U_01 = svtrn2(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_00); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_10); \
U_01 = svtrn2(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_00); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_10); \
U_01 = svtrn2(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_00); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_10); \
U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-5)); \
U_10 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-2)); \
U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(1)); \
U_01 = svtrn1(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
U_01 = svtrn1(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
U_01 = svtrn1(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
Chi_01 = svcadd_x(pg1, zero0, Chi_01, 90); \
Chi_11 = svcadd_x(pg1, zero0, Chi_11, 90); \
U_01 = svtrn2(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
U_01 = svtrn2(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
U_01 = svtrn2(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-4)); \
U_10 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-1)); \
U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(2)); \
U_01 = svtrn1(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
U_01 = svtrn1(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
U_01 = svtrn1(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
Chi_02 = svcadd_x(pg1, zero0, Chi_02, 90); \
Chi_12 = svcadd_x(pg1, zero0, Chi_12, 90); \
U_01 = svtrn2(U_00, U_00); \
UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
U_01 = svtrn2(U_10, U_10); \
UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
U_01 = svtrn2(U_20, U_20); \
UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
}
#define MULT_2SPIN_2_FMLA_A64FXf
// MULT_2SPIN
#define MULT_2SPIN_1_A64FXf(A) \
#define MULT_2SPIN_1_FCMLA_A64FXf(A) \
{ \
const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-6)); \
@@ -300,7 +385,7 @@ Author: Nils Meyer <nils.meyer@ur.de>
U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(2)); \
}
// MULT_2SPIN_BACKEND
#define MULT_2SPIN_2_A64FXf \
#define MULT_2SPIN_2_FCMLA_A64FXf \
{ \
UChi_00 = svcmla_x(pg1, UChi_00, U_01, Chi_01, 0); \
UChi_10 = svcmla_x(pg1, UChi_10, U_01, Chi_11, 0); \