split critical edges during x86 register allocation

AGENTS.md

@@ -25,39 +25,55 @@ metajit.cpp uses a generating extension for tracing.
 
 ## Project Structure
 
+- `README.md`: Project overview, architecture, and usage examples.
+- `Makefile`: Build, test, fuzzing, header generation, and coverage targets.
 - `main.cpp`: A small example program using metajit.cpp.
 - `jitir.py`: A Python script that generates JITIR-related headers from templates. It also includes all instruction definitions.
-- `interactive.hpp`: A ncurses based interactive debugger which allows a user to single step through JITIR code. The debugger is meant to be used by humans. You should not use it.
+- `interactive.hpp`: An interactive debugger that uses ncurses and allows a user to step through JITIR code. The debugger is meant to be used by humans. You should not use it.
 - `jitir.tmpl.hpp`: Template for generating the jitir.hpp header file, which defines the JITIR (Intermediate Representation) and its operations.
-- `jitir_llvmapi.tmpl.hpp`: Template for generating the jitir_llvmapi.hpp header file, which provides 
+- `jitir_llvmapi.tmpl.hpp`: Template for generating the jitir_llvmapi.hpp header file, which provides LLVM declarations for the JITIR builder API.
+- `genext.tmpl.hpp`: Template for generating genext.hpp, which implements recording, trace capability analysis, and generating extension construction.
 - `llvmgen.hpp`: Generates LLVM IR from JITIR and emits the generating extension which is used for tracing.
 - `x86gen.hpp`: Our custom x86 backend. LLVM's x86 backend is too slow for our use case, so we generate x86 code directly from JITIR.
 - `x86insts.inc.hpp`: x86 instruction definitions and encodings used by x86gen.hpp.
 - `lowerllvm.hpp`: Lowers LLVM IR to JITIR.
+- `layout2aliasing.hpp`: Defines memory layouts and derives aliasing information from them.
 - `tv.hpp`: Translation validation for JITIR using Z3. This allows us to prove refinement between JITIR programs.
 - `doc/`: Documentation, including JITIR instruction documentation generated from templates.
-- `tests/`: Unit tests
+- `tests/`: Unit tests.
   - `diff.hpp`: Differential testing. We compare the x86 backend, LLVM backend, and the interpreter against each other to find bugs.
   - `output/`: Output files from the tests. This includes JITIR and assembly output from the backends.
   - `fuzzer.cpp`: A fuzzer that generates random JITIR programs and checks that the backends and interpreter all produce the same output.
-  - `test_cfg.cpp`: Tests with control flow
-  - `test_insts.cpp`: Tests for individual instructions
+  - `coverage_to_markdown.py`: Converts LLVM coverage summaries into Markdown tables for CI reports.
+  - `test_cfg.cpp`: Tests with control flow.
+  - `test_clone.cpp`: Tests for cloning JITIR instructions and sections.
+  - `test_genext.cpp`: Tests for generating extensions, including direct tracing and recording and replay.
+  - `test_insts.cpp`: Tests for individual instructions.
+  - `test_interpreter.cpp`: Tests for the JITIR interpreter, including poison propagation.
   - `test_fuzzer.cpp`: Regression tests for all bugs previously found by the fuzzer.
   - `test_knownbits.cpp`: Tests for knownbits analysis, which is used by the `Simplify` pass.
+  - `test_mem2reg.cpp`: Tests for promoting memory operations to SSA values and block parameters.
   - `test_opt.cpp`: Tests for optimization passes.
+  - `test_reader.cpp`: Tests for reading textual JITIR and folding instructions while reading.
+  - `test_reentry.cpp`: Tests for reentry closures, captured values, and control-flow slicing.
+  - `test_source.cpp`: Tests for compiling LLVM IR generated from source fixtures.
   - `test_tv.cpp`: Tests for translation validation using Z3.
+  - `source/`: Source fixtures compiled to LLVM IR for source tests.
+- `.github/workflows/`: GitHub Actions workflows for tests, fuzzing, and coverage reports.
 
 ## Coding Guidelines
 
+- Do not add code comments.
+- Name feature branches `<author-first-name>-<short-description>`, using the author's first name followed by a dash and a short description with words separated by dashes.
 - Never edit the generated jitir.hpp and jitir_llvmapi.hpp files directly. Instead, edit the corresponding template files jitir.tmpl.hpp and jitir_llvmapi.tmpl.hpp. The instructions are specified in the jitir.py generator script.
-- Never edit any files in tests/output. They are just output files from the unit-tests used to debug failing test cases. They are also not golden tests; in fact they are .gitnored. 
-- metajit.cpp is a just in time compiler. This makes compile time a crucial metric to optimize for. Write performant compiler code and avoid allocations where possible.
+- Never edit any files in tests/output. They are just output files from the unit tests used to debug failing test cases. They are also not golden tests; in fact, they are ignored by Git.
+- metajit.cpp is a JIT compiler. This makes compile time a crucial metric to optimize for. Write performant compiler code and avoid allocations where possible.
 - Always try to keep changes scoped and reviewable. Attempt to build the minimal change needed to implement a feature or fix a bug, and avoid making unrelated formatting or refactoring changes in the same commit.
 
 ## Testing
 
 Make automatically regenerates the JITIR headers.
 
 - You can run the entire test suite using `make test`.
