1 //===- PWMAFunctionTest.cpp - Tests for PWMAFunction ----------------------===//
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 // This file contains tests for PWMAFunction.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "./Utils.h"
14 
15 #include "mlir/Analysis/Presburger/PWMAFunction.h"
16 #include "mlir/Analysis/Presburger/PresburgerSet.h"
17 #include "mlir/IR/MLIRContext.h"
18 
19 #include <gmock/gmock.h>
20 #include <gtest/gtest.h>
21 
22 namespace mlir {
23 using testing::ElementsAre;
24 
25 static Matrix makeMatrix(unsigned numRow, unsigned numColumns,
26                          ArrayRef<SmallVector<int64_t, 8>> matrix) {
27   Matrix results(numRow, numColumns);
28   assert(matrix.size() == numRow);
29   for (unsigned i = 0; i < numRow; ++i) {
30     assert(matrix[i].size() == numColumns &&
31            "Output expression has incorrect dimensionality!");
32     for (unsigned j = 0; j < numColumns; ++j)
33       results(i, j) = matrix[i][j];
34   }
35   return results;
36 }
37 
38 /// Construct a PWMAFunction given the dimensionalities and an array describing
39 /// the list of pieces. Each piece is given by a string describing the domain
40 /// and a 2D array that represents the output.
41 static PWMAFunction parsePWMAF(
42     unsigned numInputs, unsigned numOutputs,
43     ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
44         data,
45     unsigned numSymbols = 0) {
46   PWMAFunction result(numInputs - numSymbols, numSymbols, numOutputs);
47   for (const auto &pair : data) {
48     IntegerPolyhedron domain = parsePoly(pair.first);
49     result.addPiece(
50         domain, makeMatrix(numOutputs, domain.getNumIds() + 1, pair.second));
51   }
52   return result;
53 }
54 
55 TEST(PWAFunctionTest, isEqual) {
56   // The output expressions are different but it doesn't matter because they are
57   // equal in this domain.
58   PWMAFunction idAtZeros = parsePWMAF(
59       /*numInputs=*/2, /*numOutputs=*/2,
60       {
61           {"(x, y) : (y == 0)", {{1, 0, 0}, {0, 1, 0}}},             // (x, y).
62           {"(x, y) : (y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 1, 0}}}, // (x, y).
63           {"(x, y) : (-y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 1, 0}}} // (x, y).
64       });
65   PWMAFunction idAtZeros2 = parsePWMAF(
66       /*numInputs=*/2, /*numOutputs=*/2,
67       {
68           {"(x, y) : (y == 0)", {{1, 0, 0}, {0, 20, 0}}}, // (x, 20y).
69           {"(x, y) : (y - 1 >= 0, x == 0)", {{30, 0, 0}, {0, 1, 0}}}, //(30x, y)
70           {"(x, y) : (-y - 1 > =0, x == 0)", {{30, 0, 0}, {0, 1, 0}}} //(30x, y)
71       });
72   EXPECT_TRUE(idAtZeros.isEqual(idAtZeros2));
73 
74   PWMAFunction notIdAtZeros = parsePWMAF(
75       /*numInputs=*/2, /*numOutputs=*/2,
76       {
77           {"(x, y) : (y == 0)", {{1, 0, 0}, {0, 1, 0}}},              // (x, y).
78           {"(x, y) : (y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 2, 0}}},  // (x, 2y)
79           {"(x, y) : (-y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 2, 0}}}, // (x, 2y)
80       });
81   EXPECT_FALSE(idAtZeros.isEqual(notIdAtZeros));
82 
83   // These match at their intersection but one has a bigger domain.
84   PWMAFunction idNoNegNegQuadrant = parsePWMAF(
85       /*numInputs=*/2, /*numOutputs=*/2,
86       {
87           {"(x, y) : (x >= 0)", {{1, 0, 0}, {0, 1, 0}}},             // (x, y).
88           {"(x, y) : (-x - 1 >= 0, y >= 0)", {{1, 0, 0}, {0, 1, 0}}} // (x, y).
89       });
90   PWMAFunction idOnlyPosX =
91       parsePWMAF(/*numInputs=*/2, /*numOutputs=*/2,
92                  {
93                      {"(x, y) : (x >= 0)", {{1, 0, 0}, {0, 1, 0}}}, // (x, y).
94                  });
95   EXPECT_FALSE(idNoNegNegQuadrant.isEqual(idOnlyPosX));
96 
97   // Different representations of the same domain.
98   PWMAFunction sumPlusOne = parsePWMAF(
99       /*numInputs=*/2, /*numOutputs=*/1,
100       {
101           {"(x, y) : (x >= 0)", {{1, 1, 1}}},                   // x + y + 1.
102           {"(x, y) : (-x - 1 >= 0, -y - 1 >= 0)", {{1, 1, 1}}}, // x + y + 1.
103           {"(x, y) : (-x - 1 >= 0, y >= 0)", {{1, 1, 1}}}       // x + y + 1.
104       });
105   PWMAFunction sumPlusOne2 =
106       parsePWMAF(/*numInputs=*/2, /*numOutputs=*/1,
107                  {
108                      {"(x, y) : ()", {{1, 1, 1}}}, // x + y + 1.
109                  });
110   EXPECT_TRUE(sumPlusOne.isEqual(sumPlusOne2));
111 
112   // Functions with zero input dimensions.
113   PWMAFunction noInputs1 = parsePWMAF(/*numInputs=*/0, /*numOutputs=*/1,
114                                       {
115                                           {"() : ()", {{1}}}, // 1.
116                                       });
117   PWMAFunction noInputs2 = parsePWMAF(/*numInputs=*/0, /*numOutputs=*/1,
118                                       {
119                                           {"() : ()", {{2}}}, // 1.
120                                       });
121   EXPECT_TRUE(noInputs1.isEqual(noInputs1));
122   EXPECT_FALSE(noInputs1.isEqual(noInputs2));
123 
124   // Mismatched dimensionalities.
125   EXPECT_FALSE(noInputs1.isEqual(sumPlusOne));
126   EXPECT_FALSE(idOnlyPosX.isEqual(sumPlusOne));
127 
128   // Divisions.
129   // Domain is only multiples of 6; x = 6k for some k.
130   // x + 4(x/2) + 4(x/3) == 26k.
131   PWMAFunction mul2AndMul3 = parsePWMAF(
132       /*numInputs=*/1, /*numOutputs=*/1,
133       {
134           {"(x) : (x - 2*(x floordiv 2) == 0, x - 3*(x floordiv 3) == 0)",
135            {{1, 4, 4, 0}}}, // x + 4(x/2) + 4(x/3).
136       });
137   PWMAFunction mul6 = parsePWMAF(
138       /*numInputs=*/1, /*numOutputs=*/1,
139       {
140           {"(x) : (x - 6*(x floordiv 6) == 0)", {{0, 26, 0}}}, // 26(x/6).
141       });
142   EXPECT_TRUE(mul2AndMul3.isEqual(mul6));
143 
144   PWMAFunction mul6diff = parsePWMAF(
145       /*numInputs=*/1, /*numOutputs=*/1,
146       {
147           {"(x) : (x - 5*(x floordiv 5) == 0)", {{0, 52, 0}}}, // 52(x/6).
148       });
149   EXPECT_FALSE(mul2AndMul3.isEqual(mul6diff));
150 
151   PWMAFunction mul5 = parsePWMAF(
152       /*numInputs=*/1, /*numOutputs=*/1,
153       {
154           {"(x) : (x - 5*(x floordiv 5) == 0)", {{0, 26, 0}}}, // 26(x/5).
155       });
156   EXPECT_FALSE(mul2AndMul3.isEqual(mul5));
157 }
158 
159 TEST(PWMAFunction, valueAt) {
160   PWMAFunction nonNegPWAF = parsePWMAF(
161       /*numInputs=*/2, /*numOutputs=*/2,
162       {
163           {"(x, y) : (x >= 0)", {{1, 2, 3}, {3, 4, 5}}}, // (x, y).
164           {"(x, y) : (y >= 0, -x - 1 >= 0)", {{-1, 2, 3}, {-3, 4, 5}}} // (x, y)
165       });
166   EXPECT_THAT(*nonNegPWAF.valueAt({2, 3}), ElementsAre(11, 23));
167   EXPECT_THAT(*nonNegPWAF.valueAt({-2, 3}), ElementsAre(11, 23));
168   EXPECT_THAT(*nonNegPWAF.valueAt({2, -3}), ElementsAre(-1, -1));
169   EXPECT_FALSE(nonNegPWAF.valueAt({-2, -3}).hasValue());
170 }
171 
172 } // namespace mlir
173