Refactor loop

Signed-off-by: Tiotto, Ettore <[email protected]>

Author: etiotto

third_party/intel/lib/TritonIntelGPUToLLVM/DotOpToLLVM/DPAS.cpp

@@ -191,38 +191,43 @@ class DotOpDPASConversionHelper {
 
     ArrayRef<unsigned> repCluster = dpasEncoding.getRepCluster();
     unsigned rank = repCluster.size();
-    bool aggressiveReusing =
-        triton::tools::getBoolEnv("TRITON_INTEL_AGGRESSIVE_DPAS_REUSE");
+
+    auto innerLoop = [&](int b, int k, int outer, unsigned repNumM,
+                         unsigned repNumN, unsigned repInner,
+                         bool reverseLoop = false) {
+      auto body = [&](int b, int k, int outer, int inner) {
+        if (repNumM > repNumN)
+          generateDPASOp(b, inner, outer, k);
+        else
+          generateDPASOp(b, outer, inner, k);
+      };
+
+      if (reverseLoop) {
+        for (int inner = repInner - 1; inner >= 0; --inner)
+          body(b, k, outer, inner);
+        return;
+      }
+
+      for (int inner = 0; inner < repInner; ++inner)
+        body(b, k, outer, inner);
+    };
+
     // Use the smaller of the two dimensions as the outer loop for better DPAS
     // operands locality.
+    bool aggressiveReusing =
+        triton::tools::getBoolEnv("TRITON_INTEL_AGGRESSIVE_DPAS_REUSE");
     unsigned repNumM = repM * repCluster[rank - 2];
     unsigned repNumN = repN * repCluster[rank - 1];
     unsigned repOuter = repNumM > repNumN ? repNumN : repNumM;
     unsigned repInner = repNumM > repNumN ? repNumM : repNumN;
     for (int b = 0; b < repBatch; ++b)
-      for (int k = 0; k < repK; ++k) {
+      for (int k = 0; k < repK; ++k)
         for (int outer = 0; outer < repOuter; ++outer) {
-          if (aggressiveReusing && ((outer % 2) == 1)) {
-            // Change the inner loop direction in odd outer loop iteration if
-            // aggressive reuse DPAS operands.
-            for (int inner = repInner - 1; inner >= 0; --inner) {
-              if (repNumM > repNumN) {
-                generateDPASOp(b, inner, outer, k);
-              } else {
-                generateDPASOp(b, outer, inner, k);
-              }
-            }
-          } else {
-            for (int inner = 0; inner < repInner; ++inner) {
-              if (repNumM > repNumN) {
-                generateDPASOp(b, inner, outer, k);
-              } else {
-                generateDPASOp(b, outer, inner, k);
-              }
-            }
-          }
+          // Change the inner loop direction in odd outer loop iteration if
+          // aggressive reuse DPAS operands.
+          bool reverseLoop = aggressiveReusing && ((outer % 2) == 1);
+          innerLoop(b, k, outer, repNumM, repNumN, repInner, reverseLoop);
         }
-      }
 
     Value res = composeValuesToDotOperandLayoutStruct(fc, repBatch, repM, repN,
                                                       resElemTy);