xref: /llvm-project/mlir/unittests/IR/AffineMapTest.cpp (revision 6f5e5b630559f2d17bdccfab5dff3a97ac0f8c66)
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(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(0, 1, &ctx);
35   EXPECT_FALSE(map2.isProjectedPermutation());
36 
37   // 3. The result map is {0} and zero results are _allowed_.
38   auto zero = b.getAffineConstantExpr(0);
39   AffineMap map3 = AffineMap::get(1, 0, {zero}, &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(1, 0, {zero}, &ctx);
44   EXPECT_FALSE(map4.isProjectedPermutation(/*allowZeroInResults=*/false));
45 
46   // 5. The number of results > inputs
47   AffineMap map5 = AffineMap::get(1, 0, {zero, zero}, &ctx);
48   EXPECT_FALSE(map5.isProjectedPermutation(/*allowZeroInResults=*/true));
49 
50   // 6. A constant result that's not a {0}
51   auto one = b.getAffineConstantExpr(1);
52   AffineMap map6 = AffineMap::get(1, 0, {one}, &ctx);
53   EXPECT_FALSE(map6.isProjectedPermutation(/*allowZeroInResults=*/true));
54 
55   // 7. Not a dim expression
56   auto d0 = b.getAffineDimExpr(0);
57   auto d1 = b.getAffineDimExpr(1);
58 
59   auto sum = d0 + d1;
60   AffineMap map7 = AffineMap::get(2, 0, {sum}, &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(2);
65   auto d3 = b.getAffineDimExpr(3);
66   auto d4 = b.getAffineDimExpr(4);
67   auto d5 = b.getAffineDimExpr(5);
68   AffineMap map8 = AffineMap::get(6, 0, {d5, d3, d0, d1, d2, d4}, &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(6, 0, {d5, d3, sum, d2, d4}, &ctx);
73   EXPECT_FALSE(map9.isProjectedPermutation());
74 
75   // 10. (d0, d1, d2, d3, d4, d5) ->(d5, d3, d2, d4)
76   AffineMap map10 = AffineMap::get(6, 0, {d5, d3, d2, d4}, &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(0, 0, {}, &ctx);
86   AffineMap inverseMap0 = inversePermutation(map0);
87   EXPECT_TRUE(inverseMap0.isEmpty());
88 
89   auto d0 = b.getAffineDimExpr(0);
90   auto d1 = b.getAffineDimExpr(1);
91   auto d2 = b.getAffineDimExpr(2);
92 
93   // 1.   (d0, d1, d2) -> (d1, d1, d0, d2, d1, d2, d1, d0)
94   AffineMap map1 = AffineMap::get(3, 0, {d1, d1, d0, d2, d1, d2, d1, d0}, &ctx);
95   //      (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3)
96   AffineMap inverseMap1 = inversePermutation(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(resultsInv1)) {
103     AffineDimExpr expr = llvm::dyn_cast<AffineDimExpr>(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(3, 0, {d1, sum, d0, d2, d1, d2, d1, d0}, &ctx);
111   //      (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3)
112   AffineMap inverseMap2 = inversePermutation(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(resultsInv2)) {
119     AffineDimExpr expr = llvm::dyn_cast<AffineDimExpr>(res);
120     EXPECT_TRUE(expr && expr.getPosition() == expected[idx]);
121   }
122 }
123