Add support for storing basic block context in Gematria formats.

gematria/basic_block/basic_block.cc

@@ -377,11 +377,17 @@ std::ostream& operator<<(std::ostream& os, const Instruction& instruction) {
   return os;
 }
 
-BasicBlock::BasicBlock(std::vector<Instruction> instructions)
-    : instructions(std::move(instructions)) {}
+BasicBlock::BasicBlock(std::vector<Instruction> instructions,
+                       std::vector<Instruction> back_context,
+                       std::vector<Instruction> front_context)
+    : instructions(std::move(instructions)),
+      back_context(std::move(back_context)),
+      front_context(std::move(front_context)) {}
 
 bool BasicBlock::operator==(const BasicBlock& other) const {
-  return instructions == other.instructions;
+  return instructions == other.instructions &&
+         back_context == other.back_context &&
+         front_context == other.front_context;
 }
 
 std::string BasicBlock::ToString() const {
@@ -395,6 +401,24 @@ std::string BasicBlock::ToString() const {
     if (buffer.back() == ' ') buffer.pop_back();
     buffer += "))";
   }
+  if (!back_context.empty()) {
+    buffer += "back_context=InstructionList((";
+    for (const Instruction& instruction : back_context) {
+      buffer += instruction.ToString();
+      buffer += ", ";
+    }
+    if (buffer.back() == ' ') buffer.pop_back();
+    buffer += "))";
+  }
+  if (!front_context.empty()) {
+    buffer += "front_context=InstructionList((";
+    for (const Instruction& instruction : front_context) {
+      buffer += instruction.ToString();
+      buffer += ", ";
+    }
+    if (buffer.back() == ' ') buffer.pop_back();
+    buffer += "))";
+  }
   buffer.push_back(')');
   return buffer;
 }

gematria/basic_block/basic_block.h

@@ -221,7 +221,7 @@ std::ostream& operator<<(std::ostream& os, const InstructionOperand& operand);
 // Represents an annotation holding a value such as some measure/statistic
 // paired with the instruction.
 struct Annotation {
-  Annotation() : value(-1){};
+  Annotation() : value(-1) {};
 
   // Initializes all fields of the annotation.
   Annotation(std::string name, double value);
@@ -324,9 +324,12 @@ std::ostream& operator<<(std::ostream& os, const Instruction& instruction);
 struct BasicBlock {
   BasicBlock() {}
 
-  // Initializes the basic block from a list of instructions. Needed for
-  // compatibility with the Python code.
-  explicit BasicBlock(std::vector<Instruction> instructions);
+  // Initializes the basic block from a list of instructions and optional
+  // context. Needed for compatibility with the Python code.
+  explicit BasicBlock(
+      std::vector<Instruction> instructions,
+      std::vector<Instruction> back_context = std::vector<Instruction>(),
+      std::vector<Instruction> front_context = std::vector<Instruction>());
 
   BasicBlock(const BasicBlock&) = default;
   BasicBlock(BasicBlock&&) = default;
@@ -346,6 +349,11 @@ struct BasicBlock {
 
   // The list of instructions in the basic block.
   std::vector<Instruction> instructions;
+
+  // The back and front context instructions, i.e. those preceeding and
+  // following the instructions in the basic block.
+  std::vector<Instruction> back_context;
+  std::vector<Instruction> front_context;
 };
 
 std::ostream& operator<<(std::ostream& os, const BasicBlock& block);

gematria/basic_block/basic_block_protos.cc

@@ -180,8 +180,15 @@ CanonicalizedInstructionProto ProtoFromInstruction(
 
 BasicBlock BasicBlockFromProto(const BasicBlockProto& proto) {
   return BasicBlock(
-      /* instructions = */ ToVector<Instruction>(
-          proto.canonicalized_instructions(), InstructionFromProto));
+      /* instructions = */
+      ToVector<Instruction>(proto.canonicalized_instructions(),
+                            InstructionFromProto),
+      /* back_context = */
+      ToVector<Instruction>(proto.canonicalized_back_context(),
+                            InstructionFromProto),
+      /* front_context = */
+      ToVector<Instruction>(proto.canonicalized_front_context(),
+                            InstructionFromProto));
 }
 
 }  // namespace gematria

gematria/basic_block/basic_block_test.cc

@@ -277,7 +277,7 @@ TEST(InstructionOperandTest, Equality) {
 
 TEST(InstructionOperandTest, ToString) {
   const struct {
-    InstructionOperand opernad;
+    InstructionOperand operand;
     const char* expected_string;
   } kTestCases[] = {
       {InstructionOperand::Register("RAX"),
@@ -292,7 +292,7 @@ TEST(InstructionOperandTest, ToString) {
        "InstructionOperand.from_memory(32)"}};
 
   for (const auto& test_case : kTestCases) {
-    EXPECT_EQ(test_case.opernad.ToString(), test_case.expected_string);
+    EXPECT_EQ(test_case.operand.ToString(), test_case.expected_string);
   }
 }
 
@@ -318,7 +318,6 @@ TEST(InstructionOperandTest, AsTokenList) {
   }
 }
 
-// TODO(virajbshah): Add tests for Annotation.
 TEST(AnnotationTest, Constructor) {
   constexpr char kName[] = "cache_miss_freq";
   constexpr double kValue = 0.875;

gematria/basic_block/python/basic_block.cc

@@ -255,9 +255,15 @@ PYBIND11_MODULE(basic_block, m) {
 
   py::class_<BasicBlock> basic_block(m, "BasicBlock");
   basic_block
-      .def(py::init<std::vector<Instruction> /* instructions */>(),
-           py::arg("instructions") = std::vector<Instruction>())
+      .def(py::init<std::vector<Instruction> /* instructions */,
+                    std::vector<Instruction> /* back_context */,
+                    std::vector<Instruction> /* front_context */>(),
+           py::arg("instructions") = std::vector<Instruction>(),
+           py::arg("back_context") = std::vector<Instruction>(),
+           py::arg("front_context") = std::vector<Instruction>())
       .def_readwrite("instructions", &BasicBlock::instructions)
+      .def_readwrite("back_context", &BasicBlock::back_context)
+      .def_readwrite("front_context", &BasicBlock::front_context)
       .def("__repr__", &BasicBlock::ToString)
       .def("__str__", &BasicBlock::ToString)
       .def("__eq__", &BasicBlock::operator==)

gematria/proto/basic_block.proto

@@ -30,8 +30,26 @@ message BasicBlockProto {
   // same instruction.
   repeated CanonicalizedInstructionProto canonicalized_instructions = 2;
 
+  // An optional list of machine instructions preceding the basic block, used
+  // to provide context that lies before `canonicalized_instructions`. These
+  // instructions are not included in the timing measurements and predictions.
+  repeated MachineInstructionProto machine_back_context = 3;
+
+  // An optional list of machine instructions following the basic block, used
+  // to provide context lying after `canonicalized_instructions`. These
+  // instructions are not included in the timing measurements and predictions.
+  repeated MachineInstructionProto machine_front_context = 4;
+
+  // Canonicalized instructions parallel to `machine_back_context`. May be
+  // empty in case no back context is provided.
+  repeated CanonicalizedInstructionProto canonicalized_back_context = 5;
+
+  // Canonicalized instructions parallel to `machine_front_context`. May be
+  // empty in case no front context is provided.
+  repeated CanonicalizedInstructionProto canonicalized_front_context = 6;
+
   // The fingerprint-id of this basic block. Might be empty.
-  string fingerprint = 3;
+  string fingerprint = 7;
 }
 
 // Represents a raw instruction extracted from binary code.