feat: add subregisters config and allow mixed-precision fpu computations

Author: arromanoff

include/bleach/lifter/instr-impl.hpp

@@ -52,6 +52,7 @@ class instr_impl final : private std::vector<instruction> {
   std::string stack_pointer;
   std::vector<constant_reg> const_regs;
   std::vector<regclass> regclasses;
+  std::unordered_map<std::string, std::string> subregisters;
 
 public:
   instr_impl() = default;
@@ -64,13 +65,16 @@ class instr_impl final : private std::vector<instruction> {
   using vector::emplace_back;
   using vector::push_back;
 
-  StringRef get_stack_pointer() const { return stack_pointer; }
+  StringRef get_stack_pointer() const & { return stack_pointer; }
 
-  auto &get_const_regs() { return const_regs; }
-  auto &get_const_regs() const { return const_regs; }
+  auto &get_const_regs() & { return const_regs; }
+  auto &get_const_regs() const & { return const_regs; }
 
-  auto &get_regclasses() const { return regclasses; }
-  auto &get_regclasses() { return regclasses; }
+  auto &get_regclasses() const & { return regclasses; }
+  auto &get_regclasses() & { return regclasses; }
+
+  auto &get_subregs() const & { return subregisters; }
+  auto &get_subregs() & { return subregisters; }
 
   void set_stack_pointer(std::string sp) { stack_pointer = std::move(sp); }
 

include/bleach/lifter/lifter.hpp

@@ -9,7 +9,9 @@
 #include <llvm/Support/Regex.h>
 #include <llvm/Target/TargetMachine.h>
 
+#include <iostream>
 #include <map>
+#include <set>
 #include <unordered_set>
 
 namespace llvm {
@@ -56,24 +58,25 @@ class mbb2bb final
   }
 };
 
-class register_class final : private std::unordered_set<unsigned> {
+class register_class final : private std::map<unsigned, std::vector<unsigned>> {
   std::string name;
   llvm::Regex regex;
   const TargetRegisterClass *rclass = nullptr;
   unsigned reg_bitsize = 0;
 
 public:
-  using unordered_set::begin;
-  using unordered_set::contains;
-  using unordered_set::empty;
-  using unordered_set::end;
-  using unordered_set::insert;
-  using unordered_set::size;
+  using map::at;
+  using map::begin;
+  using map::contains;
+  using map::empty;
+  using map::end;
+  using map::insert;
+  using map::size;
 
   register_class(std::string &&rcname, llvm::Regex &&rx)
       : name(std::move(rcname)), regex(std::move(rx)) {}
 
-  void add_reg(unsigned reg) { insert(reg); }
+  void add_reg(unsigned reg) { try_emplace(reg); }
 
   auto &get_regex() const & { return regex; }
 
@@ -104,8 +107,14 @@ class register_stats final : std::vector<register_class> {
   }
 
   auto &get_register_class_for(unsigned reg) const {
-    auto found =
-        find_if(*this, [reg](auto &rclass) { return rclass.contains(reg); });
+    auto found = ranges::find_if(*this, [reg](auto &rclass) {
+      if (rclass.contains(reg))
+        return true;
+      auto subreg = ranges::find_if(rclass, [reg](auto &entry) {
+        return is_contained(entry.second, reg);
+      });
+      return subreg != rclass.end();
+    });
     if (found == end()) {
       throw std::invalid_argument("Unknown register encountered");
     }

lib/lifter/instr-impl.cpp

@@ -7,6 +7,7 @@
 #include <llvm/Support/SourceMgr.h>
 #include <llvm/Support/raw_ostream.h>
 
+#include <stdexcept>
 #include <yaml-cpp/emitter.h>
 #include <yaml-cpp/emittermanip.h>
 #include <yaml-cpp/yaml.h>
@@ -134,6 +135,12 @@ instr_impl load_from_yaml(std::string yaml, llvm::LLVMContext &ctx) {
   for (auto &&node : instrs_conf["register-classes"])
     instrs.get_regclasses().push_back(
         {node.first.as<std::string>(), node.second.as<std::string>()});
+  for (auto &&node : instrs_conf["subregisters"]) {
+    auto [it, inserted] = instrs.get_subregs().try_emplace(
+        node.first.as<std::string>(), node.second.as<std::string>());
+    if (!inserted)
+      throw std::invalid_argument("duplicated subregister class encountered");
+  }
   for (auto &&node : instrs_conf["instructions"]) {
     auto norm = node.as<normalized_instruction>();
     instrs.push_back(denormalize_instruction(norm, ctx));

lib/lifter/lifter.cpp

@@ -4,6 +4,7 @@
 
 #include "symtab-ir.h"
 
+#include <algorithm>
 #include <iterator>
 #include <llvm/CodeGen/MachineFunction.h>
 #include <llvm/CodeGen/MachineModuleInfo.h>
@@ -173,7 +174,8 @@ void materialize_registers(MachineFunction &mf, Function &func, reg2vals &rmap,
                                         ArrayRef<Value *>{const_zero, arr_idx},
                                         rclass.get_name());
     auto sorted_regs = std::set<unsigned>();
-    ranges::copy(rclass, std::inserter(sorted_regs, sorted_regs.end()));
+    ranges::transform(rclass, std::inserter(sorted_regs, sorted_regs.end()),
+                      [](auto &&entry) { return entry.first; });
     for (auto &&[reg_idx, reg] : sorted_regs | views::enumerate) {
       auto *array_reg_idx = ConstantInt::get(ctx, APInt(32, reg_idx));
       auto *reg_src_addr = builder.CreateInBoundsGEP(
@@ -186,6 +188,10 @@ void materialize_registers(MachineFunction &mf, Function &func, reg2vals &rmap,
           Type::getIntNTy(ctx, rclass.get_register_size()), reg_src_addr);
       builder.CreateStore(reg_val, reg_dst_addr);
       rmap.try_emplace(reg, reg_dst_addr);
+      auto &&subregs = rclass.at(reg);
+      ranges::for_each(subregs, [&](std::integral auto r) {
+        rmap.try_emplace(r, reg_dst_addr);
+      });
     }
   }
 }
@@ -430,8 +436,9 @@ void assign_register_classes(const TargetSubtargetInfo &stinfo,
   assert(rinfo);
   for (auto &&[idx, reg_class] : enumerate(stats)) {
     auto found = find_if(rinfo->regclasses(), [&reg_class](auto &rclass) {
-      return all_of(reg_class,
-                    [&rclass](auto reg) { return is_contained(*rclass, reg); });
+      return all_of(reg_class, [&rclass](auto &&reg) {
+        return is_contained(*rclass, reg.first);
+      });
     });
     if (found == rinfo->regclass_end()) {
       throw std::runtime_error(
@@ -453,17 +460,48 @@ register_stats collect_register_stats(const instr_impl &instr, Module &m,
     throw std::runtime_error(
         "register-classes were not specified in input YAML");
   register_stats stats(rclasses.begin(), rclasses.end());
+  const MCRegisterInfo *rinfo = nullptr;
+  auto *stinfo = mmi.getTarget().getSubtargetImpl(*m.begin());
   for (auto &f : m) {
     auto &mf = mmi.getOrCreateMachineFunction(f);
     collect_register_stats_for(mf, mmi.getTarget(), stats);
+    rinfo = stinfo->getRegisterInfo();
   }
   assert(!m.empty());
-  auto *stinfo = mmi.getTarget().getSubtargetImpl(*m.begin());
   assert(stinfo);
   assign_register_classes(*stinfo, stats);
   stats.erase(std::remove_if(stats.begin(), stats.end(),
                              [](auto &rclass) { return rclass.empty(); }),
               stats.end());
+  for (auto &&[subregclass, aliasto] : instr.get_subregs()) {
+    auto found = ranges::find_if(stats, [&subregclass](auto &rclass) {
+      return rclass.get_name() == subregclass;
+    });
+    if (found != stats.end()) {
+      auto found_aliasto = ranges::find_if(stats, [&aliasto](auto &rclass) {
+        return rclass.get_name() == aliasto;
+      });
+      if (found_aliasto != stats.end()) {
+        for (auto &&[entry, subreg] : views::zip(*found_aliasto, *found)) {
+          auto &&[_, subregs] = entry;
+          subregs.push_back(subreg.first);
+        }
+        stats.erase(found);
+      }
+    }
+  }
+
+  for (auto &&regclass : stats) {
+    std::cerr << "\nREGCLASS: " << regclass.get_name() << '\n';
+    for (auto &&r : regclass) {
+      std::cerr << rinfo->getName(r.first) << ": [ ";
+      for (auto subreg : r.second)
+        std::cerr << rinfo->getName(subreg) << " ";
+      std::cerr << "]\n";
+    }
+    std::cerr << '\n';
+  }
+
   return stats;
 }
 

test/integration/inputs/riscv-mixed-fp.c

@@ -1,4 +1,5 @@
 #include <STATE>
+#include <math.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <string.h>
@@ -29,6 +30,14 @@ double ref_calc1(double a, double b, int x) {
 
 double reference(double a, double b, int x) { return ref_calc1(a, b, x); }
 
+int are_equal(double a, double b) {
+  if (isnan(a) && isnan(b))
+    return 1;
+  if (isnan(a) || isnan(b))
+    return 0;
+  return fabs(a - b) <= 0.1;
+}
+
 int main() {
   printf("Hello from main\n");
   // random values
@@ -43,5 +52,6 @@ int main() {
   double res2 = bitcast_to_double(regs.FPR[10]);
   printf("result: %lf\n", res2);
   printf("reference: %lf\n", res1);
-  return res1 - res2 > 0.1;
+  printf("diff: %lf\n", res1 - res2);
+  return !are_equal(res1, res2);
 }

test/integration/inputs/riscv-mixed-precision-fp.c

@@ -0,0 +1,57 @@
+#include <STATE>
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <string.h>
+
+struct register_state regs = {};
+
+int64_t bitcast_to_int(double v) {
+  int64_t res;
+  memcpy(&res, &v, sizeof(double));
+  return res;
+}
+
+double bitcast_to_double(int64_t v) {
+  double res;
+  memcpy(&res, &v, sizeof(double));
+  return res;
+}
+
+// This should be in sync with test/integration/riscv-mixed-fp.c
+// TODO: automate this
+double ref_calc2(float prod, double sum, int x) {
+  return prod + sum - (double)x;
+}
+
+double ref_calc1(float a, float b, int x) {
+  return ref_calc2(a * b, (a + b), x + a / b);
+}
+
+double reference(double a, double b, int x) { return ref_calc1(a, b, x); }
+
+int are_equal(double a, double b) {
+  if (isnan(a) && isnan(b))
+    return 1;
+  if (isnan(a) || isnan(b))
+    return 0;
+  return fabs(a - b) <= 0.1;
+}
+
+int main() {
+  // random values that cause difference between computation in floats and
+  // doubles
+  double arg_1 = 4113444444434.33;
+  double arg_2 = 0.0000003344;
+  int arg_3 = 5;
+  regs.FPR[10] = bitcast_to_int(arg_1);
+  regs.FPR[11] = bitcast_to_int(arg_2);
+  regs.GPR[10] = arg_3;
+  top_small(&regs);
+  double res1 = reference(arg_1, arg_2, arg_3);
+  double res2 = bitcast_to_double(regs.FPR[10]);
+  printf("result: %lf\n", res2);
+  printf("reference: %lf\n", res1);
+  printf("diff: %lf\n", res1 - res2);
+  return !are_equal(res1, res2);
+}

test/integration/riscv-mixed-precision-fp.c

@@ -0,0 +1,18 @@
+// RUN: %config-gen-path/config-gen.rb --march rv64imfd \
+// RUN:   --template-dir %config-gen-path/templates -o %t.yaml
+// RUN: riscv64-unknown-linux-gnu-clang %s -O2 -c -o %t.out -march=rv64imfd
+// RUN: %bin/llvm-bleach %t.out --instructions %t.yaml -o %t.ll \
+// RUN:   --state-struct-file=%t.state.h
+// RUN: sed 's|STATE|%t.state.h|g' %S/inputs/riscv-mixed-precision-fp.c > \
+// RUN:   %t.main.c
+// RUN: clang %t.main.c %t.ll -o %t.native.out -g \
+// RUN:   -fsanitize=address,undefined
+// RUN: %t.native.out
+
+double calc2(float prod, double sum, int x) { return prod + sum - (double)x; }
+
+double calc1(float a, float b, int x) {
+  return calc2(a * b, (a + b), x + a / b);
+}
+
+double top_small(double a, double b, int x) { return calc1(a, b, x); }