Interference graph support

Dockerfile

@@ -4,4 +4,4 @@ ARG bazelisk_version=1.17.0
 RUN curl -L https://github.com/bazelbuild/bazelisk/releases/download/v${bazelisk_version}/bazelisk-linux-amd64 > /usr/bin/bazelisk && chmod +x /usr/bin/bazelisk && ln -s /usr/bin/bazelisk /usr/bin/bazel
 WORKDIR /granlte
 COPY . .
-RUN pip3 install -r requirements.in
+RUN pip3 install -r requirements.in

gematria/basic_block/basic_block.cc

@@ -60,18 +60,36 @@ std::string AddressTuple::ToString() const {
   buffer << "AddressTuple(";
   if (!base_register.empty()) {
     buffer << "base_register='" << base_register << "', ";
+    buffer << "base_register_size=" << base_register_size << ", ";
+    buffer << "base_register_intefered_register={";
+    for (const std::string& interfered_register : base_register_intefered_register) {
+      buffer << "'" << interfered_register << "', ";
+    }
+    buffer << "}, ";
   }
   if (displacement != 0) {
     buffer << "displacement=" << displacement << ", ";
   }
   if (!index_register.empty()) {
     buffer << "index_Register='" << index_register << "', ";
+    buffer << "index_register_size=" << index_register_size << ", ";
+    buffer << "index_register_intefered_register={";
+    for (const std::string& interfered_register : index_register_intefered_register) {
+      buffer << "'" << interfered_register << "', ";
+    }
+    buffer << "}, ";
   }
   if (!index_register.empty() || scaling != 0) {
     buffer << "scaling=" << scaling << ", ";
   }
   if (!segment_register.empty()) {
     buffer << "segment_register='" << segment_register << "', ";
+    buffer << "segment_register_size=" << segment_register_size << ", ";
+    buffer << "segment_register_intefered_register={";
+    for (const std::string& interfered_register : segment_register_intefered_register) {
+      buffer << "'" << interfered_register << "', ";
+    }
+    buffer << "}, ";
   }
   // If we added any keyword args to the buffer, drop the last two characters
   // (a comma and a space). This is not strictly necessary, but it looks better.
@@ -108,11 +126,12 @@ bool InstructionOperand::operator==(const InstructionOperand& other) const {
 }
 
 InstructionOperand InstructionOperand::VirtualRegister(
-    const std::string register_name, size_t size) {
+    const std::string register_name, size_t size, const std::vector<std::string>& interfered_registers) {
   InstructionOperand result;
   result.type_ = OperandType::kVirtualRegister;
   result.register_name_ = std::move(register_name);
   result.size_ = size;
+  result.interfered_registers_ = std::move(interfered_registers);
   return result;
 }
 
@@ -151,14 +170,23 @@ InstructionOperand InstructionOperand::Address(std::string base_register,
                                                int64_t displacement,
                                                std::string index_register,
                                                int scaling,
-                                               std::string segment_register) {
+                                               std::string segment_register,
+                                               int base_register_size,
+                                               int index_register_size,
+                                               int segment_register_size) {
   InstructionOperand result;
   result.type_ = OperandType::kAddress;
   result.address_.base_register = std::move(base_register);
   result.address_.index_register = std::move(index_register);
   result.address_.displacement = displacement;
   result.address_.scaling = scaling;
   result.address_.segment_register = segment_register;
+  result.address_.base_register_size = base_register_size;
+  result.address_.index_register_size = index_register_size;
+  result.address_.segment_register_size = segment_register_size;
+  result.address_.base_register_intefered_register = {};
+  result.address_.index_register_intefered_register = {};
+  result.address_.segment_register_intefered_register = {};
   return result;
 }
 

gematria/basic_block/basic_block.h

@@ -85,12 +85,23 @@ struct AddressTuple {
   AddressTuple(AddressTuple&&) = default;
   AddressTuple(std::string base_register, int64_t displacement,
                std::string index_register, int scaling,
-               std::string segment_register)
+               std::string segment_register, int base_register_size = 64, 
+               int index_register_size = 64, int segment_register_size = 64, 
+               const std::vector<std::string> base_register_intefered_register = {}, 
+               const std::vector<std::string> index_register_intefered_register = {},
+               const std::vector<std::string> segment_register_intefered_register = {})
       : base_register(std::move(base_register)),
         displacement(displacement),
         index_register(std::move(index_register)),
         scaling(scaling),
-        segment_register(std::move(segment_register)) {}
+        segment_register(std::move(segment_register)),
+        base_register_size(std::move(base_register_size)),
+        index_register_size(std::move(index_register_size)),
+        segment_register_size(std::move(segment_register_size)),
+        base_register_intefered_register(std::move(base_register_intefered_register)),
+        index_register_intefered_register(std::move(index_register_intefered_register)),
+        segment_register_intefered_register(std::move(segment_register_intefered_register))
+        {}
 
   AddressTuple& operator=(const AddressTuple&) = default;
   AddressTuple& operator=(AddressTuple&&) = default;
@@ -126,6 +137,19 @@ struct AddressTuple {
   // The name of the segment register. When empty, the default segment register
   // for the instruction is used.
   std::string segment_register;
+  // The size of the base register. Used only when base_register is non-empty.
+  size_t base_register_size;
+  // The size of the index register. Used only when index_register is non-empty.
+  size_t index_register_size;
+  // The size of the segment register. Used only when segment_register is
+  size_t segment_register_size;
+
+  // The name of the index register of the address. When empty, index register
+  std::vector<std::string> base_register_intefered_register;
+  // The name of the index register of the address. When empty, index register
+  std::vector<std::string> index_register_intefered_register;
+  // The name of the index register of the address. When empty, index register
+  std::vector<std::string> segment_register_intefered_register;
 };
 
 std::ostream& operator<<(std::ostream& os, const AddressTuple& address_tuple);
@@ -147,15 +171,18 @@ class InstructionOperand {
 
   // The operands must be created through one of the factory functions.
   static InstructionOperand VirtualRegister(std::string register_name,
-                                            size_t size);
+                                            size_t size, const std::vector<std::string>& interfered_registers);
   static InstructionOperand Register(std::string register_name);
   static InstructionOperand ImmediateValue(uint64_t immediate_value);
   static InstructionOperand FpImmediateValue(double fp_immediate_value);
   static InstructionOperand Address(AddressTuple address_tuple);
   static InstructionOperand Address(std::string base_register,
                                     int64_t displacement,
                                     std::string index_register, int scaling,
-                                    std::string segment_register);
+                                    std::string segment_register, 
+                                    int base_register_size = 64,
+                                    int index_register_size = 64,
+                                    int segment_register_size = 64);
   static InstructionOperand MemoryLocation(int alias_group_id);
 
   bool operator==(const InstructionOperand&) const;
