xref: /llvm-project/llvm/unittests/CodeGen/MLRegAllocDevelopmentFeatures.cpp (revision bb3f5e1fed7c6ba733b7f273e93f5d3930976185) 
1 //===- MLRegAllocDevelopmentFeatures.cpp - test dev MLRegAlloc features ---===//
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 "../../lib/CodeGen/MLRegAllocEvictAdvisor.h"
10 #include "llvm/Analysis/NoInferenceModelRunner.h"
11 #include "llvm/CodeGen/CodeGenTargetMachineImpl.h"
12 #include "llvm/CodeGen/MachineBasicBlock.h"
13 #include "llvm/CodeGen/MachineFunction.h"
14 #include "llvm/CodeGen/MachineModuleInfo.h"
15 #include "llvm/CodeGen/SlotIndexes.h"
16 #include "llvm/CodeGen/TargetFrameLowering.h"
17 #include "llvm/CodeGen/TargetInstrInfo.h"
18 #include "llvm/CodeGen/TargetLowering.h"
19 #include "llvm/IR/LLVMContext.h"
20 #include "llvm/IR/Module.h"
21 #include "llvm/MC/TargetRegistry.h"
22 #include "llvm/Support/Allocator.h"
23 #include "llvm/Support/CodeGen.h"
24 #include "llvm/Support/TargetSelect.h"
25 #include "llvm/Target/TargetOptions.h"
26 #include "llvm/TargetParser/Triple.h"
27 #include "gmock/gmock.h"
28 #include "gtest/gtest.h"
29 
30 #include <string>
31 #include <vector>
32 
33 using namespace llvm;
34 using testing::ContainerEq;
35 using testing::Test;
36 
37 namespace {
38 
39 #include "MFCommon.inc"
40 
41 struct LRPosInfoIndexes {
42   size_t StartIndex;
43   size_t EndIndex;
44   size_t PhysReg;
45 };
46 
47 class RegAllocDevelopmentFeaturesTest : public ::Test {
48 protected:
49   SmallVector<LRStartEndInfo>
50   setupOverlapProblem(const SmallVectorImpl<LRPosInfoIndexes> &Segments,
51                       simple_ilist<IndexListEntry> &IndexList) {
52     SmallVector<LRStartEndInfo> PositionsToReturn;
53     PositionsToReturn.reserve(Segments.size());
54     for (auto CurrentPosIndexInfo : Segments) {
55       LRStartEndInfo CurrentPosInfo = {};
56       CurrentPosInfo.Pos = CurrentPosIndexInfo.PhysReg;
57       PositionsToReturn.push_back(CurrentPosInfo);
58     }
59     size_t CurrentSegmentIndex = 0;
60     size_t CurrentIndex = 0;
61     while (CurrentSegmentIndex < Segments.size()) {
62       auto *CurrentLEMem = static_cast<IndexListEntry *>(
63           Allocator.Allocate(sizeof(IndexListEntry), alignof(IndexListEntry)));
64       auto *CurrentListEntry =
65           new (CurrentLEMem) IndexListEntry(nullptr, CurrentIndex);
66       IndexList.push_back(*CurrentListEntry);
67       for (size_t CurrentPosInfoIndex = 0;
68            CurrentPosInfoIndex < Segments.size(); ++CurrentPosInfoIndex) {
69         if ((CurrentIndex / SlotIndex::InstrDist) ==
70             Segments[CurrentPosInfoIndex].StartIndex) {
71           PositionsToReturn[CurrentPosInfoIndex].Begin =
72               SlotIndex(CurrentListEntry, 0);
73         } else if ((CurrentIndex / SlotIndex::InstrDist) ==
74                    Segments[CurrentPosInfoIndex].EndIndex) {
75           PositionsToReturn[CurrentPosInfoIndex].End =
76               SlotIndex(CurrentListEntry, 0);
77           ++CurrentSegmentIndex;
78         }
79       }
80       CurrentIndex += SlotIndex::InstrDist;
81     }
82     return PositionsToReturn;
83   }
84 
85   NoInferenceModelRunner setupModelRunner() {
86     const std::vector<TensorSpec> Inputs{
87         TensorSpec::createSpec<int64_t>("instructions", InstructionsShape),
88         TensorSpec::createSpec<int64_t>("instructions_mapping",
89                                         InstructionsMappingShape),
90         TensorSpec::createSpec<float>("mbb_frequencies", MBBFrequencyShape),
91         TensorSpec::createSpec<int64_t>("mbb_mapping", InstructionsShape)};
92     LLVMContext Ctx;
93     return NoInferenceModelRunner(Ctx, Inputs);
94   }
95 
96   std::vector<int64_t>
97   getExpectedMappingMatrix(SmallVectorImpl<LRPosInfoIndexes> &OverlapSetup) {
98     std::vector<int64_t> ExpectedMappingMatrix(
99         NumberOfInterferences * ModelMaxSupportedInstructionCount, 0);
100     for (auto NewSegment : OverlapSetup) {
101       for (size_t CurrentIndex = NewSegment.StartIndex;
102            CurrentIndex <= NewSegment.EndIndex; ++CurrentIndex) {
103         ExpectedMappingMatrix[NewSegment.PhysReg *
104                                   ModelMaxSupportedInstructionCount +
105                               CurrentIndex] = 1;
106       }
107     }
108     return ExpectedMappingMatrix;
109   }
110 
111   void runOverlapTest(SmallVectorImpl<LRPosInfoIndexes> &OverlapSetup) {
112     simple_ilist<IndexListEntry> IndexList;
113     auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList);
114     NoInferenceModelRunner ModelRunner = setupModelRunner();
115     size_t MaxIndex = 0;
116     for (size_t CurrentOverlap = 0; CurrentOverlap < OverlapSetup.size();
117          ++CurrentOverlap) {
118       if (OverlapSetup[CurrentOverlap].EndIndex >
119           OverlapSetup[MaxIndex].EndIndex) {
120         MaxIndex = CurrentOverlap;
121       }
122     }
123     SlotIndex LastIndex = OverlapProblem[MaxIndex].End;
124     extractInstructionFeatures(
125         OverlapProblem, &ModelRunner,
126         [](SlotIndex InputSlot) -> int { return 0; },
127         [](SlotIndex InputSlot) -> float { return 0.0f; },
128         [](SlotIndex InputSlot) -> MachineBasicBlock * { return nullptr; }, 0,
129         1, 2, 3, LastIndex);
130     std::vector<int64_t> MappingMatrix(
131         ModelRunner.getTensor<int64_t>(1),
132         ModelRunner.getTensor<int64_t>(1) +
133             NumberOfInterferences * ModelMaxSupportedInstructionCount);
134     ASSERT_THAT(MappingMatrix,
135                 ContainerEq(getExpectedMappingMatrix(OverlapSetup)));
136     IndexList.clear();
137   }
138 
139   BumpPtrAllocator Allocator;
140 };
141 
142 // meta tests to ensure that test setup works correctly
143 
144 TEST_F(RegAllocDevelopmentFeaturesTest,
145        MetaOverlapInstructionDistancesAreCorrect) {
146   SmallVector<LRPosInfoIndexes, 2> OverlapSetup;
147   OverlapSetup.push_back({0, 5, 0});
148   OverlapSetup.push_back({5, 10, 0});
149   simple_ilist<IndexListEntry> IndexList;
150   auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList);
151   ASSERT_EQ(OverlapProblem[0].End.distance(OverlapProblem[1].End),
152             5 * SlotIndex::InstrDist);
153   ASSERT_EQ(OverlapProblem[0].End.distance(OverlapProblem[1].Begin), 0);
154 }
155 
156 TEST_F(RegAllocDevelopmentFeaturesTest, MetaSlotIndicesAreValid) {
157   SmallVector<LRPosInfoIndexes, 1> OverlapSetup;
158   OverlapSetup.push_back({0, 10, 0});
159   simple_ilist<IndexListEntry> IndexList;
160   auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList);
161   ASSERT_TRUE(OverlapProblem[0].Begin.isValid());
162   ASSERT_TRUE(OverlapProblem[0].End.isValid());
163 }
164 
165 // Testing of feature extraction for per-instruction features
166 
167 TEST_F(RegAllocDevelopmentFeaturesTest, InstructionOpcodesAreCorrect) {
168   SmallVector<LRPosInfoIndexes, 1> OverlapSetup;
169   OverlapSetup.push_back({0, ModelMaxSupportedInstructionCount - 1, 0});
170   simple_ilist<IndexListEntry> IndexList;
171   auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList);
172   NoInferenceModelRunner ModelRunner = setupModelRunner();
173   SlotIndex LastIndex = OverlapProblem[0].End;
174   SlotIndex FirstIndex = OverlapProblem[0].Begin;
175   extractInstructionFeatures(
176       OverlapProblem, &ModelRunner,
177       [FirstIndex](SlotIndex InputSlot) -> int {
178         return FirstIndex.distance(InputSlot) / SlotIndex::InstrDist;
179       },
180       [](SlotIndex InputSlot) -> float { return 0.0f; },
181       [](SlotIndex InputSlot) -> MachineBasicBlock * { return nullptr; }, 0, 1,
182       2, 3, LastIndex);
183   for (size_t CurrentInstructionIndex = 0;
184        CurrentInstructionIndex < ModelMaxSupportedInstructionCount;
185        ++CurrentInstructionIndex) {
186     ASSERT_EQ(
187         (size_t)ModelRunner.getTensor<int64_t>(0)[CurrentInstructionIndex],
188         CurrentInstructionIndex);
189   }
190 }
191 
192 TEST_F(RegAllocDevelopmentFeaturesTest, FullOverlap) {
193   SmallVector<LRPosInfoIndexes, 2> OverlapSetup;
194   OverlapSetup.push_back({0, ModelMaxSupportedInstructionCount - 1, 0});
195   OverlapSetup.push_back({0, ModelMaxSupportedInstructionCount - 1, 1});
196   runOverlapTest(OverlapSetup);
197 }
198 
199 TEST_F(RegAllocDevelopmentFeaturesTest, PartialOverlap) {
200   SmallVector<LRPosInfoIndexes, 2> OverlapSetup;
201   OverlapSetup.push_back({0, 20, 0});
202   OverlapSetup.push_back({15, 30, 1});
203   runOverlapTest(OverlapSetup);
204 }
205 
206 TEST_F(RegAllocDevelopmentFeaturesTest, PartialOverlapOpposite) {
207   SmallVector<LRPosInfoIndexes, 2> OverlapSetup;
208   OverlapSetup.push_back({15, 30, 1});
209   OverlapSetup.push_back({0, 20, 0});
210   runOverlapTest(OverlapSetup);
211 }
212 
213 TEST_F(RegAllocDevelopmentFeaturesTest, InternalOverlap) {
214   SmallVector<LRPosInfoIndexes, 2> OverlapSetup;
215   OverlapSetup.push_back({0, 30, 0});
216   OverlapSetup.push_back({10, 20, 1});
217   runOverlapTest(OverlapSetup);
218 }
219 
220 TEST_F(RegAllocDevelopmentFeaturesTest, TripleInternalOverlap) {
221   SmallVector<LRPosInfoIndexes, 3> OverlapSetup;
222   OverlapSetup.push_back({0, 30, 0});
223   OverlapSetup.push_back({10, 25, 1});
224   OverlapSetup.push_back({15, 20, 2});
225   runOverlapTest(OverlapSetup);
226 }
227 
228 TEST_F(RegAllocDevelopmentFeaturesTest, InternalMultiOverlap) {
229   SmallVector<LRPosInfoIndexes, 3> OverlapSetup;
230   OverlapSetup.push_back({0, 45, 0});
231   OverlapSetup.push_back({30, 40, 1});
232   OverlapSetup.push_back({35, 60, 2});
233   runOverlapTest(OverlapSetup);
234 }
235 
236 TEST_F(RegAllocDevelopmentFeaturesTest, SingleMBBTest) {
237   NoInferenceModelRunner ModelRunner = setupModelRunner();
238   SlotIndex CurrentIndex;
239   // set index to 1 so we can ensure that the mapping actually get set
240   std::map<MachineBasicBlock *, size_t> VisitedMBBs = {{nullptr, 1}};
241   extractMBBFrequency(
242       CurrentIndex, 0, VisitedMBBs,
243       [](SlotIndex InputSlot) -> float { return 1.0f; }, nullptr, &ModelRunner,
244       2, 3);
245   ASSERT_FLOAT_EQ(ModelRunner.getTensor<float>(2)[1], 1.0f);
246   ASSERT_EQ(ModelRunner.getTensor<int64_t>(3)[0], 1);
247 }
248 
249 TEST_F(RegAllocDevelopmentFeaturesTest, MBBFullTruncated) {
250   SmallVector<LRPosInfoIndexes, 1> OverlapSetup;
251   OverlapSetup.push_back({0, ModelMaxSupportedInstructionCount - 1, 0});
252   simple_ilist<IndexListEntry> IndexList;
253   auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList);
254   NoInferenceModelRunner ModelRunner = setupModelRunner();
255   SlotIndex LastIndex = OverlapProblem[0].End;
256   SlotIndex FirstIndex = OverlapProblem[0].Begin;
257 
258   LLVMContext Ctx;
259   Module Mod("Module", Ctx);
260   auto MF = createMachineFunction(Ctx, Mod);
261   std::array<MachineBasicBlock *, ModelMaxSupportedInstructionCount>
262       MBBsForTest;
263   for (size_t I = 0; I < ModelMaxSupportedInstructionCount; ++I) {
264     MBBsForTest[I] = MF->CreateMachineBasicBlock();
265   }
266 
267   extractInstructionFeatures(
268       OverlapProblem, &ModelRunner,
269       [](SlotIndex InputSlot) -> int { return 0; },
270       [FirstIndex](SlotIndex InputSlot) -> float {
271         return static_cast<float>(FirstIndex.distance(InputSlot) /
272                                   SlotIndex::InstrDist);
273       },
274       [FirstIndex, MBBsForTest](SlotIndex InputSlot) -> MachineBasicBlock * {
275         return MBBsForTest[FirstIndex.distance(InputSlot) /
276                            SlotIndex::InstrDist];
277       },
278       0, 1, 2, 3, LastIndex);
279   for (size_t MBBIndex = 0; MBBIndex < ModelMaxSupportedMBBCount; ++MBBIndex) {
280     ASSERT_FLOAT_EQ(ModelRunner.getTensor<float>(2)[MBBIndex],
281                     static_cast<float>(MBBIndex));
282     ASSERT_EQ(ModelRunner.getTensor<int64_t>(3)[MBBIndex],
283               static_cast<int64_t>(MBBIndex));
284   }
285   // the rest of the mapping values should be zero (truncated to 100 MBBs)
286   for (size_t MBBIndex = ModelMaxSupportedMBBCount;
287        MBBIndex < ModelMaxSupportedInstructionCount; ++MBBIndex) {
288     ASSERT_EQ(ModelRunner.getTensor<int64_t>(3)[MBBIndex],
289               static_cast<int64_t>(0));
290   }
291 }
292 
293 } // end namespace
294