[MLIR] Simplify division extraction unit testing.

[pashu123]

Summary

The new getLocalReprs function also outputs dividends and
denominators and hence the CheckDivisionRepresentation fn is
modified to take the newer getLocalReprs function into account.

Signed-off-by: Prashant Kumar <[e-mail]>

[pashu123]

Clang format.

[pashu123]

Removing unnecessary changes.

[Groverkss]

Thank you for this simplification!

Could you change the revision title to better reflect the changes? Something like [MLIR] Simplify division extraction unit testing maybe?

mlir/unittests/Analysis/AffineStructuresTest.cpp

654–656

Please use camelBack for variables (https://mlir.llvm.org/getting_started/DeveloperGuide/)

659–661

I don't think you need to extract the inequalities here (repr).

663–667

These asserts should be at top of the function body.

669–683

I think you can remove this. We don't need to check repr at all. Checking dividends and denominators is enough.

676–679

You can replace this with EXPECT_EQ(trueDividends == dividends).

681–683

You can replace this with EXPECT_EQ(trueDenominators == denominators).

[Groverkss]

mlir/unittests/Analysis/AffineStructuresTest.cpp

654–656

I think expectedDividends, expectedDenominators is a better name. Also, doc comment for this function needs to be updated with this.

[pashu123]

mlir/unittests/Analysis/AffineStructuresTest.cpp

681–683

I think it's EXPECT_TRUE(...), EXPECT_EQ macro expects two arguments.

[pashu123]

Addressing Groverkss comments.

[pashu123]

mlir/unittests/Analysis/AffineStructuresTest.cpp

663–667

The size asserts can only happen after passing the variables through getLocalRepr fn.

[pashu123]

Renaming variable names in comments.

[Groverkss]

mlir/unittests/Analysis/AffineStructuresTest.cpp

682–686

Since this is a test, we can get rid of these asserts also. EXPECT_TRUE below will catch this.

687–692

[pashu123]

Removing asserts.

[Groverkss]

LGTM. Thank you!

mlir/unittests/Analysis/AffineStructuresTest.cpp

 /// Check if the expected division representation of local variables matches the
 /// computed representation. The expected division representation is given as
 /// a vector of expressions set in `divisions` and the corressponding
 /// denominator in `denoms`. If expected denominator for a variable is
 /// non-positive, the local variable is expected to not have a computed
 /// representation.
 static void checkDivisionRepresentation(
     FlatAffineConstraints &fac,
-    const std::vector<SmallVector<int64_t, 8>> &divisions,
-    const SmallVector<int64_t, 8> &denoms) {
+    const std::vector<SmallVector<int64_t, 8>> &true_dividends,
+    const SmallVector<int64_t, 8> &true_denominators) {
 
-  assert(divisions.size() == fac.getNumLocalIds() &&
-         "Size of expected divisions does not match number of local variables");
-  assert(
-      denoms.size() == fac.getNumLocalIds() &&
-      "Size of expected denominators does not match number of local variables");
-
-  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> res(
-      fac.getNumLocalIds(), llvm::None);
-  fac.getLocalReprs(res);
+  std::vector<SmallVector<int64_t, 8>> dividends;
+  SmallVector<unsigned, 4> denominators;
+  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr;
+
+  fac.getLocalReprs(dividends, denominators, repr);
+
+  assert(true_dividends.size() == dividends.size() &&
+         "Size of expected dividends does not match the calculated dividends");
+  assert(true_denominators.size() == denominators.size() &&
+         "Size of expected denominators does not match the calculated "
+         "denominators");
 
   // Check if all expected divisions are computed.
   for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i)
-    if (denoms[i] > 0)
-      EXPECT_TRUE(res[i].hasValue());
+    if (true_denominators[i] > 0)
+      EXPECT_TRUE(repr[i].hasValue());
     else
-      EXPECT_FALSE(res[i].hasValue());
+      EXPECT_FALSE(repr[i].hasValue());
 
-  unsigned divOffset = fac.getNumDimAndSymbolIds();
-  for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i) {
-    if (!res[i])
-      continue;
-
-    // Check if the bounds are of the form:
-    //      0 <= expr - divisor * id <= divisor - 1
-    // Rearranging, we have:
-    //       divisor * id - expr + (divisor - 1) >= 0  <-- Lower bound for 'id'
-    //      -divisor * id + expr                 >= 0  <-- Upper bound for 'id'
-    // where `id = expr floordiv divisor`.
-    unsigned ubPos = res[i]->first, lbPos = res[i]->second;
-    const SmallVector<int64_t, 8> &expr = divisions[i];
-
-    // Check if lower bound is of the correct form.
-    int64_t computedDivisorLb = fac.atIneq(lbPos, i + divOffset);
-    EXPECT_EQ(computedDivisorLb, denoms[i]);
-    for (unsigned c = 0, f = fac.getNumLocalIds(); c < f; ++c) {
-      if (c == i + divOffset)
-        continue;
-      EXPECT_EQ(fac.atIneq(lbPos, c), -expr[c]);
-    }
-    // Check if constant term of lower bound matches expected constant term.
-    EXPECT_EQ(fac.atIneq(lbPos, fac.getNumCols() - 1),
-              -expr.back() + (denoms[i] - 1));
-
-    // Check if upper bound is of the correct form.
-    int64_t computedDivisorUb = fac.atIneq(ubPos, i + divOffset);
-    EXPECT_EQ(computedDivisorUb, -denoms[i]);
-    for (unsigned c = 0, f = fac.getNumLocalIds(); c < f; ++c) {
-      if (c == i + divOffset)
-        continue;
-      EXPECT_EQ(fac.atIneq(ubPos, c), expr[c]);
-    }
-    // Check if constant term of upper bound matches expected constant term.
-    EXPECT_EQ(fac.atIneq(ubPos, fac.getNumCols() - 1), expr.back());
-  }
+  // Check that the `dividends` and `true dividends` match.
+  for (unsigned i = 0, e = dividends.size(); i < e; ++i)
+    for (unsigned j = 0, m = dividends[i].size(); j < m; ++j)
+      EXPECT_EQ(true_dividends[i][j], dividends[i][j]);
+
+  // Check that the `denominators` and `true_denominators` match.
+  for (unsigned i = 0, e = denominators.size(); i < e; ++i)
+    EXPECT_EQ(true_denominators[i], denominators[i]);
 }
 
 TEST(FlatAffineConstraintsTest, computeLocalReprSimple) {
   FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});
 
   fac.addLocalFloorDiv({1, 4}, 10);
   fac.addLocalFloorDiv({1, 0, 100}, 10);
 
   std::vector<SmallVector<int64_t, 8>> divisions = {{1, 0, 0, 4},
                                                     {1, 0, 0, 100}};
   SmallVector<int64_t, 8> denoms = {10, 10};
 
   // Check if floordivs can be computed when no other inequalities exist
   // and floor divs do not depend on each other.
   checkDivisionRepresentation(fac, divisions, denoms);
 }
 
