diff --git a/mlir/unittests/Dialect/SparseTensor/CMakeLists.txt b/mlir/unittests/Dialect/SparseTensor/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/unittests/Dialect/SparseTensor/CMakeLists.txt @@ -0,0 +1,7 @@ +add_mlir_unittest(MLIRSparseTensorTests + MergerTest.cpp +) +target_link_libraries(MLIRSparseTensorTests + PRIVATE + MLIRSparseTensorUtils +) diff --git a/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp b/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Dialect/SparseTensor/MergerTest.cpp @@ -0,0 +1,46 @@ +#include "mlir/Dialect/SparseTensor/Utils/Merger.h" +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +using namespace mlir; +using namespace mlir::sparse_tensor; + +TEST(MergerTest, VectorAdd4) { + // Vector addition of 4 vectors, as seen in Aart's post here: + // https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020/2 + // 5 tensors (4 inputs and a single output), 1 loop. + Merger merger = Merger(5, 1); + auto b = merger.addExp(Kind::kTensor, 0); + auto c = merger.addExp(Kind::kTensor, 1); + auto d = merger.addExp(Kind::kTensor, 2); + auto e = merger.addExp(Kind::kTensor, 3); + merger.setDim(b, 0, Dim::kSparse); + merger.setDim(c, 0, Dim::kSparse); + merger.setDim(d, 0, Dim::kSparse); + merger.setDim(e, 0, Dim::kSparse); + auto add0 = merger.addExp(Kind::kAddF, b, c); + auto add1 = merger.addExp(Kind::kAddF, d, e); + auto add2 = merger.addExp(Kind::kAddF, add0, add1); + auto l = merger.optimizeSet(merger.buildLattices(add2, 0)); + EXPECT_THAT(merger.set(l).size(), 15); +} + +TEST(MergerTest, VectorMul4) { + // Vector multiplication of 4 vectors, as seen in Aart's post here: + // https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020/2 + // 5 tensors (4 inputs and a single output), 1 loop. + Merger merger = Merger(5, 1); + auto b = merger.addExp(Kind::kTensor, 0); + auto c = merger.addExp(Kind::kTensor, 1); + auto d = merger.addExp(Kind::kTensor, 2); + auto e = merger.addExp(Kind::kTensor, 3); + merger.setDim(b, 0, Dim::kSparse); + merger.setDim(c, 0, Dim::kSparse); + merger.setDim(d, 0, Dim::kSparse); + merger.setDim(e, 0, Dim::kSparse); + auto mul0 = merger.addExp(Kind::kMulF, b, c); + auto mul1 = merger.addExp(Kind::kMulF, d, e); + auto mul2 = merger.addExp(Kind::kMulF, mul0, mul1); + auto l = merger.optimizeSet(merger.buildLattices(mul2, 0)); + EXPECT_THAT(merger.set(l).size(), 1); +}