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
13using namespace mlir;
14
15// Test AffineMap replace API for the zero result case.
16TEST(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
25TEST(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
80TEST(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

source code of mlir/unittests/IR/AffineMapTest.cpp