Add support for using regalloc3 as the allocator back-end

Author: Amanieu

.github/workflows/rust.yml

@@ -14,6 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
+    - uses: dtolnay/rust-toolchain@stable
     - name: Install rustfmt
       run: rustup component add rustfmt
     - name: Run rustfmt and check there's no difference
@@ -24,6 +25,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
+    - uses: dtolnay/rust-toolchain@stable
     - name: Build
       run: cargo build
     - name: Run tests
@@ -34,6 +36,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
+    - uses: dtolnay/rust-toolchain@stable
     - name: Check with all features
       run: cargo check --all-features
 
@@ -42,6 +45,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
+    - uses: dtolnay/rust-toolchain@stable
     - name: Install thumbv6m-none-eabi target
       run: rustup target add thumbv6m-none-eabi
     - name: Check no_std build

Cargo.toml

@@ -8,7 +8,7 @@ authors = [
     "Chris Fallin <[email protected]>",
     "Mozilla SpiderMonkey Developers",
 ]
-edition = "2018"
+edition = "2021"
 license = "Apache-2.0 WITH LLVM-exception"
 description = "Backtracking register allocator inspired from IonMonkey"
 repository = "https://github.com/bytecodealliance/regalloc2"
@@ -18,6 +18,7 @@ log = { version = "0.4.8", default-features = false }
 smallvec = { version = "1.6.1", features = ["union"] }
 rustc-hash = { version = "2.0.0", default-features = false }
 hashbrown = { version = "0.15", default-features = false, features = [] }
+regalloc3 = { git = "https://github.com/Amanieu/regalloc3.git" }
 
 # Optional serde support, enabled by feature below.
 serde = { version = "1.0.136", features = [
@@ -46,7 +47,7 @@ std = []
 checker = []
 
 # Enables detailed logging which can be somewhat expensive.
-trace-log = []
+trace-log = ["regalloc3/trace-log"]
 
 # Exposes the internal API for fuzzing.
 fuzzing = ["libfuzzer-sys", "checker", "trace-log"]

fuzz/Cargo.toml

@@ -55,6 +55,12 @@ path = "fuzz_targets/fastalloc_checker.rs"
 test = false
 doc = false
 
+[[bin]]
+name = "ra3_checker"
+path = "fuzz_targets/ra3_checker.rs"
+test = false
+doc = false
+
 # Enable debug assertions and overflow checks when fuzzing
 [profile.release]
 debug = true

fuzz/fuzz_targets/ra3_checker.rs

@@ -0,0 +1,59 @@
+/*
+ * Released under the terms of the Apache 2.0 license with LLVM
+ * exception. See `LICENSE` for details.
+ */
+
+#![no_main]
+use regalloc2::fuzzing::arbitrary::{Arbitrary, Result, Unstructured};
+use regalloc2::fuzzing::checker::Checker;
+use regalloc2::fuzzing::func::{Func, Options};
+use regalloc2::fuzzing::fuzz_target;
+
+#[derive(Clone, Debug)]
+struct TestCase {
+    func: Func,
+}
+
+impl Arbitrary<'_> for TestCase {
+    fn arbitrary(u: &mut Unstructured) -> Result<TestCase> {
+        Ok(TestCase {
+            func: Func::arbitrary_with_options(
+                u,
+                &Options {
+                    reused_inputs: true,
+                    fixed_regs: true,
+                    fixed_nonallocatable: true,
+                    clobbers: true,
+                    reftypes: true,
+                },
+            )?,
+        })
+    }
+}
+
+fuzz_target!(|testcase: TestCase| {
+    let func = testcase.func;
+    let _ = env_logger::try_init();
+    log::trace!("func:\n{:?}", func);
+    let env = regalloc2::fuzzing::func::machine_env();
+
+    thread_local! {
+        // We test that ctx is cleared properly between runs.
+        static CTX: std::cell::RefCell<regalloc2::fuzzing::ion::Ctx> = std::cell::RefCell::default();
+    }
+
+    CTX.with(|ctx| {
+        let options = regalloc2::RegallocOptions {
+            verbose_log: true,
+            validate_ssa: true,
+            algorithm: regalloc2::Algorithm::Regalloc3,
+        };
+        let mut ctx = ctx.borrow_mut();
+        let out = regalloc2::run_with_ctx(&func, &env, &options, &mut *ctx)
+            .expect("regalloc did not succeed");
+
+        let mut checker = Checker::new(&func, &env);
+        checker.prepare(&out);
+        checker.run().expect("checker failed");
+    });
+});

src/fuzzing/func.rs

@@ -313,7 +313,10 @@ impl Func {
         let mut out_blocks = vec![];
         let mut in_blocks = vec![];
         while from < num_blocks {
-            in_blocks.push(from);
+            // regalloc3 doesn't allow the entry block to have predecessors.
+            if from != 0 {
+                in_blocks.push(from);
+            }
             if num_blocks > 3 && from < num_blocks - 3 && bool::arbitrary(u)? {
                 // To avoid critical edges, we use from+1 as an edge
                 // block, and advance `from` an extra block; `from+2`

src/fuzzing/mod.rs

@@ -27,6 +27,9 @@ pub mod ion {
 pub mod fastalloc {
     pub use crate::fastalloc::*;
 }
+pub mod regalloc3 {
+    pub use crate::regalloc3::*;
+}
 pub mod checker {
     pub use crate::checker::*;
 }

src/ion/data_structures.rs

@@ -16,6 +16,7 @@ use super::liveranges::SpillWeight;
 use crate::cfg::{CFGInfo, CFGInfoCtx};
 use crate::index::ContainerComparator;
 use crate::indexset::IndexSet;
+use crate::regalloc3::Regalloc3Ctx;
 use crate::Vec2;
 use crate::{
     define_index, Allocation, Block, Bump, Edit, Function, FxHashMap, FxHashSet, MachineEnv,
@@ -489,6 +490,8 @@ pub struct Ctx {
     pub(crate) scratch_workqueue_set: FxHashSet<Block>,
 
     pub(crate) scratch_bump: Bump,
+
+    pub(crate) ra3_ctx: Regalloc3Ctx,
 }
 
 impl Ctx {

src/lib.rs

@@ -50,6 +50,7 @@ pub mod indexset;
 pub(crate) mod ion;
 pub(crate) mod moves;
 pub(crate) mod postorder;
+pub(crate) mod regalloc3;
 pub mod ssa;
 
 #[macro_use]
@@ -1598,7 +1599,7 @@ pub fn run<F: Function>(
     options: &RegallocOptions,
 ) -> Result<Output, RegAllocError> {
     match options.algorithm {
-        Algorithm::Ion => {
+        Algorithm::Ion | Algorithm::Regalloc3 => {
             let mut ctx = Ctx::default();
             run_with_ctx(func, env, options, &mut ctx)?;
             Ok(ctx.output)
@@ -1623,6 +1624,13 @@ pub fn run_with_ctx<'a, F: Function>(
         Algorithm::Fastalloc => {
             ctx.output = fastalloc::run(func, env, options.verbose_log, options.validate_ssa)?
         }
+        Algorithm::Regalloc3 => ctx.ra3_ctx.run(
+            func,
+            env,
+            &mut ctx.cfginfo,
+            &mut ctx.cfginfo_ctx,
+            &mut ctx.output,
+        )?,
     }
     Ok(&ctx.output)
 }
@@ -1632,6 +1640,7 @@ pub enum Algorithm {
     #[default]
     Ion,
     Fastalloc,
+    Regalloc3,
 }
 
 /// Options for allocation.