@@ -169,6 +196,11 @@ class InstructionOperand {
   // Returns the list of tokens representing this instruction.
   std::vector<std::string> AsTokenList() const;
 
+  std::vector<std::string> getInterferedRegisters() const {
+    assert(type_ == OperandType::kVirtualRegister);
+    return interfered_registers_;
+  }
+
   // Returns a human-readable representation of the operand.
   //
   // This method implements the __str__() and __repr__() methods in the Python
@@ -225,6 +257,7 @@ class InstructionOperand {
   double fp_immediate_value_ = 0.0;
   AddressTuple address_;
   int alias_group_id_ = 0;
+  std::vector<std::string> interfered_registers_;
 };
 
 std::ostream& operator<<(std::ostream& os, const InstructionOperand& operand);

gematria/basic_block/basic_block_protos.cc

@@ -25,14 +25,34 @@
 
 namespace gematria {
 
+namespace {
+  std::vector<std::string> ToVector(
+    const google::protobuf::RepeatedPtrField<std::string>& protos) {
+  return std::vector<std::string>(protos.begin(), protos.end());
+  }
+}
+
 AddressTuple AddressTupleFromProto(
     const CanonicalizedOperandProto::AddressTuple& proto) {
-  return AddressTuple(
+  auto result = AddressTuple(
       /* base_register = */ proto.base_register(),
       /* displacement = */ proto.displacement(),
       /* index_register = */ proto.index_register(),
       /* scaling = */ proto.scaling(),
       /* segment_register = */ proto.segment());
+  if (proto.base_register()[0] == '%'){
+    result.base_register_size = proto.base_register_size();
+    result.base_register_intefered_register = ToVector(proto.base_register_intefered_register());
+  }
+  if (proto.index_register()[0] == '%'){
+    result.index_register_size = proto.index_register_size();
+    result.index_register_intefered_register = ToVector(proto.index_register_intefered_register());
+  }
+  if (proto.segment()[0] == '%'){
+    result.segment_register_size = proto.segment_size();
+    result.segment_register_intefered_register = ToVector(proto.segment_intefered_register());
+  }
+  return result;
 }
 
 CanonicalizedOperandProto::AddressTuple ProtoFromAddressTuple(
@@ -43,6 +63,24 @@ CanonicalizedOperandProto::AddressTuple ProtoFromAddressTuple(
   proto.set_index_register(address_tuple.index_register);
   proto.set_scaling(address_tuple.scaling);
   proto.set_segment(address_tuple.segment_register);
+  if (!address_tuple.base_register.empty() && address_tuple.base_register[0] == '%') {
+    proto.set_base_register_size(address_tuple.base_register_size);
+    for (auto interfered_register : address_tuple.base_register_intefered_register){
+      proto.add_base_register_intefered_register(std::move(interfered_register));
+    }
+  }
+  if (!address_tuple.index_register.empty() && address_tuple.index_register[0] == '%') {
+    proto.set_index_register_size(address_tuple.index_register_size);
+    for (auto interfered_register : address_tuple.index_register_intefered_register){
+      proto.add_index_register_intefered_register(std::move(interfered_register));
+    }
+  }
+  if (!address_tuple.segment_register.empty() && address_tuple.segment_register[0] == '%') {
+    proto.set_segment_size(address_tuple.segment_register_size);
+    for (auto interfered_register : address_tuple.segment_register_intefered_register){
+      proto.add_segment_intefered_register(std::move(interfered_register));
+    }
+  }
   return proto;
 }
 
@@ -64,8 +102,12 @@ InstructionOperand InstructionOperandFromProto(
       return InstructionOperand::MemoryLocation(
           proto.memory().alias_group_id());
     case CanonicalizedOperandProto::kVirtualRegister:
-      return InstructionOperand::VirtualRegister(
-          proto.virtual_register().name(), proto.virtual_register().size());
+      {
+        std::vector<std::string> interfered_registers = ToVector(proto.intefered_register());
+        return InstructionOperand::VirtualRegister(
+          proto.virtual_register().name(), proto.virtual_register().size(), interfered_registers);
+      }
+
   }
 }
 
@@ -102,7 +144,6 @@ CanonicalizedOperandProto ProtoFromInstructionOperand(
 }
 
 namespace {
-
 std::vector<InstructionOperand> ToVector(
     const google::protobuf::RepeatedPtrField<CanonicalizedOperandProto>&
         protos) {

gematria/basic_block/basic_block_protos_test.cc

@@ -239,7 +239,11 @@ TEST(BasicBlockFromProtoTest, VRegInstructions) {
     canonicalized_instructions {
       mnemonic: "CMP64RI32"
       llvm_mnemonic: "CMP64ri32"
-      input_operands { virtual_register { name: "%60" size: 64 } }
+      input_operands { 
+        virtual_register { name: "%60" size: 64 } 
+        intefered_register: "%61"
+        intefered_register: "%62"
+      }
       input_operands { immediate_value: 0 }
       implicit_output_operands { register_name: "EFLAGS" }
     }
@@ -250,7 +254,7 @@ TEST(BasicBlockFromProtoTest, VRegInstructions) {
                 /* mnemonic = */ "CMP64RI32", /* llvm_mnemonic = */ "CMP64ri32",
                 /* prefixes = */ {},
                 /* input_operands = */
-                {InstructionOperand::VirtualRegister("%60", 64),
+                {InstructionOperand::VirtualRegister("%60", 64, {"%61, %62"}),
                  InstructionOperand::ImmediateValue(0)},
                 /* implicit_input_operands = */ {},
                 /* output_operands = */ {},

gematria/basic_block/basic_block_test.cc

@@ -50,39 +50,49 @@ TEST(AddressTupleTest, ToString) {
   const struct {
     AddressTuple address;
     const char* expected_string;
-  } kTestCases[] = {
-      {AddressTuple(/* base_register = */ "RAX",
-                    /* displacement = */ 0,
-                    /* index_register = */ "",
-                    /* scaling = */ 0,
-                    /* segment_register = */ ""),
-       "AddressTuple(base_register='RAX')"},
-      {AddressTuple(/* base_register = */ "RAX",
-                    /* displacement = */ 16,
-                    /* index_register = */ "",
-                    /* scaling = */ 0,
-                    /* segment_register = */ ""),
-       "AddressTuple(base_register='RAX', displacement=16)"},
-      {AddressTuple(/* base_register = */ "RAX",
-                    /* displacement = */ 0,
-                    /* index_register = */ "RSI",
-                    /* scaling = */ 0,
-                    /* segment_register = */ ""),
-       "AddressTuple(base_register='RAX', index_Register='RSI', scaling=0)"},
-      {AddressTuple(/* base_register = */ "RAX",
-                    /* displacement = */ -8,
-                    /* index_register = */ "RSI",
-                    /* scaling = */ 1,
-                    /* segment_register = */ ""),
-       "AddressTuple(base_register='RAX', displacement=-8, "
-       "index_Register='RSI', scaling=1)"},
-      {AddressTuple(/* base_register = */ "RAX",
-                    /* displacement = */ -123,
-                    /* index_register = */ "RSI",
-                    /* scaling = */ 1,
-                    /* segment_register = */ "ES"),
-       "AddressTuple(base_register='RAX', displacement=-123, "
-       "index_Register='RSI', scaling=1, segment_register='ES')"}};
+  } kTestCases[] = {{AddressTuple(/* base_register = */ "RAX",
+                                  /* displacement = */ 0,
+                                  /* index_register = */ "",
+                                  /* scaling = */ 0,
+                                  /* segment_register = */ ""),
+                     "AddressTuple(base_register='RAX', base_register_size=64, "
+                     "base_register_intefered_register={})"},
+                    {AddressTuple(/* base_register = */ "RAX",
+                                  /* displacement = */ 16,
+                                  /* index_register = */ "",
+                                  /* scaling = */ 0,
+                                  /* segment_register = */ ""),
+                     "AddressTuple(base_register='RAX', base_register_size=64, "
+                     "base_register_intefered_register={}, displacement=16)"},
+                    {AddressTuple(/* base_register = */ "RAX",
+                                  /* displacement = */ 0,
+                                  /* index_register = */ "RSI",
+                                  /* scaling = */ 0,
+                                  /* segment_register = */ ""),
+                     "AddressTuple(base_register='RAX', base_register_size=64, "
+                     "base_register_intefered_register={}, "
+                     "index_Register='RSI', index_register_size=64, "
+                     "index_register_intefered_register={}, scaling=0)"},
+                    {AddressTuple(/* base_register = */ "RAX",
+                                  /* displacement = */ -8,
+                                  /* index_register = */ "RSI",
+                                  /* scaling = */ 1,
+                                  /* segment_register = */ ""),
+                     "AddressTuple(base_register='RAX', base_register_size=64, "
+                     "base_register_intefered_register={}, displacement=-8, "
+                     "index_Register='RSI', index_register_size=64, "
+                     "index_register_intefered_register={}, scaling=1)"},
+                    {AddressTuple(/* base_register = */ "RAX",
+                                  /* displacement = */ -123,
+                                  /* index_register = */ "RSI",
+                                  /* scaling = */ 1,
+                                  /* segment_register = */ "ES"),
+                     "AddressTuple(base_register='RAX', base_register_size=64, "
+                     "base_register_intefered_register={}, displacement=-123, "
+                     "index_Register='RSI', index_register_size=64, "
+                     "index_register_intefered_register={}, scaling=1, "
+                     "segment_register='ES', segment_register_size=64, "
+                     "segment_register_intefered_register={})"}};
 
   for (const auto& test_case : kTestCases) {
     SCOPED_TRACE(test_case.expected_string);
@@ -287,7 +297,8 @@ TEST(InstructionOperandTest, ToString) {
       {InstructionOperand::FpImmediateValue(3.14),
        "InstructionOperand.from_fp_immediate_value(3.14)"},
       {InstructionOperand::Address("RAX", 0, "", 0, ""),
-       "InstructionOperand.from_address(AddressTuple(base_register='RAX'))"},
+       "InstructionOperand.from_address(AddressTuple(base_register='RAX', "
+       "base_register_size=3, base_register_intefered_register={}))"},
       {InstructionOperand::MemoryLocation(32),
        "InstructionOperand.from_memory(32)"}};
 

gematria/basic_block/python/basic_block.cc

@@ -156,15 +156,17 @@ PYBIND11_MODULE(basic_block, m) {
                   py::arg("fp_immediate_value"))
       .def_static("from_virtual_register",
                   &InstructionOperand::VirtualRegister,
-                  py::arg("register_name"), py::arg("size") = 0)
+                  py::arg("register_name"), py::arg("size"), py::arg("interfered_registers"))
       .def_static<InstructionOperand (*)(
           std::string /* base_register */, int64_t /* displacement */,
           std::string /* index_register */, int /* scaling */,
-          std::string /* segment_register */)>(
+          std::string /* segment_register */, int /* base_register_size */,
+          int /* index_register_size */, int /* segment_register_size */)>(
           "from_address", &InstructionOperand::Address,
           py::arg("base_register") = std::string(), py::arg("displacement") = 0,
           py::arg("index_register") = std::string(), py::arg("scaling") = 0,
-          py::arg("segment_register") = std::string())
+          py::arg("segment_register") = std::string(), py::arg("base_register_size") = 64,
+          py::arg("index_register_size") = 64, py::arg("segment_register_size") = 64)
       .def_static<InstructionOperand (*)(AddressTuple)>(
           "from_address", &InstructionOperand::Address,
           py::arg("address_tuple"))