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

source code of llvm/unittests/CodeGen/MLRegAllocDevelopmentFeatures.cpp