Simplify the logic and fix comments

Author: kasuga-fj

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

@@ -73,15 +73,6 @@ using LoopVector = SmallVector<Loop *, 8>;
 // TODO: Check if we can use a sparse matrix here.
 using CharMatrix = std::vector<std::vector<char>>;
 
-// Classification of a direction vector by the leftmost element after removing
-// '=' and 'I' from it.
-enum class DirectionVectorPattern {
-  Zero,     ///< The direction vector contains only '=' or 'I'.
-  Positive, ///< The leftmost element after removing '=' and 'I' is '<'.
-  Negative, ///< The leftmost element after removing '=' and 'I' is '>'.
-  All,      ///< The leftmost element after removing '=' and 'I' is '*'.
-};
-
 /// Types of rules used in profitability check.
 enum class RuleTy {
   PerLoopCacheAnalysis,
@@ -240,115 +231,102 @@ static void interChangeDependencies(CharMatrix &DepMatrix, unsigned FromIndx,
     std::swap(DepMatrix[I][ToIndx], DepMatrix[I][FromIndx]);
 }
 
-// Classify the direction vector into the four patterns. The target vector is
-// [DV[Left], DV[Left+1], ..., DV[Right-1]], not the whole of \p DV.
-static DirectionVectorPattern
-classifyDirectionVector(const std::vector<char> &DV, unsigned Left,
-                        unsigned Right) {
-  assert(Left <= Right && "Left must be less or equal to Right");
-  for (unsigned I = Left; I < Right; I++) {
-    unsigned char Direction = DV[I];
-    switch (Direction) {
-    case '<':
-      return DirectionVectorPattern::Positive;
-    case '>':
-      return DirectionVectorPattern::Negative;
-    case '*':
-      return DirectionVectorPattern::All;
-    case '=':
-    case 'I':
-      break;
-    default:
-      llvm_unreachable("Unknown element in direction vector");
-    }
-  }
-  return DirectionVectorPattern::Zero;
+/// Returns the leftmost non-'=' element. If such a element doesn't exist,
+/// returns nullopt. 'I' is treated same as '='.
+std::optional<char> getLeftmostNonEqElement(const std::vector<char> &Dep) {
+  for (char C : Dep)
+    if (C != '=' && C != 'I')
+      return C;
+  return std::nullopt;
 }
 
-// Check whether the requested interchange is legal or not. The interchange is
-// valid if the following condition holds:
-//
-// [Cond] For two instructions that can access the same location, the execution
-// order of the instructions before and after interchanged is the same.
-//
-// If the direction vector doesn't contain '*', the above Cond is equivalent to
-// one of the following:
-//
-// - The leftmost non-'=' element is '<' before and after interchanging.
-// - The leftmost non-'=' element is '>' before and after interchanging.
-// - All the elements in the direction vector is '='.
-//
-// As for '*', we must treat it as having dependency in all directions. It could
-// be '<', it could be '>', it could be '='. We can eliminate '*'s from the
-// direction vector by enumerating all possible patterns by replacing '*' with
-// '<' or '>' or '=', and then doing the above checks for all of them. The
-// enumeration can grow exponentially, so it is not practical to run it as it
-// is. Fortunately, we can perform the following pruning.
-//
-// - For '*' to the left of \p OuterLoopId, replacing it with '=' is allowed.
-//
-// This is because, for patterns where '<' (or '>') is assigned to some '*' to
-// the left of \p OuterLoopId, the first (or second) condition above holds
-// regardless of interchanging. After doing this pruning, the interchange is
-// legal if the leftmost non-'=' element is the same before and after swapping
-// the element of \p OuterLoopId and \p InnerLoopId.
-//
-//
-// Example: Consider the following loop.
-//
-// ```
-// for (i=0; i<=32; i++)
-//   for (j=0; j<N-1; j++)
-//     for (k=0; k<N-1; k++) {
-// Src:   A[i][j][k] = ...;
-// Dst:   use(A[32-i][j+1][k+1]);
-//     }
-// ```
-//
-// In this case, the direction vector is [* < <] (if the analysis is powerful
-// enough). The enumeration of all possible patterns by replacing '*' is as
-// follows:
-//
-// - [< < <] : when i < 16
-// - [= < <] : when i = 16
-// - [> < <] : when i > 16
-//
-// We can prove that it is safe to interchange the innermost two loops here,
-// because the interchange doesn't change the leftmost non-'=' element for all
-// enumerated vectors.
-//
-// TODO: There are cases where the interchange is legal but rejected. At least
-// the following patterns are legal:
-// - If both Dep[OuterLoopId] and Dep[InnerLoopId] are '=', the interchange is
-// legal regardless of any other elements.
-// - If the loops are adjacent to each other and at least one of them is '=',
-// the interchange is legal even if the other is '*'.
+/// Check whether the requested interchange is legal or not. The interchange is
+/// valid if the following condition holds:
+///
+/// [Cond] For all two instructions pairs that can access the same location, the
+/// execution order of them before and after exchanging is the same.
+///
+/// If the direction vector doesn't contain '*', the above [Cond] is equivalent
+/// to one of the following:
+///
+/// - The leftmost non-'=' element is '<' before and after interchanging.
+/// - The leftmost non-'=' element is '>' before and after interchanging.
+/// - All the elements in the direction vector is '='.
+///
+/// This fact implies that it is allowed to "decompose" a direction vector with
+/// a '*' into three direction vectors where the original '*' is replaced with
+/// '<', '>' or '=', to perform the legality check on each of them, and to use
+/// the logical and of those values as the result of the legality check for the
+/// original direction vector. For example, if we have a direction vector
+/// [* < =], then the following holds:
+///
+/// isLegal([* < =], OuterId, InnerId) is true if all the following is true.
+/// - isLegal([< < =], OuterId, InnerId)
+/// - isLegal([> < =], OuterId, InnerId)
+/// - isLegal([= < =], OuterId, InnerId)
+///
+/// This can be easily extended for a vector with multiple '*'s by enumerating
+/// all possible vectors (let N be the number of '*', then there are 3^N
+/// possible vectors).
+///
+/// In practice, such a combinatorial explosion is undesirable. To address this,
+/// we will only consider replacing '*'s to the left of OuterId. This allows us
+/// to omit vectors where any '*' is replaced with either '<' or '>'. This is
+/// because, exchanging the given two loops obviously doesn't change the
+/// leftmost non-'=' element ('<' or '>'). For example, consider the direction
+/// vector is [* < =] as in the previous example. If OuterId is 2 and InnerId is
+/// 3, then isLegal([< < =], ...) and isLegal([> < =], ...) is trivially true
+/// because the first element ensures that the leftmost non-'=' element doesn't
+/// change. Therefore, it is sufficient to check the legality for [= < =].
+///
+/// Therefore, it is sufficient to consider only a single direction vector where
+/// all '*'s to the left of OuterId in the original vector are replaced with
+/// '='. For the same reason, if there are one or more '<' or '>' to the left of
+/// OuterId, it is legal to exchange the given two loops.
+///
+/// This function performs the following algorithm, which is sound due to the
+/// above facts:
+///
+///   Step 1: If there is '<' or '>' to the left of OuterId, then interchanging
+///           the two loops is legal.
+///
+///   Step 2: If all elements to the left of OuterId are '=' or '*', then we
+///           check that the leftmost non-'=' element is not '*' and swapping
+///           the two loops doesn't change it.
 static bool isLegalToInterchangeLoopsForRow(std::vector<char> Dep,
                                             unsigned InnerLoopId,
                                             unsigned OuterLoopId) {
-  // Replace '*' to the left of OuterLoopId with '='. The presence of '<' means
-  // that the direction vector is something like [= = = < ...], where the
-  // interchange is safe.
+  assert(OuterLoopId + 1 == InnerLoopId && "The two loops must be adjacent.");
+
+  // Step 1: Replace '*' to the left of OuterLoopId with '='. If we find '<' or
+  // '>' on the way, then returns true.
   for (unsigned I = 0; I < OuterLoopId; I++) {
     if (Dep[I] == '<' || Dep[I] == '>')
       return true;
     Dep[I] = '=';
   }
 
-  // From this point on, all elements to the left of OuterLoopId are considered
-  // to be '='.
+  // Step 2: Check if the leftmost non-'=' element is same before and after the
+  // exchange. If either one of them is '*', then return false conseratively.
+  // TODO: There are some cases where the exchange is legal even if one of them
+  // is '*', e.g., the two loops are adjacent and the other one element '='.
+  std::optional<char> Before = getLeftmostNonEqElement(Dep);
 
-  // Perform legality checks by comparing the leftmost non-'=' element between
-  // before and after the interchange. If either one is '*', then the
-  // interchange is unsafe. Otherwise it is safe if the element is equal.
-  auto BeforePattern =
-      classifyDirectionVector(Dep, OuterLoopId, InnerLoopId + 1);
-  if (BeforePattern == DirectionVectorPattern::All)
+  // The vector only contains '=' or 'I'. Exchanging the two loops is legal in
+  // this case.
+  if (!Before)
+    return true;
+  if (*Before == '*')
     return false;
+
+  assert((*Before == '<' || *Before == '>') && "Unexpected element.");
   std::swap(Dep[InnerLoopId], Dep[OuterLoopId]);
-  auto AfterPattern =
-      classifyDirectionVector(Dep, OuterLoopId, InnerLoopId + 1);
-  return BeforePattern == AfterPattern;
+  std::optional<char> After = getLeftmostNonEqElement(Dep);
+  assert(After.has_value() && "Something unexpected happened.");
+
+  // At this point, Before is either '<' or '>'. So we don't need to check if
+  // After is '*' because this comparison is false in such a case.
+  return *Before == *After;
 }
 
 // Checks if it is legal to interchange 2 loops.