[mlir][sparse] Add Merger unit tests

[gussmith23]

Summary

We opt to use unit tests rather than check tests as the lattice/merger code is a small C++ component with a well-defined API. Testing this API via check tests would be far less direct and readable. In addition, as the check tests will only be able to test the API indirectly, the tests may break based on unrelated changes; e.g. changes in linalg.

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

10–11

Links like that may go stale in the (far?) future.

I think it would be more illustrative to just show the "tensor index expression" in human readable format in the comments

26

this is a good start, but we probably also want to make sure the lattices have the expected structure

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

25

how about a before opt after opt check as well?

[gussmith23]

More thorough checking in tests

[gussmith23]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

10–11

Great point. Let me know if my new comments make sense.

25

I will add this -- didn't add it in my latest update, but I'll get to it!

26

Let me know if the updated tests are thorough enough, and if they're testing for the right things.

[mehdi_amini]

Can you expand in the commit description why this is more suitable as a unit-tests and can't be tested on IR change (folding on simple sparse operations for example that would exercises all theses cases?).

[aartbik]

In D104956#2845894, @mehdi_amini wrote:

Can you expand in the commit description why this is more suitable as a unit-tests and can't be tested on IR change.

I'll take this, since I was suggesting this to Gus as part of the refactoring :-)

To me, unit testing classes with a well-defined API for a large variety of expected input/output scenarios (and the lattice code falls into that category) seems much more preferable over setting up elaborate IR in an attempt to test those scenarios?

But, made curious by by your feedback, reading our testing guide, I indeed see

Unit tests are written using Google Test and are located in the unittests/ directory. Tests of these form should be limited to API tests that cannot be reasonably written as Check tests, e.g. those for data structures. It is important to keep in mind that the C++ APIs are not stable, and evolve over time. As such, directly testing the C++ IR interfaces makes the tests more fragile as those C++ APIs evolve over time. This makes future API refactorings, which may happen frequently, much more cumbersome as the number of tests scale.

I suppose this is the first time I categorically disagree with the llvm guidelines ;-)

But, even despite my objections, in this case, I suppose we could make a good case for unit testing, since testing through checks would require testing the generated code rather than expected lattices (so hunting for while-loops and for-loops of a particular form, and with a particular body) rather than simply checking for expected lattices with tensor subexpressions. I will work with Gus on putting a rationale in the commit message.

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

33

isn't this a bit too strict? isn't it possible to have two nodes for the same tensor?

so tensorExp0.e0 == tensorExpr1.e0

seems a better way of comparing two tensor exprs

92–107

