xe2: jit: gemm: use atomic loads for fused post-op readback

Author: petercad

src/gpu/intel/gemm/jit/generator/pieces/atomic_fusions.cxx

@@ -530,8 +530,11 @@ bool Generator<hw>::gemmFusedPostOpsFinalize(Label &labelLateExit, GEMMProblem &
     modState.ra.safeRelease(zero);
 
     status << "Load completed C tile" << status_stream::endl;
-    modStrategy.C.atomic = modStrategy.CO.atomic = false;
-    modState.Cext_strategy.atomic = false;
+    if (hw < HW::Xe2) {
+        /* Xe2 + later need atomic loads */
+        modStrategy.C.atomic = modStrategy.CO.atomic = false;
+        modState.Cext_strategy.atomic = false;
+    }
     gemmAccessC(COperation::Load, modProblem, modStrategy, modState);
 
     if (strategy.zeroTempC) {

src/gpu/intel/gemm/jit/generator/pieces/c_update.cxx

@@ -172,8 +172,6 @@ bool Generator<hw>::gemmAccessC(COperation op, const GEMMProblem &problem, const
     bool stdCRemainder = !(altCRemainder && (strategy.remHandling[LoopM] == RemainderHandling::KnownRemainder)
                                          && (strategy.remHandling[LoopN] == RemainderHandling::KnownRemainder));
 
-    if ((op != COperation::UpdateStore) && strategy.C.atomic) stub();
-
     if (state.allowEmptyC && (remainderM || remainderN)) {
         if (!state.isNested) stub();
         int simt = strategy.fused ? 16 : 1;
@@ -846,6 +844,7 @@ void Generator<hw>::updateCLayout(const RegisterLayout &layoutExt, const GRFRang
 #define FOR_EACH_C for (int q = 0; q < C_count; q++)
     auto Tc = problem.Tc, Tc_ext = problem.Tc_ext, Ts = problem.Ts;
     bool loadOnly = (op == COperation::Load);
+    bool atomicUpdate = (op == COperation::UpdateStore) && strategy.C.atomic;
     bool beta0 = problem.beta0();
     bool needLoad = (!beta0 && !loadOnly);
     bool copyC = state.copyC;
@@ -876,7 +875,7 @@ void Generator<hw>::updateCLayout(const RegisterLayout &layoutExt, const GRFRang
     }
 
     // Prepare for late C conversion.
-    bool lateCConvert = (!loadOnly && !strategy.C.atomic && problem.needsTsConvert() && state.Tacc != Ts);
+    bool lateCConvert = (!loadOnly && !atomicUpdate && problem.needsTsConvert() && state.Tacc != Ts);
     bool copyCLoad = needLoad && (copyC || lateCConvert);
     if (lateCConvert && Tc.isComplex()) stub();
 
@@ -986,7 +985,7 @@ void Generator<hw>::updateCLayout(const RegisterLayout &layoutExt, const GRFRang
             setupAddr(C_addrsWith0, state.effC[q], sublayoutWith0, state.inputs.ldc[q], strategy, state, C_params, state.ldcMultiples[q], 1);
         }
 
-        if (strategy.C.atomic) {
+        if (atomicUpdate) {
             // Atomic update.
             // Alpha scaling is done earlier; beta scaling isn't supported.
             if (!problem.alpha1() || !problem.beta1()) stub();
@@ -1185,6 +1184,7 @@ bool Generator<hw>::doStdCRemainder(RegisterLayout &layoutExt, RegisterLayout &l
     if (!C_blockUnmasked0 && !layoutExtUnmasked.empty()) C_blockUnmasked0 = &layoutExtUnmasked[0];
 
     bool canEOT = !state.isNested && (op == COperation::UpdateStore);
+    bool atomicUpdate = strategy.C.atomic && (op == COperation::UpdateStore);
 
     Label lEnd;
 
@@ -1199,7 +1199,7 @@ bool Generator<hw>::doStdCRemainder(RegisterLayout &layoutExt, RegisterLayout &l
     status << status_stream::endl;
 
     // Allocate temporaries for emulated atomic addition if needed.
-    if (!inside && strategy.C.atomic) allocEAtomicAddRegs(hw, Tc_ext, layoutExt, problem.C, strategy.C, state);
+    if (!inside && atomicUpdate) allocEAtomicAddRegs(hw, Tc_ext, layoutExt, problem.C, strategy.C, state);
 
     // Handle a subproblem. Return true if successful.
     auto descend = [&](RegisterLayout &sublayoutExt, RegisterLayout &sublayoutExtUnmasked, bool full = false) -> bool {
@@ -1586,7 +1586,7 @@ bool Generator<hw>::doStdCRemainder(RegisterLayout &layoutExt, RegisterLayout &l
     mark(lEnd);
     success ? appendCurrentStream() : discardStream();
 
-    if (!inside && strategy.C.atomic) freeEAtomicAddRegs(state);
+    if (!inside && atomicUpdate) freeEAtomicAddRegs(state);
 
     return success;
 }
@@ -1602,7 +1602,8 @@ void Generator<hw>::doAlternateCRemainder(COperation op, const GEMMProblem &prob
 #define FOR_EACH_C_REV for (int q = C_count - 1; q >= 0; q--)
 
     bool lateYLoopCheck = false;
-    bool atomic = strategy.C.atomic;
+    bool atomicUpdate = strategy.C.atomic && (op == COperation::UpdateStore);
+    bool atomicLoad   = strategy.C.atomic && !atomicUpdate;
 
     bool surface = !strategy.C.base.isStateless();
     bool loadOnly = (op == COperation::Load);
@@ -1620,7 +1621,7 @@ void Generator<hw>::doAlternateCRemainder(COperation op, const GEMMProblem &prob
         nec = nbytes >> 2;
 
     // 8-byte+ types can use scattered qword. Only atomic for now.
-    bool nativeAtomic = atomic && hasNativeAtomicAdd(hw, Tc_ext.real(), problem.C, strategy.C);
+    bool nativeAtomic = atomicUpdate && hasNativeAtomicAdd(hw, Tc_ext.real(), problem.C, strategy.C);
     bool qword = false;
     int rshift = qword ? 3 : 2;     // log2(data stride in regs)
     int rsimd = 64 >> rshift;
@@ -1904,7 +1905,7 @@ void Generator<hw>::doAlternateCRemainder(COperation op, const GEMMProblem &prob
 
 #undef IGNORE_SWSB
 
-    if (atomic) {
+    if (atomicUpdate) {
         // Atomic update. Requires beta = 0/1, alpha prescaled.
         if (!problem.alpha1() || !problem.beta1()) stub();
         if (C_count > 1) stub();
@@ -1958,7 +1959,12 @@ void Generator<hw>::doAlternateCRemainder(COperation op, const GEMMProblem &prob
         // Regular update.
         if (loadOnly || !problem.beta0()) {
             doReadSuppressionWA(strategy, state);
-            if (strategy.C.newDP) {
+            if (atomicLoad && hw >= HW::Xe2) {
+                if (!strategy.C.newDP) stub();
+                !byte_access         ? atomic(AtomicOp::load, 16 | mod, Cload, D32    | strategy.C.cachingR, strategy.C.base, header[0]) :
+                (Tc_ext.size() == 2) ? atomic(AtomicOp::load, 16 | mod, Cload, D16U32 | strategy.C.cachingR, strategy.C.base, header[0])
+                                     : stub();
+            } else if (strategy.C.newDP) {
                 !byte_access         ? load(16 | mod, Cload, D32    | strategy.C.cachingR, strategy.C.base, header[0]) :
                 (Tc_ext.size() == 2) ? load(16 | mod, Cload, D16U32 | strategy.C.cachingR, strategy.C.base, header[0])
                                      : load(16 | mod, Cload, D8U32  | strategy.C.cachingR, strategy.C.base, header[0]);
@@ -2203,9 +2209,9 @@ void Generator<hw>::gemmAccessSums(COperation op, const GEMMProblem &problem, co
     auto Tco = problem.Tco;
     auto cor = sumA ? strategy.unroll[LoopM] : 1;
     auto coc = sumB ? strategy.unroll[LoopN] : 1;
-    bool atomic = strategy.CO.atomic;
+    bool atomicUpdate = strategy.CO.atomic && (op == COperation::UpdateStore);
     bool loadOnly = (op == COperation::Load);
-    bool load = (op != COperation::Store && !problem.beta0() && !(problem.beta1() && atomic));
+    bool load = (op != COperation::Store && !problem.beta0() && !(problem.beta1() && atomicUpdate));
 
     auto CO = problem.CO;
     auto CO_strategy = strategy.CO;
@@ -2282,7 +2288,7 @@ void Generator<hw>::gemmAccessSums(COperation op, const GEMMProblem &problem, co
         }
 
         auto &effCO_regs = share ? Xs_regs : CO_regs;
-        if (atomic) {
+        if (atomicUpdate) {
             allocEAtomicAddRegs(hw, Tco, CO_layout, CO, CO_strategy, state, state.flagAP);
             atomicAddMatrix(effCO_regs, CO_layout, CO_addrs, problem, strategy, state);
             freeEAtomicAddRegs(state, state.flagAP);

src/gpu/intel/gemm/jit/generator/pieces/matrix_access.cxx

@@ -104,7 +104,9 @@ void Generator<hw>::loadMatrixBlock(const Register &dest, const RegisterBlock &b
         case AccessType::Scattered:
         case AccessType::ChannelScattered: {
             auto spec = getDataSpecLSC(atype, astrategy, block, AccessClass::Read);
-            if (block.descAssigned) {
+            if (astrategy.atomic && hw >= HW::Xe2)
+                atomic(AtomicOp::load, mod, dest, spec, astrategy.base, getAddress(addr, block, astrategy));
+            else if (block.descAssigned) {
                 MessageDescriptor desc;
                 ExtendedMessageDescriptor exdesc;
                 encodeLoadDescriptors(hw, desc, exdesc, block.simdSize, r0, spec, astrategy.base, null);