123 lines · cpp
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-exception6//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(replacements);22 EXPECT_EQ(map, map);23}24 25TEST(AffineMapTest, isProjectedPermutation) {26 MLIRContext ctx;27 OpBuilder b(&ctx);28 29 // 1. Empty map30 AffineMap map1 = b.getEmptyAffineMap();31 EXPECT_TRUE(map1.isProjectedPermutation());32 33 // 2. Map with a symbol34 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 > inputs47 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 expression56 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 80TEST(AffineMapTest, getInversePermutation) {81 MLIRContext ctx;82 OpBuilder b(&ctx);83 84 // 0. Empty map85 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