This seems a very elaborate way of hunting for a particular subexpression
(1) not easy to read (for humans when glancing over the test
(2) you scan for the expression in the set, which seems expensive, but also does not take order into account (which matters in this case, making "set" a bit of a misnomer I am afraid)

Let's brainstorm offline a bit on how to make these tests more concise

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

74

It seems all tests will consist of building a tensor expression, followed by testing if the computed lattices are as expected. I suggest to use a test fixture (subclass ::testing::Test and put all shared logic to setup and test all unit tests there, then use TEST_F instead of TEST).

Also, I would like to make the tests readable at quick glance, so something like

EXPECT_THAT(lattice(0), "tensor_0 + tensor_1")

a quick and dirty way would be string based, so either

(1) add a method to the fixture that compares while parsing the expected string, or
(2) add toString instead of dump methods to the lattice code and compare that one against expected

or, we would add some intuitive methods to the fixture to make the expected more readable,
e.g.

EXPECT_THAT(lattice(0), addf(tensor(0), tensor(1))

The short methods have all the ugly details hidden inside the fixture (use of merger), but give very concise testing

Then finally, we need a way to make it less sensitive to a particular order, but while still making sure the required order is preserved.

Lastly, let's start with unit tests for simple cases, e.g. 1 or 2 tensors only, and then later add the more elaborate ones.

[gussmith23]

Address one of Aart's comments

[gussmith23]

Rebase

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

33

Yes, I wanted to ask you about this! I didn't fully understand this when I originally saw it. Do some tensor expressions point to themselves? That is, let's say expression with id 0 has Kind kTensor; will this tensor have e0 = 0?

92–107

Re: (1), I agree, I would like a way to make it more readable! Re: (2), I didn't realize order was important, my mistake. That should simplify things considerably!

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

33

Ah, I see the confusion. Note that perhaps we should make this a union.

For binary ops, the e0/e1 are indices to other subexpressions.
For tensors, e0 simply denotes the tensor value (interpreted as int)

[gussmith23]

• Simplify tests (test 2-arg add and mul, rather than 4-arg)
• Use a test fixture, move logic into test fixture class
• Make tests a bit more readable (though still needs work)
  • The main change here is that we now use helper functions to check that lattice point expressions and bits equal what we expect them to equal. Those helpers are designed to take a more readable set of arguments, e.g. a prefix-notation expression to test against, in the case of the expression checker.
• Check for a specific ordering of lattice points (this may break in the future!)
  • If we were really thorough, we'd check for a partial ordering on the lattice points. However, we just check for an absolute ordering.

[gussmith23]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

33

Fixed by the union refactor

74

I think I have addressed all points except this one:

Then finally, we need a way to make it less sensitive to a particular order, but while still making sure the required order is preserved.

This I could imagine doing by building a function which, given two lattice points, checks for the given ordering between the lattice points. We'll need an ergonomic way to "find" a lattice point, e.g. by its expression. Once we "find" all lattice points, we can then do pairwise checking between the points, to make sure they are partially ordered correctly.

92–107

Removed this function. See comment above re: checking partial ordering.

[gussmith23]

Rebase and fix clang-tidy

[aartbik]

mlir/unittests/Dialect/SparseTensor/CMakeLists.txt

2

I believe you also have to cd .. and add

add_subdirectory(SparseTensor)

to the CMakeList.txt at the one-higher level.

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

15

How about having a generic base fixture class MergerTest that takes these two parameters in the constructor, but specific fixture classes for #tensors and #loops, just so that we can later add tests for other situations by just instantiating a new case with a a few lines of code, e.g. MergerTest48 : MergerTest(4,8), MergerTest11 : MergerTest(1,1) etc. with some tensor setup code inside the constructor (perhaps better name than this but you get the idea)?

79–80

I should have put this in the getter comments, but in general holding on to the const ref for too long can have very subtle bugs. Here it is okay since you don't insert anything

but using merger.exp(e0).children.e0 inline is preferred in general

81–83

no braces {} for single line if

91–96

replace by just default: ?

126

I find these comments a bit distracting.
How about adding a similar method as tensor()
e.g. lat()

so that this just reads as lat(0) instead of /*lat point*/ 0

Also, this makes the test a bit implementation specific (ie. a complete order is tested, but we should really just expect a partial order on the output).

Do you have any good ideas on how to deal with that?

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

30

Wait, by the way this is used below, should this not add a full tensor node?
The way it is now, the children will "contain" the tensor value?

Also, there is something impure about using the same merger for testing if an expression exists.
Do you see a cleaner way around that?

[gussmith23]

Address some of Aart's comments (not all)

[gussmith23]

Address all of Aart's comments except for partial ordering

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

15

Let me know if the way I did this makes sense!

30

Wait, by the way this is used below, should this not add a full tensor node?

You're totally right -- fixed now. This came from me misunderstanding how the expressions worked in the Merger (back when I thought tensornum == expression id)

Also, there is something impure about using the same merger for testing if an expression exists.

Let me know if the new pattern approach works for you.

79–80

Alright, I'll leave it if it's okay here, but I'll keep that in mind for the future!

79–80

Good to know -- I'll leave this here if it's okay, but I'll keep this in mind for the future.

91–96

For some reason I've always been very paranoid about bugs arising from adding new enum variants and having those variants be handled by a "default", rather than throwing a compile-time error. I'm happy to use default here if you don't think it'll be an issue!

126

Made it a bit more readable I think, removed the comments. Let me know. Working on partial ordering now.

[aartbik]

I really like the direction this is going. Add the partial order, and the tests have become very readable!!!

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

15

love it!

91–96

Fair enough.

[gussmith23]

Address partial ordering; clean up unused code

[gussmith23]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

74

I implemented your idea; let me know if it looks the way you were expecting!

126

Let me know if the new partial-ordering checks work for you.

[gussmith23]

Fix comments

[aartbik]

almost there!

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

46–50

I would save the vertical space here, and just put the comments "against" the constructor, even without the leading empty ///
(bit subjective of course)

203

actually, one last subtle point, note that since we have

a(i) = b(i) + c(i)

we really have *three* tensors in this example, with the third one (a), unused but dense

This will matter when we start testing stuff like

a(i) = a(i) + .....

and also avoids dumping tensor_out for input tensors during debugging. So please add one etra tensor in the constructor, and mark output as dense

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

121

why do we need to clone here, can't we just the pointer?

129

shall we make a convenience wrapper

expectLatPoint(s, p ... ) -> expectLaPointWithinRange(s, p, p+1, ....)

for the single cases, just to make those even more readable?

[gussmith23]

Convert to shared_ptr, add output tensor

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

121

Fixed by changing to shared_ptr!

203

Hm, adding the output tensor actually changed one of the tests. Does that seem right?

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

203

yes, we optimized a(i) = a(i) away before, now fixed ;-)

[aartbik]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

13–16

one last nit, can we do without this, just to keep boilerplate down?

[gussmith23]

Remove children struct

[gussmith23]

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

13–16

Done!

[aartbik]

Nice solid work, Guss! Thank you.

Ship it!

[mehdi_amini]

Hey, sorry I had to revert to unbreak gcc 5 builds: https://buildkite.com/mlir/mlir-core/builds/15063#d35a9e3f-da90-4e84-86be-c69cb67b2b5c

mlir/unittests/Dialect/CMakeLists.txt

 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
 )
 target_link_libraries(MLIRDialectTests
   PRIVATE
   MLIRIR
   MLIRDialect)
 
 add_subdirectory(Quant)
