[MLGO] Add per-instruction MBB frequencies to regalloc dev features

[aidengrossman]

Summary

This commit adds in two new features to the ML regalloc eviction
analysis that can be used in ML models, a vector of MBB frequencies and
a vector of indicies mapping instructions to their corresponding basic
blocks. This will allow for further experimentation with per-instruction
features and give a lot more flexibility for future experimentation over
how we're extracting MBB frequency data currently.

[aidengrossman]

Here's the patch for the MBB based features. I've set it up in a way that isn't necessarily the most efficienct, but should allow for extra flexibility for experimentation on the ML side. I've opted to set up the main MBB feature extraction function to get called at every loop iteration in the current per-instruction feature extraction code, but due to how we currently have unit testing setup, this adds a couple more parameters to that function and leaves things a little bit less elegant than they probably could be. I think this should be fine for now, but I can definitely refactor to a different paradigm if necessary, but figured it would be better to discuss solutions first (and see if it's even necessary) before doing that work.

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

1057

Why not use the MBB address?

[aidengrossman]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

1057

Give me a little bit of time to do some experimentation/investigation, but I tried calling getSectionIDNum() and utilizing that, but it would only return a single value for all the MBBs in the MF (only a single value would appear in the map).

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

1057

I mean &MBB. I.e. MachineBasicBlock*

[aidengrossman]

Moved from mapping MBBs to indicies through the use of MBB pointers rather
than grabbing their names as strings.

[aidengrossman]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

1057

Mapping using the MBB pointer as the key works perfectly and avoids passing around/generating strings. I've switched over. I had to do some slightly unorthodox casting to generate some test values in the unittest. I've added a comment explaining it, but if necessary I can modify it.

[aidengrossman]

Refactored test pointer generation/arithmetic in MBB test case to avoid
compiler warning.

[mtrofin]

llvm/unittests/CodeGen/MLRegallocDevelopmentFeatures.cpp

263

Try rather making real MBB objects instead - the problem is that, if later we start dereferencing the MBB, this test will probably fail, and probably with non-deterministic failures - because the addressed memory will be invalid.

Should be like a 2-liner, make a Function, a MachineFunction, and create MBBs. I think some unittests have examples.

[aidengrossman]

Switched from using dummy pointer values in development regalloc feature
test to using pointers to actual MBBs in case we ever end up dereferencing
these values at some point in the future.

[mtrofin]

lgtm, and I apologize, but I should have added Jacob earlier to ptal, since he's also working on this. Please let him also take a look before landing.

[jacobhegna]

lgtm, this will be really useful!

just so I understand, mbb_frequencies will be of shape (1, 100) and if there are less than 100 mbbs, the rest of the entries will be 0? and mbb_mapping will of shape (1, 300) and contains the indices of the mbb for each instruction?

just a small question/nit: why use a one-hot encoding for the instruction opcode features but an index encoding for the frequencies?

[aidengrossman]

just so I understand, mbb_frequencies will be of shape (1, 100) and if there are less than 100 mbbs, the rest of the entries will be 0? and mbb_mapping will of shape (1, 300) and contains the indices of the mbb for each instruction?

Yes. Up to 100 MBBs have their frequency data stored and otherwise everything is zero padded. The mapping vector is intended to work exactly as you describe. I settled on 100 MBBs after monitoring a LLVM compile after setting up some debug logging to see how many MBBs there were in each eviction problem and didn't see any with more than 100.

just a small question/nit: why use a one-hot encoding for the instruction opcode features but an index encoding for the frequencies?

I did this to give added flexibility in terms of how we integrate these features on the ML side. Currently just encoding a vector of length (1, ModelMaxSupportedInstructionCount) with the MBB frequencies would work for what I have been experimenting with, but if we want to do anything else in the future, the way this is setup gives added flexibility in that area (ie grouping instructions by MBB). This is the same reason I set up passing the instructions through an opcode vector and a binary mapping matrix rather than just passing the opcodes directly in a (33, 300) matrix.

[jacobhegna]

sounds good!

llvm/lib/CodeGen/MLRegallocEvictAdvisor.h

 // testing.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_MLREGALLOCEVICTIONADVISOR_H
 #define LLVM_CODEGEN_MLREGALLOCEVICTIONADVISOR_H
 
 #include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 
 using namespace llvm;
 
 // LRStartEndInfo contains the start and end of a specific live range as
 // slot indices as well as storing the index of the physical register it
 // is assigned to (or 1 above the phys reg count if its the candidate).
 // Used when extracting per-instruction features in the context of a
 // specific eviction problem.
 struct LRStartEndInfo {
   SlotIndex Begin;
   SlotIndex End;
   size_t Pos = 0;
 };
 
 void extractInstructionFeatures(
     llvm::SmallVectorImpl<LRStartEndInfo> &LRPosInfo,
     MLModelRunner *RegallocRunner, function_ref<int(SlotIndex)> GetOpcode,
+    function_ref<float(SlotIndex)> GetMBBFreq,
+    function_ref<MachineBasicBlock *(SlotIndex)> GetMBBReference,
     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<MachineBasicBlock *, size_t> &VisitedMBBs,
+                         function_ref<float(SlotIndex)> GetMBBFreq,
+                         MachineBasicBlock *CurrentMBBReference,
+                         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.
 // TODO: find a way to get this, statically, in a programmatic way.
 static const int64_t MaxInterferences = 32;
 
 // Logically, we can think of the feature set given to the evaluator as a 2D
 // matrix. The rows are the features (see next). The columns correspond to the

 // instructions relevant to the eviction problem, and a NumberOfInterferences *
 // ModelMaxSupportedInstructionCount matrix that maps LRs to the instructions
 // that they span.
 static const std::vector<int64_t> InstructionsShape{
     1, ModelMaxSupportedInstructionCount};
 static const std::vector<int64_t> 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<int64_t> MBBFrequencyShape{1,
+                                                    ModelMaxSupportedMBBCount};
+
 #endif // LLVM_CODEGEN_MLREGALLOCEVICTIONADVISOR_H

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 
 #ifdef LLVM_HAVE_TF_API
 #define RA_EVICT_FIRST_DEVELOPMENT_FEATURE(M)                                  \
   M(int64_t, instructions, InstructionsShape,                                  \
     "Opcodes of the instructions covered by the eviction problem")
 
 #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)
 #endif
 
 // The model learns to pick one of the mask == 1 interferences. This is the
 // name of the output tensor. The contract with the model is that the output
 // will be guaranteed to be to a mask == 1 position. Using a macro here to
extractInstructionFeatures(
         [this](SlotIndex InputIndex) -> int {
           auto *CurrentMachineInstruction =
               LIS->getInstructionFromIndex(InputIndex);
           if (!CurrentMachineInstruction) {
             return -1;
           }
           return CurrentMachineInstruction->getOpcode();
         },
+        [this](SlotIndex InputIndex) -> float {
+          auto *CurrentMachineInstruction =
+              LIS->getInstructionFromIndex(InputIndex);
+          return MBFI.getBlockFreqRelativeToEntryBlock(
+              CurrentMachineInstruction->getParent());
+        },
+        [this](SlotIndex InputIndex) -> MachineBasicBlock * {
+          auto *CurrentMachineInstruction =
+              LIS->getInstructionFromIndex(InputIndex);
+          return CurrentMachineInstruction->getParent();
+        },
         FeatureIDs::instructions, FeatureIDs::instructions_mapping,
+        FeatureIDs::mbb_frequencies, FeatureIDs::mbb_mapping,
         LIS->getSlotIndexes()->getLastIndex());
   }
 #endif // #ifdef LLVM_HAVE_TF_API
   // Normalize the features.
   for (auto &V : Largest)
     V = V ? V : 1.0;
   for (size_t FeatureIndex = 0; FeatureIndex < FeatureIDs::FeatureCount;
        ++FeatureIndex) {
#define SET(ID, TYPE, VAL)                                                     \
   SET(hottest_bb_freq_by_max, float, HottestBlockFreq);
   SET(liverange_size, float, Size);
   SET(use_def_density, float, TotalWeight);
   SET(max_stage, int64_t, MaxStage);
   SET(min_stage, int64_t, MinStage);
 #undef SET
 }
 
-void extractInstructionFeatures(SmallVectorImpl<LRStartEndInfo> &LRPosInfo,
-                                MLModelRunner *RegallocRunner,
-                                function_ref<int(SlotIndex)> GetOpcode,
-                                const int InstructionsIndex,
-                                const int InstructionsMappingIndex,
-                                const SlotIndex LastIndex) {
+void extractInstructionFeatures(
+    SmallVectorImpl<LRStartEndInfo> &LRPosInfo, MLModelRunner *RegallocRunner,
+    function_ref<float(SlotIndex)> GetMBBFreq,
+    function_ref<MachineBasicBlock *(SlotIndex)> GetMBBReference,
+    const int InstructionsIndex, const int InstructionsMappingIndex,
+    const int MBBFreqIndex, const int MBBMappingIndex,
   // This function extracts instruction based features relevant to the eviction
   // problem currently being solved. This function ends up extracting two
   // tensors.
   // 1 - A vector of size max instruction count. It contains the opcodes of the
   // instructions spanned by all the intervals in the current instance of the
   // eviction problem.
   // 2 - A binary mapping matrix of size (LR count * max
   // instruction count) which maps where the LRs are live to the actual opcodes
   // for which they are live.
+  // 3 - A vector of size max supported MBB count storing MBB frequencies,
+  // encompassing all of the MBBs covered by the eviction problem.
+  // 4 - A vector of size max instruction count of indices to members of the MBB
+  // frequency vector, mapping each instruction to its associated MBB.
 
   // Start off by sorting the segments based on the beginning slot index.
   std::sort(
       LRPosInfo.begin(), LRPosInfo.end(),
       [](LRStartEndInfo A, LRStartEndInfo B) { return A.Begin < B.Begin; });
   size_t InstructionIndex = 0;
   size_t CurrentSegmentIndex = 0;
   SlotIndex CurrentIndex = LRPosInfo[0].Begin;
+  std::map<MachineBasicBlock *, size_t> 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
   // overlaps with segments that start at greater slot indices). After hitting
   // that end index, the current segment being processed gets bumped until they
   // are all processed or the max instruction count is hit, where everything is
   // just truncated.
   while (true) {
     // If the index that we are currently at is within the current segment and
     // we haven't hit the max instruction count, continue processing the current
     // segment.
     while (CurrentIndex <= LRPosInfo[CurrentSegmentIndex].End &&
            InstructionIndex < ModelMaxSupportedInstructionCount) {
       int CurrentOpcode = GetOpcode(CurrentIndex);
       // If the current machine instruction is null, skip it
       if (CurrentOpcode == -1) {
         // If we're currently at the last index in the SlotIndex analysis,
         // we can't go any further, so return from the function
         if (CurrentIndex >= LastIndex) {
           return;
         }
         CurrentIndex = CurrentIndex.getNextIndex();
         continue;
       }
+      MachineBasicBlock *CurrentMBBReference = GetMBBReference(CurrentIndex);
+      if (VisitedMBBs.count(CurrentMBBReference) == 0) {
+        VisitedMBBs[CurrentMBBReference] = CurrentMBBIndex;
+        ++CurrentMBBIndex;
+      }
+      extractMBBFrequency(CurrentIndex, InstructionIndex, VisitedMBBs,
+                          GetMBBFreq, CurrentMBBReference, RegallocRunner,
+                          MBBFreqIndex, MBBMappingIndex);
       // Current code assumes we're not going to get any disjointed segments
       assert(LRPosInfo[CurrentSegmentIndex].Begin <= CurrentIndex);
       RegallocRunner->getTensor<int64_t>(InstructionsIndex)[InstructionIndex] =
           CurrentOpcode < OpcodeValueCutoff ? CurrentOpcode : 0;
       // set value in the binary mapping matrix for the current instruction
       auto CurrentSegmentPosition = LRPosInfo[CurrentSegmentIndex].Pos;
       RegallocRunner->getTensor<int64_t>(
           InstructionsMappingIndex)[CurrentSegmentPosition *
while (true) {
     if (LRPosInfo[CurrentSegmentIndex + 1].Begin >
         LRPosInfo[CurrentSegmentIndex].End) {
       CurrentIndex = LRPosInfo[CurrentSegmentIndex + 1].Begin;
     }
     ++CurrentSegmentIndex;
   }
 }
 
+void extractMBBFrequency(const SlotIndex CurrentIndex,
+                         const size_t CurrentInstructionIndex,
+                         std::map<MachineBasicBlock *, size_t> &VisitedMBBs,
+                         function_ref<float(SlotIndex)> GetMBBFreq,
+                         MachineBasicBlock *CurrentMBBReference,
+                         MLModelRunner *RegallocRunner, const int MBBFreqIndex,
+                         const int MBBMappingIndex) {
+  size_t CurrentMBBIndex = VisitedMBBs[CurrentMBBReference];
+  float CurrentMBBFreq = GetMBBFreq(CurrentIndex);
+  if (CurrentMBBIndex < ModelMaxSupportedMBBCount) {
+    RegallocRunner->getTensor<float>(MBBFreqIndex)[CurrentMBBIndex] =
+        CurrentMBBFreq;
+    RegallocRunner->getTensor<int64_t>(
+        MBBMappingIndex)[CurrentInstructionIndex] = CurrentMBBIndex;
+  }
+}
+
 // Development mode-specific implementations
 #ifdef LLVM_HAVE_TF_API
 
 RegAllocEvictionAdvisorAnalysis *llvm::createDevelopmentModeAdvisor() {
   return new DevelopmentModeEvictionAdvisorAnalysis();
 }
 
 int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(

llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll

 ; that not all the values between the 10th LR and the candidate are zero.
 ; CHECK-COUNT-6600: value: 0
 ; CHECK-SAME: value: 1
 ; Ensure that we can still go through the mapping matrices for the rest of the
 ; eviction problems to make sure we haven't hit the end of the matrix above.
 ; 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

llvm/unittests/CodeGen/MLRegallocDevelopmentFeatures.cpp

 //===- MLRegAllocDevelopmentFeatures.cpp - test dev MLRegalloc features ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "../../lib/CodeGen/MLRegallocEvictAdvisor.h"
+#include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Target/TargetOptions.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
+#include <string>
 #include <vector>
 
 using testing::ContainerEq;
 using testing::Test;
 
+namespace {
+
+#include "MFCommon.inc"
+
 struct LRPosInfoIndexes {
   size_t StartIndex;
   size_t EndIndex;
   size_t PhysReg;
 };
 
 class RegallocDevelopmentFeaturesTest : public ::Test {
 protected:
setupOverlapProblem(const SmallVectorImpl<LRPosInfoIndexes> &Segments,
     }
     return PositionsToReturn;
   }
 
   NoInferenceModelRunner setupModelRunner() {
     const std::vector<TensorSpec> Inputs{
         TensorSpec::createSpec<int64_t>("instructions", InstructionsShape),
         TensorSpec::createSpec<int64_t>("instructions_mapping",
-                                        InstructionsMappingShape)};
+                                        InstructionsMappingShape),
+        TensorSpec::createSpec<float>("mbb_frequencies", MBBFrequencyShape),
+        TensorSpec::createSpec<int64_t>("mbb_mapping", InstructionsShape)};
     LLVMContext Ctx;
     return NoInferenceModelRunner(Ctx, Inputs);
   }
 
   std::vector<int64_t>
   getExpectedMappingMatrix(SmallVectorImpl<LRPosInfoIndexes> &OverlapSetup) {
     std::vector<int64_t> ExpectedMappingMatrix(
         NumberOfInterferences * ModelMaxSupportedInstructionCount, 0);
for (size_t CurrentOverlap = 0; CurrentOverlap < OverlapSetup.size();
       if (OverlapSetup[CurrentOverlap].EndIndex >
           OverlapSetup[MaxIndex].EndIndex) {
         MaxIndex = CurrentOverlap;
       }
     }
     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) -> MachineBasicBlock * { return nullptr; }, 0,
+        1, 2, 3, LastIndex);
     std::vector<int64_t> MappingMatrix(
         ModelRunner.getTensor<int64_t>(1),
         ModelRunner.getTensor<int64_t>(1) +
             NumberOfInterferences * ModelMaxSupportedInstructionCount);
     ASSERT_THAT(MappingMatrix,
                 ContainerEq(getExpectedMappingMatrix(OverlapSetup)));
     IndexList.clearAndLeakNodesUnsafely();
   }
TEST_F(RegallocDevelopmentFeaturesTest, InstructionOpcodesAreCorrect) {
   NoInferenceModelRunner ModelRunner = setupModelRunner();
   SlotIndex LastIndex = OverlapProblem[0].End;
   SlotIndex FirstIndex = OverlapProblem[0].Begin;
   extractInstructionFeatures(
       OverlapProblem, &ModelRunner,
       [FirstIndex](SlotIndex InputSlot) -> int {
         return FirstIndex.distance(InputSlot) / SlotIndex::InstrDist;
       },
-      0, 1, LastIndex);
+      [](SlotIndex InputSlot) -> float { return 0.0f; },
+      [](SlotIndex InputSlot) -> MachineBasicBlock * { return nullptr; }, 0, 1,
+      2, 3, LastIndex);
   for (size_t CurrentInstructionIndex = 0;
        CurrentInstructionIndex < ModelMaxSupportedInstructionCount;
        ++CurrentInstructionIndex) {
     ASSERT_EQ(
         (size_t)ModelRunner.getTensor<int64_t>(0)[CurrentInstructionIndex],
         CurrentInstructionIndex);
   }
 }

 
 TEST_F(RegallocDevelopmentFeaturesTest, InternalMultiOverlap) {
   SmallVector<LRPosInfoIndexes, 3> OverlapSetup;
   OverlapSetup.push_back({0, 45, 0});
   OverlapSetup.push_back({30, 40, 1});
   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<MachineBasicBlock *, size_t> VisitedMBBs = {{nullptr, 1}};
+  extractMBBFrequency(
+      CurrentIndex, 0, VisitedMBBs,
+      [](SlotIndex InputSlot) -> float { return 1.0f; }, nullptr, &ModelRunner,
+      2, 3);
+  ASSERT_FLOAT_EQ(ModelRunner.getTensor<float>(2)[1], 1.0f);
+  ASSERT_EQ(ModelRunner.getTensor<int64_t>(3)[0], 1);
+}
+
+TEST_F(RegallocDevelopmentFeaturesTest, MBBFullTruncated) {
+  SmallVector<LRPosInfoIndexes, 1> OverlapSetup;
+  OverlapSetup.push_back({0, ModelMaxSupportedInstructionCount - 1, 0});
+  ilist<IndexListEntry> IndexList;
+  auto OverlapProblem = setupOverlapProblem(OverlapSetup, IndexList);
+  NoInferenceModelRunner ModelRunner = setupModelRunner();
+  SlotIndex LastIndex = OverlapProblem[0].End;
+  SlotIndex FirstIndex = OverlapProblem[0].Begin;
+
+  LLVMContext Ctx;
+  Module Mod("Module", Ctx);
+  auto MF = createMachineFunction(Ctx, Mod);
+  std::array<MachineBasicBlock *, ModelMaxSupportedInstructionCount>
+      MBBsForTest;
+  for (size_t I = 0; I < ModelMaxSupportedInstructionCount; ++I) {
+    MBBsForTest[I] = MF->CreateMachineBasicBlock();
+  }
+
+  extractInstructionFeatures(
+      OverlapProblem, &ModelRunner,
+      [](SlotIndex InputSlot) -> int { return 0; },
+      [FirstIndex](SlotIndex InputSlot) -> float {
+        return static_cast<float>(FirstIndex.distance(InputSlot) /
+                                  SlotIndex::InstrDist);
+      },
+      [FirstIndex, MBBsForTest](SlotIndex InputSlot) -> MachineBasicBlock * {
+        return MBBsForTest[FirstIndex.distance(InputSlot) /
+                           SlotIndex::InstrDist];
+      },
+      0, 1, 2, 3, LastIndex);
+  for (size_t MBBIndex = 0; MBBIndex < ModelMaxSupportedMBBCount; ++MBBIndex) {
+    ASSERT_FLOAT_EQ(ModelRunner.getTensor<float>(2)[MBBIndex],
+                    static_cast<float>(MBBIndex));
+    ASSERT_EQ(ModelRunner.getTensor<int64_t>(3)[MBBIndex],
+              static_cast<int64_t>(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<int64_t>(3)[MBBIndex],
+              static_cast<int64_t>(0));
+  }
+}
+
+} // end namespace