TEST(FlatAffineConstraintsTest, computeLocalReprRecursive) {
 
   fac.addLocalFloorDiv({0, -2, 7, 2, 10}, 3);
   fac.addLocalFloorDiv({3, 0, 9, 2, 2, 10}, 5);
   fac.addLocalFloorDiv({0, 1, -123, 2, 0, -4, 10}, 3);
 
   fac.addInequality({1, 2, -2, 1, -5, 0, 6, 100});
   fac.addInequality({1, 2, -8, 1, 3, 7, 0, -9});
 
-  std::vector<SmallVector<int64_t, 8>> divisions = {{0, -2, 7, 2, 0, 0, 0, 10},
-                                                    {3, 0, 9, 2, 2, 0, 0, 10},
-                                                    {0, 1, -123, 2, 0, -4, 10}};
+  std::vector<SmallVector<int64_t, 8>> divisions = {
+      {0, -2, 7, 2, 0, 0, 0, 10},
+      {0, 1, -123, 2, 0, -4, 0, 10}};
+
   SmallVector<int64_t, 8> denoms = {3, 5, 3};
 
   // Check if floordivs which may depend on other floordivs can be computed.
   checkDivisionRepresentation(fac, divisions, denoms);
 }
 
 TEST(FlatAffineConstraintsTest, removeIdRange) {
   FlatAffineConstraints fac(3, 2, 1);
TEST(FlatAffineConstraintsTest, mergeDivisionsConstants) {
     // Local space should be same.
     EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(fac1.getNumLocalIds(), 2u);
     EXPECT_EQ(fac2.getNumLocalIds(), 2u);
   }
 }
-
-} // namespace mlir
+}