+add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)

mlir/unittests/Dialect/SparseTensor/CMakeLists.txt

+add_mlir_unittest(MLIRSparseTensorTests
+  MergerTest.cpp
+)
+target_link_libraries(MLIRSparseTensorTests
+  PRIVATE
+  MLIRSparseTensorUtils
+)

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

+#include "mlir/Dialect/SparseTensor/Utils/Merger.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include <memory>
+
+using namespace mlir::sparse_tensor;
+
+namespace {
+
+/// Simple recursive data structure used to match expressions in Mergers.
+struct Pattern {
+  Kind kind;
+
+  /// Expressions representing tensors simply have a tensor number.
+  unsigned tensorNum;
+
+  /// Tensor operations point to their children.
+  std::shared_ptr<Pattern> e0;
+  std::shared_ptr<Pattern> e1;
+
+  /// Constructors.
+  /// Rather than using these, please use the readable helper constructor
+  /// functions below to make tests more readable.
+  Pattern(unsigned tensorNum) : kind(Kind::kTensor), tensorNum(tensorNum) {}
+  Pattern(Kind kind, std::shared_ptr<Pattern> e0, std::shared_ptr<Pattern> e1)
+      : kind(kind), e0(e0), e1(e1) {
+    assert(kind >= Kind::kMulF);
+    assert(e0 && e1);
+  }
+};
+
+///
+/// Readable Pattern builder functions.
+/// These should be preferred over the actual constructors.
+///
+
+static std::shared_ptr<Pattern> tensorPattern(unsigned tensorNum) {
+  return std::make_shared<Pattern>(tensorNum);
+}
+
+static std::shared_ptr<Pattern> addfPattern(std::shared_ptr<Pattern> e0,
+                                            std::shared_ptr<Pattern> e1) {
+  return std::make_shared<Pattern>(Kind::kAddF, e0, e1);
+}
+
+static std::shared_ptr<Pattern> mulfPattern(std::shared_ptr<Pattern> e0,
+                                            std::shared_ptr<Pattern> e1) {
+  return std::make_shared<Pattern>(Kind::kMulF, e0, e1);
+}
+
+class MergerTestBase : public ::testing::Test {
+protected:
+  MergerTestBase(unsigned numTensors, unsigned numLoops)
+      : numTensors(numTensors), numLoops(numLoops),
+        merger(numTensors, numLoops) {}
+
+  ///
+  /// Expression construction helpers.
+  ///
+
+  unsigned tensor(unsigned tensor) {
+    return merger.addExp(Kind::kTensor, tensor);
+  }
+
+  unsigned addf(unsigned e0, unsigned e1) {
+    return merger.addExp(Kind::kAddF, e0, e1);
+  }
+
+  unsigned mulf(unsigned e0, unsigned e1) {
+    return merger.addExp(Kind::kMulF, e0, e1);
+  }
+
+  ///
+  /// Comparison helpers.
+  ///
+
+  /// For readability of tests.
+  unsigned lat(unsigned lat) { return lat; }
+
+  /// Returns true if a lattice point with an expression matching the given
+  /// pattern and bits matching the given bits is present in lattice points
+  /// [p, p+n) of lattice set s. This is useful for testing partial ordering
+  /// constraints between lattice points. We generally know how contiguous
+  /// groups of lattice points should be ordered with respect to other groups,
+  /// but there is no required ordering within groups.
+  bool latPointWithinRange(unsigned s, unsigned p, unsigned n,
+                           std::shared_ptr<Pattern> pattern,
+                           llvm::BitVector bits) {
+    for (unsigned i = p; i < p + n; ++i) {
+      if (compareExpression(merger.lat(merger.set(s)[i]).exp, pattern) &&
+          compareBits(s, i, bits))
+        return true;
+    }
+    return false;
+  }
+
+  /// Wrapper over latPointWithinRange for readability of tests.
+  void expectLatPointWithinRange(unsigned s, unsigned p, unsigned n,
+                                 std::shared_ptr<Pattern> pattern,
+                                 llvm::BitVector bits) {
+    EXPECT_TRUE(latPointWithinRange(s, p, n, pattern, bits));
+  }
+
+  /// Wrapper over expectLatPointWithinRange for a single lat point.
+  void expectLatPoint(unsigned s, unsigned p, std::shared_ptr<Pattern> pattern,
+                      llvm::BitVector bits) {
+    EXPECT_TRUE(latPointWithinRange(s, p, 1, pattern, bits));
+  }
+
+  /// Converts a vector of (loop, tensor) pairs to a bitvector with the
+  /// corresponding bits set.
+  llvm::BitVector
+  loopsToBits(std::vector<std::tuple<unsigned, unsigned>> loops) {
+    llvm::BitVector testBits = llvm::BitVector(numTensors + 1, false);
+    for (auto l : loops) {
+      auto loop = std::get<0>(l);
+      auto tensor = std::get<1>(l);
+      testBits.set(numTensors * loop + tensor);
+    }
+    return testBits;
+  }
+
+  /// Returns true if the bits of lattice point p in set s match the given bits.
+  bool compareBits(unsigned s, unsigned p, llvm::BitVector bits) {
+    return merger.lat(merger.set(s)[p]).bits == bits;
+  }
+
+  /// Check that there are n lattice points in set s.
+  void expectNumLatPoints(unsigned s, unsigned n) {
+    EXPECT_THAT(merger.set(s).size(), n);
+  }
+
+  /// Compares expressions for equality. Equality is defined recursively as:
+  /// - Two expressions can only be equal if they have the same Kind.
+  /// - Two binary expressions are equal if they have the same Kind and their
+  ///     children are equal.
+  /// - Expressions with Kind invariant or tensor are equal if they have the
+  ///     same expression id.
+  bool compareExpression(unsigned e, std::shared_ptr<Pattern> pattern) {
+    auto tensorExp = merger.exp(e);
+    if (tensorExp.kind != pattern->kind)
+      return false;
+    assert(tensorExp.kind != Kind::kInvariant &&
+           "Invariant comparison not yet supported");
+    switch (tensorExp.kind) {
+    case Kind::kTensor:
+      return tensorExp.tensor == pattern->tensorNum;
+    case Kind::kZero:
+      return true;
+    case Kind::kMulF:
+    case Kind::kMulI:
+    case Kind::kAddF:
+    case Kind::kAddI:
+    case Kind::kSubF:
+    case Kind::kSubI:
+      return compareExpression(tensorExp.children.e0, pattern->e0) &&
+             compareExpression(tensorExp.children.e1, pattern->e1);
+    default:
+      llvm_unreachable("Unhandled Kind");
+    }
+  }
+
+  unsigned numTensors;
+  unsigned numLoops;
+  Merger merger;
+};
+
+class MergerTest3T1L : public MergerTestBase {
+protected:
+  // Our three tensors (two inputs, one output).
+  const unsigned t0 = 0, t1 = 1, t2 = 2;
+
+  // Our single loop.
+  const unsigned l0 = 0;
+
+  MergerTest3T1L() : MergerTestBase(3, 1) {
+    // Tensor 0: sparse input vector.
+    merger.addExp(Kind::kTensor, t0, -1u);
+    merger.setDim(t0, l0, Dim::kSparse);
+
+    // Tensor 1: sparse input vector.
+    merger.addExp(Kind::kTensor, t1, -1u);
+    merger.setDim(t1, l0, Dim::kSparse);
+
+    // Tensor 2: dense output vector.
+    merger.addExp(Kind::kTensor, t2, -1u);
+    merger.setDim(t2, l0, Dim::kDense);
+  }
+};
+
+} // anonymous namespace
+
+/// Vector addition of 2 vectors, i.e.:
+///   a(i) = b(i) + c(i)
+/// which should form the 3 lattice points
+/// {
+///   lat( i_00 i_01 / (tensor_0 + tensor_1) )
+///   lat( i_00 / tensor_0 )
+///   lat( i_01 / tensor_1 )
+/// }
+/// and after optimization, will reduce to the 2 lattice points
+/// {
+///   lat( i_00 i_01 / (tensor_0 + tensor_1) )
+///   lat( i_00 / tensor_0 )
+/// }
+TEST_F(MergerTest3T1L, VectorAdd2) {
+  // Construct expression.
+  auto e = addf(tensor(t0), tensor(t1));
+
+  // Build lattices and check.
+  auto s = merger.buildLattices(e, l0);
+  expectNumLatPoints(s, 3);
+  expectLatPoint(s, lat(0), addfPattern(tensorPattern(t0), tensorPattern(t1)),
+                 loopsToBits({{l0, t0}, {l0, t1}}));
+  expectLatPointWithinRange(s, lat(1), 2, tensorPattern(t0),
+                            loopsToBits({{l0, t0}}));
+  expectLatPointWithinRange(s, lat(1), 2, tensorPattern(t1),
+                            loopsToBits({{l0, t1}}));
+
+  // Optimize lattices and check.
+  s = merger.optimizeSet(s);
+  expectNumLatPoints(s, 3);
+  expectLatPoint(s, lat(0), addfPattern(tensorPattern(t0), tensorPattern(t1)),
+                 loopsToBits({{l0, t0}, {l0, t1}}));
+  expectLatPointWithinRange(s, lat(1), 2, tensorPattern(t0),
+                            loopsToBits({{l0, t0}}));
+  expectLatPointWithinRange(s, lat(1), 2, tensorPattern(t1),
+                            loopsToBits({{l0, t1}}));
+}
+
+/// Vector multiplication of 2 vectors, i.e.:
+///   a(i) = b(i) * c(i)
+/// which should form the single lattice point
+/// {
+///   lat( i_00 i_01 / (tensor_0 * tensor_1) )
+/// }
+TEST_F(MergerTest3T1L, VectorMul2) {
+  // Construct expression.
+  auto e = mulf(t0, t1);
+
+  // Build lattices and check.
+  auto s = merger.buildLattices(e, l0);
+  expectNumLatPoints(s, 1);
+  expectLatPoint(s, lat(0), mulfPattern(tensorPattern(t0), tensorPattern(t1)),
+                 loopsToBits({{l0, t0}, {l0, t1}}));
+
+  // Optimize lattices and check.
+  s = merger.optimizeSet(s);
+  expectNumLatPoints(s, 1);
+  expectLatPoint(s, lat(0), mulfPattern(tensorPattern(t0), tensorPattern(t1)),
+                 loopsToBits({{l0, t0}, {l0, t1}}));
+}