-- You can run individual tests using `make tests/test_* && ./tests/test_*`
+- You can run individual tests using `make tests/test_* && ./tests/test_*`.
 - You can run the fuzzer using `make fuzz`.

jitir.tmpl.hpp

@@ -5287,10 +5287,10 @@ namespace metajit {
         } else {
           Block* succ = frame.successors.back();
           frame.successors.pop_back();
-          if (_dt.dominates(succ, frame.block)) {
-            _seen_loop = true;
-          } else if (_visited.insert(succ).second) {
+          if (_visited.insert(succ).second) {
             stack.push_back({succ, succ->successors()});
+          } else if (_dt.dominates(succ, frame.block)) {
+            _seen_loop = true;
           }
         }
       }

tests/test_cfg.cpp

@@ -61,6 +61,71 @@ int main(int argc, char** argv) {
 
   });
 
+  suite.diff_test("critical_edge").aot(false).run([](Builder& builder, TestData& data) {
+    Block* a = builder.build_block();
+    Block* b = builder.build_block();
+    Block* other = builder.build_block();
+    Block* merge = builder.build_block();
+
+    Value* first_cond = data.input(Type::Bool);
+    Value* second_cond = data.input(Type::Bool);
+    std::vector<Value*> values;
+    for (size_t index = 0; index < 24; index++) {
+      values.push_back(data.input(Type::Int64));
+    }
+    builder.build_branch(first_cond, a, b);
+
+    builder.move_to_end(a);
+    builder.build_jump(merge);
+
+    builder.move_to_end(b);
+    for (Value* value : values) {
+      data.output(builder.build_add(value, builder.build_const(Type::Int64, 1)));
+    }
+    builder.build_branch(second_cond, merge, other);
+
+    builder.move_to_end(other);
+    data.output(builder.build_const(Type::Int64, 42));
+    builder.build_jump(merge);
+
+    builder.move_to_end(merge);
+    for (Value* value : values) {
+      data.output(value);
+    }
+  });
+
+  suite.diff_test("critical_backedge").aot(false).run([](Builder& builder, TestData& data) {
+    Block* header = builder.build_block();
+    Block* body = builder.build_block();
+    Block* end = builder.build_block();
+
+    Value* count = data.input(RandomRange(Type::Int64, 1, 8));
+    Value* counter = builder.build_alloca(builder.build_const(Type::Int64, 8), 8);
+    builder.build_store(counter, count, AliasingGroup(0), 0);
+    std::vector<Value*> values;
+    for (size_t index = 0; index < 24; index++) {
+      values.push_back(data.input(Type::Int64));
+    }
+    builder.build_jump(header);
+
+    builder.move_to_end(header);
+    Value* current = builder.build_load(counter, Type::Int64, LoadFlags::None, AliasingGroup(0), 0);
+    builder.build_jump(body);
+
+    builder.move_to_end(body);
+    Value* next = builder.build_sub(current, builder.build_const(Type::Int64, 1));
+    builder.build_store(counter, next, AliasingGroup(0), 0);
+    for (Value* value : values) {
+      data.output(builder.build_add(value, next));
+    }
+    builder.build_branch(builder.build_lt_u(builder.build_const(Type::Int64, 0), next), header, end);
+
+    builder.move_to_end(end);
+    for (Value* value : values) {
+      data.output(value);
+    }
+  });
+
   suite.diff_test("sum_to").run([](Builder& builder, TestData& data) {
     Block* loop_header = builder.build_block({Type::Int64, Type::Int64}); // (i, sum)
     Block* loop_body = builder.build_block();

x86gen.hpp

@@ -1720,7 +1720,8 @@ namespace metajit {
       }
       _blocks[0]->set_regalloc(initial_state);
 
-      for (X86Block* block : _blocks) {
+      for (size_t block_index = 0; block_index < _blocks.size(); block_index++) {
+        X86Block* block = _blocks[block_index];
         if (block->regalloc()) {
           load_state(reg_file, block->regalloc());
         }
@@ -1848,8 +1849,21 @@ namespace metajit {
           if (std::holds_alternative<X86Block*>(inst->imm())) {
             X86Block* target = std::get<X86Block*>(inst->imm());
             if (target->regalloc()) {
+              if (inst->kind() != X86Inst::Kind::Jmp) {
+                X86Block* edge = _builder.build_block();
+                edge->set_name(block->name());
+                Reg* state = (Reg*) _allocator.alloc(sizeof(Reg) * reg_file.size(), alignof(Reg));
+                for (size_t it = 0; it < reg_file.size(); it++) {
+                  state[it] = reg_file[Reg::phys(it)];
+                }
+                edge->set_regalloc(state);
+                _builder.move_before(edge, nullptr);
+                _builder.jmp(target)->set_name(inst->name());
+                _blocks.insert(_blocks.begin() + block_index + 1, edge);
+                inst->set_imm(edge);
+                continue;
+              }
               // Restore regalloc state
-              assert(inst->kind() == X86Inst::Kind::Jmp); // Merges may only be unconditional jumps
               if (target->name() < block->name()) {
                 // Backedge
                 assert(target->loop());
@@ -1906,6 +1920,10 @@ namespace metajit {
         }
       }
 
+      for (size_t it = 0; it < _blocks.size(); it++) {
+        _blocks[it]->set_name(it);
+      }
+
       #ifdef METAJIT_DEBUG
       for (X86Block* block : _blocks) {
         if (block->regalloc()) {