mlir/unittests/IR/AffineExprTest.cpp

   1 //===- AffineExprTest.cpp - unit tests for affine expression 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 <cstdint>
  10 #include <limits>
  11
  12 #include "mlir/IR/AffineExpr.h"
  13 #include "mlir/IR/Builders.h"
  14 #include "gtest/gtest.h"
  15
  16 using namespace mlir;
  17
  18 static std::string toString(AffineExpr expr) {
  19   std::string s;
  20   llvm::raw_string_ostream ss(s);
  21   ss << expr;
  22   return s;
  23 }
  24
  25 // Test creating AffineExprs using the overloaded binary operators.
  26 TEST(AffineExprTest, constructFromBinaryOperators) {
  27   MLIRContext ctx;
  28   OpBuilder b(&ctx);
  29
  30   auto d0 = b.getAffineDimExpr(0);
  31   auto d1 = b.getAffineDimExpr(1);
  32
  33   auto sum = d0 + d1;
  34   auto difference = d0 - d1;
  35   auto product = d0 * d1;
  36   auto remainder = d0 % d1;
  37
  38   ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
  39   ASSERT_EQ(difference.getKind(), AffineExprKind::Add);
  40   ASSERT_EQ(product.getKind(), AffineExprKind::Mul);
  41   ASSERT_EQ(remainder.getKind(), AffineExprKind::Mod);
  42 }
  43
  44 TEST(AffineExprTest, constantFolding) {
  45   MLIRContext ctx;
  46   OpBuilder b(&ctx);
  47   auto cn1 = b.getAffineConstantExpr(-1);
  48   auto c0 = b.getAffineConstantExpr(0);
  49   auto c1 = b.getAffineConstantExpr(1);
  50   auto c2 = b.getAffineConstantExpr(2);
  51   auto c3 = b.getAffineConstantExpr(3);
  52   auto c6 = b.getAffineConstantExpr(6);
  53   auto cmax = b.getAffineConstantExpr(std::numeric_limits<int64_t>::max());
  54   auto cmin = b.getAffineConstantExpr(std::numeric_limits<int64_t>::min());
  55
  56   ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::Add, c1, c2), c3);
  57   ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::Mul, c2, c3), c6);
  58   ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::FloorDiv, c3, c2), c1);
  59   ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::CeilDiv, c3, c2), c2);
  60
  61   // Test division by zero:
  62   auto c3ceildivc0 = getAffineBinaryOpExpr(AffineExprKind::CeilDiv, c3, c0);
  63   ASSERT_EQ(c3ceildivc0.getKind(), AffineExprKind::CeilDiv);
  64
  65   auto c3floordivc0 = getAffineBinaryOpExpr(AffineExprKind::FloorDiv, c3, c0);
  66   ASSERT_EQ(c3floordivc0.getKind(), AffineExprKind::FloorDiv);
  67
  68   auto c3modc0 = getAffineBinaryOpExpr(AffineExprKind::Mod, c3, c0);
  69   ASSERT_EQ(c3modc0.getKind(), AffineExprKind::Mod);
  70
  71   // Test overflow:
  72   auto cmaxplusc1 = getAffineBinaryOpExpr(AffineExprKind::Add, cmax, c1);
  73   ASSERT_EQ(cmaxplusc1.getKind(), AffineExprKind::Add);
  74
  75   auto cmaxtimesc2 = getAffineBinaryOpExpr(AffineExprKind::Mul, cmax, c2);
  76   ASSERT_EQ(cmaxtimesc2.getKind(), AffineExprKind::Mul);
  77
  78   auto cminceildivcn1 =
  79       getAffineBinaryOpExpr(AffineExprKind::CeilDiv, cmin, cn1);
  80   ASSERT_EQ(cminceildivcn1.getKind(), AffineExprKind::CeilDiv);
  81
  82   auto cminfloordivcn1 =
  83       getAffineBinaryOpExpr(AffineExprKind::FloorDiv, cmin, cn1);
  84   ASSERT_EQ(cminfloordivcn1.getKind(), AffineExprKind::FloorDiv);
  85 }
  86
  87 TEST(AffineExprTest, divisionSimplification) {
  88   MLIRContext ctx;
  89   OpBuilder b(&ctx);
  90   auto cn6 = b.getAffineConstantExpr(-6);
  91   auto c6 = b.getAffineConstantExpr(6);
  92   auto d0 = b.getAffineDimExpr(0);
  93   auto d1 = b.getAffineDimExpr(1);
  94
  95   ASSERT_EQ(c6.floorDiv(-1), cn6);
  96   ASSERT_EQ((d0 * 6).floorDiv(2), d0 * 3);
  97   ASSERT_EQ((d0 * 6).floorDiv(4).getKind(), AffineExprKind::FloorDiv);
  98   ASSERT_EQ((d0 * 6).floorDiv(-2), d0 * -3);
  99   ASSERT_EQ((d0 * 6 + d1).floorDiv(2), d0 * 3 + d1.floorDiv(2));
 100   ASSERT_EQ((d0 * 6 + d1).floorDiv(-2), d0 * -3 + d1.floorDiv(-2));
 101   ASSERT_EQ((d0 * 6 + d1).floorDiv(4).getKind(), AffineExprKind::FloorDiv);
 102
 103   ASSERT_EQ(c6.ceilDiv(-1), cn6);
 104   ASSERT_EQ((d0 * 6).ceilDiv(2), d0 * 3);
 105   ASSERT_EQ((d0 * 6).ceilDiv(4).getKind(), AffineExprKind::CeilDiv);
 106   ASSERT_EQ((d0 * 6).ceilDiv(-2), d0 * -3);
 107 }
 108
 109 TEST(AffineExprTest, modSimplificationRegression) {
 110   MLIRContext ctx;
 111   OpBuilder b(&ctx);
 112   auto d0 = b.getAffineDimExpr(0);
 113   auto sum = d0 + d0.floorDiv(3).floorDiv(-3);
 114   ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
 115 }
 116
 117 TEST(AffineExprTest, divisorOfNegativeFloorDiv) {
 118   MLIRContext ctx;
 119   OpBuilder b(&ctx);
 120   ASSERT_EQ(b.getAffineDimExpr(0).floorDiv(-1).getLargestKnownDivisor(), 1);
 121 }
 122
 123 TEST(AffineExprTest, d0PlusD0FloorDivNeg2) {
 124   // Regression test for a bug where this was rewritten to d0 mod -2. We do not
 125   // support a negative RHS for mod in LowerAffinePass.
 126   MLIRContext ctx;
 127   OpBuilder b(&ctx);
 128   auto d0 = b.getAffineDimExpr(0);
 129   auto sum = d0 + d0.floorDiv(-2) * 2;
 130   ASSERT_EQ(toString(sum), "d0 + (d0 floordiv -2) * 2");
 131 }