mlir/unittests/IR/AffineMapTest.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-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 }