1 | //===- AffineMapTest.cpp - unit tests for affine map API ------------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | #include "mlir/IR/AffineMap.h" |
10 | #include "mlir/IR/Builders.h" |
11 | #include "gtest/gtest.h" |
12 | |
13 | using namespace mlir; |
14 | |
15 | // Test AffineMap replace API for the zero result case. |
16 | TEST(AffineMapTest, inferMapFromAffineExprs) { |
17 | MLIRContext ctx; |
18 | OpBuilder b(&ctx); |
19 | AffineMap map = b.getEmptyAffineMap(); |
20 | DenseMap<AffineExpr, AffineExpr> replacements; |
21 | map.replace(map: replacements); |
22 | EXPECT_EQ(map, map); |
23 | } |
24 | |
25 | TEST(AffineMapTest, isProjectedPermutation) { |
26 | MLIRContext ctx; |
27 | OpBuilder b(&ctx); |
28 | |
29 | // 1. Empty map |
30 | AffineMap map1 = b.getEmptyAffineMap(); |
31 | EXPECT_TRUE(map1.isProjectedPermutation()); |
32 | |
33 | // 2. Map with a symbol |
34 | AffineMap map2 = AffineMap::get(dimCount: 0, symbolCount: 1, context: &ctx); |
35 | EXPECT_FALSE(map2.isProjectedPermutation()); |
36 | |
37 | // 3. The result map is {0} and zero results are _allowed_. |
38 | auto zero = b.getAffineConstantExpr(constant: 0); |
39 | AffineMap map3 = AffineMap::get(dimCount: 1, symbolCount: 0, results: {zero}, context: &ctx); |
40 | EXPECT_TRUE(map3.isProjectedPermutation(/*allowZeroInResults=*/true)); |
41 | |
42 | // 4. The result map is {0} and zero results are _not allowed_ |
43 | AffineMap map4 = AffineMap::get(dimCount: 1, symbolCount: 0, results: {zero}, context: &ctx); |
44 | EXPECT_FALSE(map4.isProjectedPermutation(/*allowZeroInResults=*/false)); |
45 | |
46 | // 5. The number of results > inputs |
47 | AffineMap map5 = AffineMap::get(dimCount: 1, symbolCount: 0, results: {zero, zero}, context: &ctx); |
48 | EXPECT_FALSE(map5.isProjectedPermutation(/*allowZeroInResults=*/true)); |
49 | |
50 | // 6. A constant result that's not a {0} |
51 | auto one = b.getAffineConstantExpr(constant: 1); |
52 | AffineMap map6 = AffineMap::get(dimCount: 1, symbolCount: 0, results: {one}, context: &ctx); |
53 | EXPECT_FALSE(map6.isProjectedPermutation(/*allowZeroInResults=*/true)); |
54 | |
55 | // 7. Not a dim expression |
56 | auto d0 = b.getAffineDimExpr(position: 0); |
57 | auto d1 = b.getAffineDimExpr(position: 1); |
58 | |
59 | auto sum = d0 + d1; |
60 | AffineMap map7 = AffineMap::get(dimCount: 2, symbolCount: 0, results: {sum}, context: &ctx); |
61 | EXPECT_FALSE(map7.isProjectedPermutation()); |
62 | |
63 | // 8. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d0, d1, d2, d4) |
64 | auto d2 = b.getAffineDimExpr(position: 2); |
65 | auto d3 = b.getAffineDimExpr(position: 3); |
66 | auto d4 = b.getAffineDimExpr(position: 4); |
67 | auto d5 = b.getAffineDimExpr(position: 5); |
68 | AffineMap map8 = AffineMap::get(dimCount: 6, symbolCount: 0, results: {d5, d3, d0, d1, d2, d4}, context: &ctx); |
69 | EXPECT_TRUE(map8.isProjectedPermutation()); |
70 | |
71 | // 9. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d0 + d1, d2, d4) |
72 | AffineMap map9 = AffineMap::get(dimCount: 6, symbolCount: 0, results: {d5, d3, sum, d2, d4}, context: &ctx); |
73 | EXPECT_FALSE(map9.isProjectedPermutation()); |
74 | |
75 | // 10. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d2, d4) |
76 | AffineMap map10 = AffineMap::get(dimCount: 6, symbolCount: 0, results: {d5, d3, d2, d4}, context: &ctx); |
77 | EXPECT_TRUE(map10.isProjectedPermutation()); |
78 | } |
79 | |
80 | TEST(AffineMapTest, getInversePermutation) { |
81 | MLIRContext ctx; |
82 | OpBuilder b(&ctx); |
83 | |
84 | // 0. Empty map |
85 | AffineMap map0 = AffineMap::get(dimCount: 0, symbolCount: 0, results: {}, context: &ctx); |
86 | AffineMap inverseMap0 = inversePermutation(map: map0); |
87 | EXPECT_TRUE(inverseMap0.isEmpty()); |
88 | |
89 | auto d0 = b.getAffineDimExpr(position: 0); |
90 | auto d1 = b.getAffineDimExpr(position: 1); |
91 | auto d2 = b.getAffineDimExpr(position: 2); |
92 | |
93 | // 1. (d0, d1, d2) -> (d1, d1, d0, d2, d1, d2, d1, d0) |
94 | AffineMap map1 = AffineMap::get(dimCount: 3, symbolCount: 0, results: {d1, d1, d0, d2, d1, d2, d1, d0}, context: &ctx); |
95 | // (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3) |
96 | AffineMap inverseMap1 = inversePermutation(map: map1); |
97 | auto resultsInv1 = inverseMap1.getResults(); |
98 | EXPECT_EQ(resultsInv1.size(), 3UL); |
99 | |
100 | // Expect (d2, d0, d3) |
101 | SmallVector<unsigned> expected = {2, 0, 3}; |
102 | for (auto [idx, res] : llvm::enumerate(First&: resultsInv1)) { |
103 | AffineDimExpr expr = llvm::dyn_cast<AffineDimExpr>(Val: res); |
104 | EXPECT_TRUE(expr && expr.getPosition() == expected[idx]); |
105 | } |
106 | |
107 | // 2. (d0, d1, d2) -> (d1, d0 + d1, d0, d2, d1, d2, d1, d0) |
108 | auto sum = d0 + d1; |
109 | AffineMap map2 = |
110 | AffineMap::get(dimCount: 3, symbolCount: 0, results: {d1, sum, d0, d2, d1, d2, d1, d0}, context: &ctx); |
111 | // (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3) |
112 | AffineMap inverseMap2 = inversePermutation(map: map2); |
113 | auto resultsInv2 = inverseMap2.getResults(); |
114 | EXPECT_EQ(resultsInv2.size(), 3UL); |
115 | |
116 | // Expect (d2, d0, d3) |
117 | expected = {2, 0, 3}; |
118 | for (auto [idx, res] : llvm::enumerate(First&: resultsInv2)) { |
119 | AffineDimExpr expr = llvm::dyn_cast<AffineDimExpr>(Val: res); |
120 | EXPECT_TRUE(expr && expr.getPosition() == expected[idx]); |
121 | } |
122 | } |
123 | |