diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.h b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.h --- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.h +++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.h @@ -33,9 +33,20 @@ void extractInstructionFeatures( llvm::SmallVectorImpl &LRPosInfo, MLModelRunner *RegallocRunner, function_ref GetOpcode, + function_ref GetMBBFreq, + function_ref GetMBBName, const int InstructionsIndex, const int InstructionsMappingIndex, + const int MBBFreqIndex, const int MBBMappingIndex, const SlotIndex LastIndex); +void extractMBBFrequency(const SlotIndex CurrentIndex, + const size_t CurrentInstructionIndex, + std::map &VisitedMBBs, + function_ref GetMBBFreq, + const std::string CurrentMBBName, + MLModelRunner *RegallocRunner, const int MBBFreqIndex, + const int MBBMappingIndex); + // This is the maximum number of interfererring ranges. That's the number of // distinct AllocationOrder values, which comes from MCRegisterClass::RegsSize. // For X86, that's 32. @@ -69,4 +80,13 @@ static const std::vector InstructionsMappingShape{ 1, NumberOfInterferences, ModelMaxSupportedInstructionCount}; +// When extracting mappings between MBBs and individual instructions, we create +// a vector of MBB frequencies, currently of size 100, which was a value +// determined through experimentation to encompass the vast majority of eviction +// problems. The actual mapping is the same shape as the instruction opcodes +// vector. +static const int64_t ModelMaxSupportedMBBCount = 100; +static const std::vector MBBFrequencyShape{1, + ModelMaxSupportedMBBCount}; + #endif // LLVM_CODEGEN_MLREGALLOCEVICTIONADVISOR_H diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp --- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp @@ -197,7 +197,11 @@ #define RA_EVICT_REST_DEVELOPMENT_FEATURES(M) \ M(int64_t, instructions_mapping, InstructionsMappingShape, \ - "A binary matrix mapping LRs to instruction opcodes") + "A binary matrix mapping LRs to instruction opcodes") \ + M(float, mbb_frequencies, MBBFrequencyShape, \ + "A vector of machine basic block frequencies") \ + M(int64_t, mbb_mapping, InstructionsShape, \ + "A vector of indicies mapping instructions to MBBs") #else #define RA_EVICT_FIRST_DEVELOPMENT_FEATURE(M) #define RA_EVICT_REST_DEVELOPMENT_FEATURES(M) @@ -722,7 +726,19 @@ } return CurrentMachineInstruction->getOpcode(); }, + [this](SlotIndex InputIndex) -> float { + auto *CurrentMachineInstruction = + LIS->getInstructionFromIndex(InputIndex); + return MBFI.getBlockFreqRelativeToEntryBlock( + CurrentMachineInstruction->getParent()); + }, + [this](SlotIndex InputIndex) -> std::string { + auto *CurrentMachineInstruction = + LIS->getInstructionFromIndex(InputIndex); + return CurrentMachineInstruction->getParent()->getFullName(); + }, FeatureIDs::instructions, FeatureIDs::instructions_mapping, + FeatureIDs::mbb_frequencies, FeatureIDs::mbb_mapping, LIS->getSlotIndexes()->getLastIndex()); } #endif // #ifdef LLVM_HAVE_TF_API @@ -910,8 +926,12 @@ void extractInstructionFeatures(SmallVectorImpl &LRPosInfo, MLModelRunner *RegallocRunner, function_ref GetOpcode, + function_ref GetMBBFreq, + function_ref GetMBBName, const int InstructionsIndex, const int InstructionsMappingIndex, + const int MBBFreqIndex, + const int MBBMappingIndex, const SlotIndex LastIndex) { // This function extracts instruction based features relevant to the eviction // problem currently being solved. This function ends up extracting two @@ -930,6 +950,8 @@ size_t InstructionIndex = 0; size_t CurrentSegmentIndex = 0; SlotIndex CurrentIndex = LRPosInfo[0].Begin; + std::map VisitedMBBs; + size_t CurrentMBBIndex = 0; // This loop processes all the segments sequentially by starting at the // beginning slot index of the first segment, iterating through all the slot // indices before the end slot index of that segment (while checking for @@ -954,6 +976,14 @@ CurrentIndex = CurrentIndex.getNextIndex(); continue; } + std::string CurrentMBBName = GetMBBName(CurrentIndex); + if (VisitedMBBs.count(CurrentMBBName) == 0) { + VisitedMBBs[CurrentMBBName] = CurrentMBBIndex; + ++CurrentMBBIndex; + } + extractMBBFrequency(CurrentIndex, InstructionIndex, VisitedMBBs, + GetMBBFreq, CurrentMBBName, RegallocRunner, + MBBFreqIndex, MBBMappingIndex); // Current code assumes we're not going to get any disjointed segments assert(LRPosInfo[CurrentSegmentIndex].Begin <= CurrentIndex); RegallocRunner->getTensor(InstructionsIndex)[InstructionIndex] = @@ -1008,6 +1038,23 @@ } } +void extractMBBFrequency(const SlotIndex CurrentIndex, + const size_t CurrentInstructionIndex, + std::map &VisitedMBBs, + function_ref GetMBBFreq, + const std::string CurrentMBBName, + MLModelRunner *RegallocRunner, const int MBBFreqIndex, + const int MBBMappingIndex) { + size_t CurrentMBBIndex = VisitedMBBs[CurrentMBBName]; + float CurrentMBBFreq = GetMBBFreq(CurrentIndex); + if (CurrentMBBIndex < ModelMaxSupportedMBBCount) { + RegallocRunner->getTensor(MBBFreqIndex)[CurrentMBBIndex] = + CurrentMBBFreq; + RegallocRunner->getTensor( + MBBMappingIndex)[CurrentInstructionIndex] = CurrentMBBIndex; + } +} + // Development mode-specific implementations #ifdef LLVM_HAVE_TF_API diff --git a/llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll b/llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll --- a/llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll +++ b/llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll @@ -50,3 +50,16 @@ ; There are a total of 23 eviction problems with this test. ; CHECK-COUNT-22: int64_list ; CHECK: key: \"is_free\" +; Make sure that we're exporting the mbb_frequencies. Don't actually check +; values due to all values being floating point/liable to change very easily. +; CHECK: key: \"mbb_frequencies\" +; Make sure that we have the mbb_mapping feature, and that the first couple +; of values are correct. +; CHECK: key: \"mbb_mapping\" +; CHECK-NEXT: 0 +; CHECK-SAME: 0 +; CHECK-SAME: 0 +; CHECK-SAME: 0 +; CHECK-SAME: 0 +; CHECK-SAME: 1 +; CHECK-SAME: 1 diff --git a/llvm/unittests/CodeGen/MLRegallocDevelopmentFeatures.cpp b/llvm/unittests/CodeGen/MLRegallocDevelopmentFeatures.cpp --- a/llvm/unittests/CodeGen/MLRegallocDevelopmentFeatures.cpp +++ b/llvm/unittests/CodeGen/MLRegallocDevelopmentFeatures.cpp @@ -15,6 +15,7 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include using testing::ContainerEq; @@ -68,7 +69,9 @@ const std::vector Inputs{ TensorSpec::createSpec("instructions", InstructionsShape), TensorSpec::createSpec("instructions_mapping", - InstructionsMappingShape)}; + InstructionsMappingShape), + TensorSpec::createSpec("mbb_frequencies", MBBFrequencyShape), + TensorSpec::createSpec("mbb_mapping", InstructionsShape)}; LLVMContext Ctx; return NoInferenceModelRunner(Ctx, Inputs); } @@ -103,7 +106,10 @@ SlotIndex LastIndex = OverlapProblem[MaxIndex].End; extractInstructionFeatures( OverlapProblem, &ModelRunner, - [](SlotIndex InputSlot) -> int { return 0; }, 0, 1, LastIndex); + [](SlotIndex InputSlot) -> int { return 0; }, + [](SlotIndex InputSlot) -> float { return 0.0f; }, + [](SlotIndex InputSlot) -> std::string { return ""; }, 0, 1, 2, 3, + LastIndex); std::vector MappingMatrix( ModelRunner.getTensor(1), ModelRunner.getTensor(1) + @@ -154,7 +160,9 @@ [FirstIndex](SlotIndex InputSlot) -> int { return FirstIndex.distance(InputSlot) / SlotIndex::InstrDist; }, - 0, 1, LastIndex); + [](SlotIndex InputSlot) -> float { return 0.0f; }, + [](SlotIndex InputSlot) -> std::string { return ""; }, 0, 1, 2, 3, + LastIndex); for (size_t CurrentInstructionIndex = 0; CurrentInstructionIndex < ModelMaxSupportedInstructionCount; ++CurrentInstructionIndex) { @@ -207,3 +215,50 @@ OverlapSetup.push_back({35, 60, 2}); runOverlapTest(OverlapSetup); } + +TEST_F(RegallocDevelopmentFeaturesTest, SingleMBBTest) { + NoInferenceModelRunner ModelRunner = setupModelRunner(); + SlotIndex CurrentIndex; + // set index to 1 so we can ensure that the mapping actually get set + std::map VisitedMBBs = {{"testing", 1}}; + extractMBBFrequency( + CurrentIndex, 0, VisitedMBBs, + [](SlotIndex InputSlot) -> float { return 1.0f; }, "testing", + &ModelRunner, 2, 3); + ASSERT_FLOAT_EQ(ModelRunner.getTensor(2)[1], 1.0f); + ASSERT_EQ(ModelRunner.getTensor(3)[0], 1); +} + +TEST_F(RegallocDevelopmentFeaturesTest, MBBFullTruncated) { + SmallVector OverlapSetup; + OverlapSetup.push_back({0, ModelMaxSupportedInstructionCount - 1, 0}); + ilist IndexList; + auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList); + NoInferenceModelRunner ModelRunner = setupModelRunner(); + SlotIndex LastIndex = OverlapProblem[0].End; + SlotIndex FirstIndex = OverlapProblem[0].Begin; + extractInstructionFeatures( + OverlapProblem, &ModelRunner, + [](SlotIndex InputSlot) -> int { return 0; }, + [FirstIndex](SlotIndex InputSlot) -> float { + return static_cast(FirstIndex.distance(InputSlot) / + SlotIndex::InstrDist); + }, + [FirstIndex](SlotIndex InputSlot) -> std::string { + return std::to_string(FirstIndex.distance(InputSlot) / + SlotIndex::InstrDist); + }, + 0, 1, 2, 3, LastIndex); + for (size_t MBBIndex = 0; MBBIndex < ModelMaxSupportedMBBCount; ++MBBIndex) { + ASSERT_FLOAT_EQ(ModelRunner.getTensor(2)[MBBIndex], + static_cast(MBBIndex)); + ASSERT_EQ(ModelRunner.getTensor(3)[MBBIndex], + static_cast(MBBIndex)); + } + // the rest of the mapping values should be zero (truncated to 100 MBBs) + for (size_t MBBIndex = ModelMaxSupportedMBBCount; + MBBIndex < ModelMaxSupportedInstructionCount; ++MBBIndex) { + ASSERT_EQ(ModelRunner.getTensor(3)[MBBIndex], + static_cast(0)); + } +}