[Mlir-commits] [mlir] [mlir][affine] Add unit tests for `isProjectedPermutation` (PR #114775)

[Andrzej Warzyński]

https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/114775

>From 9f0363923913ffd6f4bb6e6107490cc8883a5ea8 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 4 Nov 2024 11:21:41 +0000
Subject: [PATCH 1/2] [mlir][affine] Add unit tests for
 `isProjectedPermutation`

The only way to test `isProjectedPermutation` is through unit tests. The
concept of "projected permutations" is tricky to document and these
tests are a good source documentation of the expected/intended
behavoiur. Hence these additional unit tests.
---
 mlir/unittests/IR/AffineMapTest.cpp | 41 +++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/mlir/unittests/IR/AffineMapTest.cpp b/mlir/unittests/IR/AffineMapTest.cpp
index 081afadd632f1b..569f0db8c55c9c 100644
--- a/mlir/unittests/IR/AffineMapTest.cpp
+++ b/mlir/unittests/IR/AffineMapTest.cpp
@@ -21,3 +21,44 @@ TEST(AffineMapTest, inferMapFromAffineExprs) {
   map.replace(replacements);
   EXPECT_EQ(map, map);
 }
+
+TEST(AffineMapTest, isProjectedPermutation) {
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+
+  // 1. Empty map - a projected permutation.
+  AffineMap map1 = b.getEmptyAffineMap();
+  EXPECT_TRUE(map1.isProjectedPermutation());
+
+  // 2. Contains a symbol - not a projected permutation.
+  AffineMap map2 = AffineMap::get(0, 1, &ctx);
+  EXPECT_FALSE(map2.isProjectedPermutation());
+
+  // 3. The result map is {0} - since zero results are _allowed_, this _is_ a
+  // projected permutation.
+  auto zero = b.getAffineConstantExpr(0);
+  AffineMap map3 = AffineMap::get(1, 0, {zero}, &ctx);
+  EXPECT_TRUE(map3.isProjectedPermutation(/*allowZeroInResults=*/true));
+
+  // 4. The result map is {0} - since zero results are _not allowed_, this _is
+  // not_ a projected permutation.
+  AffineMap map4 = AffineMap::get(1, 0, {zero}, &ctx);
+  EXPECT_FALSE(map4.isProjectedPermutation(/*allowZeroInResults=*/false));
+
+  // 5. The number of results > inputs, not a projected permutation.
+  AffineMap map5 = AffineMap::get(1, 0, {zero, zero}, &ctx);
+  EXPECT_FALSE(map5.isProjectedPermutation(/*allowZeroInResults=*/true));
+
+  // 6. A constant result that's not a {0} - not a projected permutation.
+  auto one = b.getAffineConstantExpr(1);
+  AffineMap map6 = AffineMap::get(1, 0, {one}, &ctx);
+  EXPECT_FALSE(map6.isProjectedPermutation(/*allowZeroInResults=*/true));
+
+  // 7. Not a dim expression - not a projected permutation.
+  auto d0 = b.getAffineDimExpr(0);
+  auto d1 = b.getAffineDimExpr(1);
+
+  auto sum = d0 + d1;
+  AffineMap map7 = AffineMap::get(2, 0, {sum}, &ctx);
+  EXPECT_FALSE(map7.isProjectedPermutation());
+}

>From f06eb8120de3ee78be3bd0786bd572db77c84dfc Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 5 Nov 2024 11:50:27 +0000
Subject: [PATCH 2/2] fixup! [mlir][affine] Add unit tests for
 `isProjectedPermutation`

Fix comments, add more cases
---
 mlir/unittests/IR/AffineMapTest.cpp | 32 +++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/mlir/unittests/IR/AffineMapTest.cpp b/mlir/unittests/IR/AffineMapTest.cpp
index 569f0db8c55c9c..c272e9d5093709 100644
--- a/mlir/unittests/IR/AffineMapTest.cpp
+++ b/mlir/unittests/IR/AffineMapTest.cpp
@@ -26,39 +26,53 @@ TEST(AffineMapTest, isProjectedPermutation) {
   MLIRContext ctx;
   OpBuilder b(&ctx);
 
-  // 1. Empty map - a projected permutation.
+  // 1. Empty map
   AffineMap map1 = b.getEmptyAffineMap();
   EXPECT_TRUE(map1.isProjectedPermutation());
 
-  // 2. Contains a symbol - not a projected permutation.
+  // 2. Map with a symbol
   AffineMap map2 = AffineMap::get(0, 1, &ctx);
   EXPECT_FALSE(map2.isProjectedPermutation());
 
-  // 3. The result map is {0} - since zero results are _allowed_, this _is_ a
-  // projected permutation.
+  // 3. The result map is {0} and zero results are _allowed_.
   auto zero = b.getAffineConstantExpr(0);
   AffineMap map3 = AffineMap::get(1, 0, {zero}, &ctx);
   EXPECT_TRUE(map3.isProjectedPermutation(/*allowZeroInResults=*/true));
 
-  // 4. The result map is {0} - since zero results are _not allowed_, this _is
-  // not_ a projected permutation.
+  // 4. The result map is {0} and zero results are _not allowed_
   AffineMap map4 = AffineMap::get(1, 0, {zero}, &ctx);
   EXPECT_FALSE(map4.isProjectedPermutation(/*allowZeroInResults=*/false));
 
-  // 5. The number of results > inputs, not a projected permutation.
+  // 5. The number of results > inputs
   AffineMap map5 = AffineMap::get(1, 0, {zero, zero}, &ctx);
   EXPECT_FALSE(map5.isProjectedPermutation(/*allowZeroInResults=*/true));
 
-  // 6. A constant result that's not a {0} - not a projected permutation.
+  // 6. A constant result that's not a {0}
   auto one = b.getAffineConstantExpr(1);
   AffineMap map6 = AffineMap::get(1, 0, {one}, &ctx);
   EXPECT_FALSE(map6.isProjectedPermutation(/*allowZeroInResults=*/true));
 
-  // 7. Not a dim expression - not a projected permutation.
+  // 7. Not a dim expression
   auto d0 = b.getAffineDimExpr(0);
   auto d1 = b.getAffineDimExpr(1);
 
   auto sum = d0 + d1;
   AffineMap map7 = AffineMap::get(2, 0, {sum}, &ctx);
   EXPECT_FALSE(map7.isProjectedPermutation());
+
+  // 8. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d0, d1, d2, d4)
+  auto d2 = b.getAffineDimExpr(2);
+  auto d3 = b.getAffineDimExpr(3);
+  auto d4 = b.getAffineDimExpr(4);
+  auto d5 = b.getAffineDimExpr(5);
+  AffineMap map8 = AffineMap::get(5, 0, {d5, d3, d0, d1, d2, d4}, &ctx);
+  EXPECT_TRUE(map8.isProjectedPermutation());
+
+  // 9. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d0 + d1, d2, d4)
+  AffineMap map9 = AffineMap::get(5, 0, {d5, d3, sum, d2, d4}, &ctx);
+  EXPECT_FALSE(map9.isProjectedPermutation());
+
+  // 10. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d2, d4)
+  AffineMap map10 = AffineMap::get(5, 0, {d5, d3, d2, d4}, &ctx);
+  EXPECT_TRUE(map10.isProjectedPermutation());
 }