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