llvm-project

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
// Test creating AffineExprs using the overloaded binary operators.
19
TEST(AffineExprTest, constructFromBinaryOperators) {
20
  MLIRContext ctx;
21
  OpBuilder b(&ctx);
22

23
  auto d0 = b.getAffineDimExpr(0);
24
  auto d1 = b.getAffineDimExpr(1);
25

26
  auto sum = d0 + d1;
27
  auto difference = d0 - d1;
28
  auto product = d0 * d1;
29
  auto remainder = d0 % d1;
30

31
  ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
32
  ASSERT_EQ(difference.getKind(), AffineExprKind::Add);
33
  ASSERT_EQ(product.getKind(), AffineExprKind::Mul);
34
  ASSERT_EQ(remainder.getKind(), AffineExprKind::Mod);
35
}
36

37
TEST(AffineExprTest, constantFolding) {
38
  MLIRContext ctx;
39
  OpBuilder b(&ctx);
40
  auto cn1 = b.getAffineConstantExpr(-1);
41
  auto c0 = b.getAffineConstantExpr(0);
42
  auto c1 = b.getAffineConstantExpr(1);
43
  auto c2 = b.getAffineConstantExpr(2);
44
  auto c3 = b.getAffineConstantExpr(3);
45
  auto c6 = b.getAffineConstantExpr(6);
46
  auto cmax = b.getAffineConstantExpr(std::numeric_limits<int64_t>::max());
47
  auto cmin = b.getAffineConstantExpr(std::numeric_limits<int64_t>::min());
48

49
  ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::Add, c1, c2), c3);
50
  ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::Mul, c2, c3), c6);
51
  ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::FloorDiv, c3, c2), c1);
52
  ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::CeilDiv, c3, c2), c2);
53

54
  // Test division by zero:
55
  auto c3ceildivc0 = getAffineBinaryOpExpr(AffineExprKind::CeilDiv, c3, c0);
56
  ASSERT_EQ(c3ceildivc0.getKind(), AffineExprKind::CeilDiv);
57

58
  auto c3floordivc0 = getAffineBinaryOpExpr(AffineExprKind::FloorDiv, c3, c0);
59
  ASSERT_EQ(c3floordivc0.getKind(), AffineExprKind::FloorDiv);
60

61
  auto c3modc0 = getAffineBinaryOpExpr(AffineExprKind::Mod, c3, c0);
62
  ASSERT_EQ(c3modc0.getKind(), AffineExprKind::Mod);
63

64
  // Test overflow:
65
  auto cmaxplusc1 = getAffineBinaryOpExpr(AffineExprKind::Add, cmax, c1);
66
  ASSERT_EQ(cmaxplusc1.getKind(), AffineExprKind::Add);
67

68
  auto cmaxtimesc2 = getAffineBinaryOpExpr(AffineExprKind::Mul, cmax, c2);
69
  ASSERT_EQ(cmaxtimesc2.getKind(), AffineExprKind::Mul);
70

71
  auto cminceildivcn1 =
72
      getAffineBinaryOpExpr(AffineExprKind::CeilDiv, cmin, cn1);
73
  ASSERT_EQ(cminceildivcn1.getKind(), AffineExprKind::CeilDiv);
74

75
  auto cminfloordivcn1 =
76
      getAffineBinaryOpExpr(AffineExprKind::FloorDiv, cmin, cn1);
77
  ASSERT_EQ(cminfloordivcn1.getKind(), AffineExprKind::FloorDiv);
78
}
79

80
TEST(AffineExprTest, divisionSimplification) {
81
  MLIRContext ctx;
82
  OpBuilder b(&ctx);
83
  auto cn6 = b.getAffineConstantExpr(-6);
84
  auto c6 = b.getAffineConstantExpr(6);
85
  auto d0 = b.getAffineDimExpr(0);
86
  auto d1 = b.getAffineDimExpr(1);
87

88
  ASSERT_EQ(c6.floorDiv(-1), cn6);
89
  ASSERT_EQ((d0 * 6).floorDiv(2), d0 * 3);
90
  ASSERT_EQ((d0 * 6).floorDiv(4).getKind(), AffineExprKind::FloorDiv);
91
  ASSERT_EQ((d0 * 6).floorDiv(-2), d0 * -3);
92
  ASSERT_EQ((d0 * 6 + d1).floorDiv(2), d0 * 3 + d1.floorDiv(2));
93
  ASSERT_EQ((d0 * 6 + d1).floorDiv(-2), d0 * -3 + d1.floorDiv(-2));
94
  ASSERT_EQ((d0 * 6 + d1).floorDiv(4).getKind(), AffineExprKind::FloorDiv);
95

96
  ASSERT_EQ(c6.ceilDiv(-1), cn6);
97
  ASSERT_EQ((d0 * 6).ceilDiv(2), d0 * 3);
98
  ASSERT_EQ((d0 * 6).ceilDiv(4).getKind(), AffineExprKind::CeilDiv);
99
  ASSERT_EQ((d0 * 6).ceilDiv(-2), d0 * -3);
100
}
101