[WIP / RFC] Cranelift: Basic support for EGraph roundtripping.

This is a work-in-progress, and meant to sketch the direction I've been
thinking in for a mid-end framework. A proper BA RFC will come soon.

This PR builds a phase in the optimization pipeline that converts a CLIF
CFG into an egraph representing the function body. Each node represents
an original instruction or operator. The "skeleton" of side-effecting
instructions is retained, but non-side-effecting (pure) operators are
allowed to "float": the egraph will naturally deduplicate them during
build, and we will determine their proper place when we convert back to
a CFG representation.

The conversion from the egraph back to the CFG is done via a new
algorithm I call "scoped elaboration". The basic idea is to do a
preorder traversal of the domtree, and at each level, evaluate the
values of the eclasses called upon by the side-effect skeleton,
memoizing in an eclass-to-SSA-value map. This map is a scoped hashmap,
with scopes at each domtree level. In this way, (i) when a value is
computed in a location that dominates another instance of that value,
the first replacees the second; but (ii) we never produce "partially
dead" computations, i.e. we never hoist to a level in the domtree where
a node is not "anticipated" (always eventually computed). This exactly
matches what GVN does today. With a small tweak, it can also subsume
LICM: we need to be loop-nest-aware in our recursive eclass elaboration,
and potentially place nodes higher up the domtree (and higher up in the
scoped hashmap).

Unlike what I had been thinking in Monday's meeting, this produces CLIF
out of the egraph and then allows that to be lowered. It's overall
simpler and a better starting point (thanks @abrown for tipping me over
the edge in this). The way it produces CLIF now could be made more
efficient: it could reuse instructions already in the DFG for nodes that
are *not* duplicated (likely most of them) rather than clearing all and
repopulating.

This PR does *not* do anything to actually rewrite in the egraph. That's
the next step! I need to work out exactly how to integrate ISLE with
some sort of rewrite machinery. I have some ideas about efficient
dispatch with an "operand-tree discriminants shape analysis" on the
egraph and indexing rules by their matched shape; more to come.

Author: cfallin

Cargo.lock

@@ -25,6 +25,7 @@ gimli = { version = "0.26.0", default-features = false, features = ["write"], op
 smallvec = { version = "1.6.1" }
 regalloc2 = { version = "0.2.0", features = ["checker"] }
 souper-ir = { version = "2.1.0", optional = true }
+egg = { git = "https://github.com/cfallin/egg", branch = "egg-lite", default-features = false }
 # It is a goal of the cranelift-codegen crate to have minimal external dependencies.
 # Please don't add any unless they are essential to the task of creating binary
 # machine code. Integration tests that need external dependencies can be

cranelift/codegen/Cargo.toml

@@ -60,36 +60,51 @@ fn gen_formats(formats: &[&InstructionFormat], fmt: &mut Formatter) {
     fmt.empty_line();
 }
 
-/// Generate the InstructionData enum.
+/// Generate the InstructionData and InstructionImms enums.
+///
+/// `InstructionImms` stores everything about an instruction except for the arguments: in other
+/// words, the `Opcode` and any immediates or other parameters. `InstructionData` stores this, plus
+/// the SSA `Value` arguments.
 ///
 /// Every variant must contain an `opcode` field. The size of `InstructionData` should be kept at
 /// 16 bytes on 64-bit architectures. If more space is needed to represent an instruction, use a
 /// `ValueList` to store the additional information out of line.
 fn gen_instruction_data(formats: &[&InstructionFormat], fmt: &mut Formatter) {
-    fmt.line("#[derive(Clone, Debug)]");
-    fmt.line(r#"#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]"#);
-    fmt.line("#[allow(missing_docs)]");
-    fmt.line("pub enum InstructionData {");
-    fmt.indent(|fmt| {
-        for format in formats {
-            fmtln!(fmt, "{} {{", format.name);
-            fmt.indent(|fmt| {
-                fmt.line("opcode: Opcode,");
-                if format.has_value_list {
-                    fmt.line("args: ValueList,");
-                } else if format.num_value_operands == 1 {
-                    fmt.line("arg: Value,");
-                } else if format.num_value_operands > 0 {
-                    fmtln!(fmt, "args: [Value; {}],", format.num_value_operands);
-                }
-                for field in &format.imm_fields {
-                    fmtln!(fmt, "{}: {},", field.member, field.kind.rust_type);
-                }
-            });
-            fmtln!(fmt, "},");
+    for (name, include_args) in &[("InstructionData", true), ("InstructionImms", false)] {
+        fmt.line("#[derive(Clone, Debug)]");
+        if !include_args {
+            // `InstructionImms` gets some extra derives: it acts like a sort of extended opcode
+            // and we want to allow for hashconsing, sorting, and the like.
+            fmt.line("#[derive(Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]");
         }
-    });
-    fmt.line("}");
+        fmt.line(r#"#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]"#);
+        fmt.line("#[allow(missing_docs)]");
+        // generate `enum InstructionData` or `enum InstructionImms`.
+        // (This comment exists so one can grep for `enum InstructionData`!)
+        fmtln!(fmt, "pub enum {} {{", name);
+        fmt.indent(|fmt| {
+            for format in formats {
+                fmtln!(fmt, "{} {{", format.name);
+                fmt.indent(|fmt| {
+                    fmt.line("opcode: Opcode,");
+                    if *include_args {
+                        if format.has_value_list {
+                            fmt.line("args: ValueList,");
+                        } else if format.num_value_operands == 1 {
+                            fmt.line("arg: Value,");
+                        } else if format.num_value_operands > 0 {
+                            fmtln!(fmt, "args: [Value; {}],", format.num_value_operands);
+                        }
+                    }
+                    for field in &format.imm_fields {
+                        fmtln!(fmt, "{}: {},", field.member, field.kind.rust_type);
+                    }
+                });
+                fmtln!(fmt, "},");
+            }
+        });
+        fmt.line("}");
+    }
 }
 
 fn gen_arguments_method(formats: &[&InstructionFormat], fmt: &mut Formatter, is_mut: bool) {
@@ -150,6 +165,122 @@ fn gen_arguments_method(formats: &[&InstructionFormat], fmt: &mut Formatter, is_
     fmtln!(fmt, "}");
 }
 
+/// Generate the conversion from `InstructionData` to `InstructionImms`, stripping out the
+/// `Value`s.
+fn gen_instruction_data_to_instruction_imms(formats: &[&InstructionFormat], fmt: &mut Formatter) {
+    fmt.line("impl std::convert::From<&InstructionData> for InstructionImms {");
+    fmt.indent(|fmt| {
+        fmt.doc_comment("Convert an `InstructionData` into an `InstructionImms`.");
+        fmt.line("fn from(data: &InstructionData) -> InstructionImms {");
+        fmt.indent(|fmt| {
+            fmt.line("match data {");
+            fmt.indent(|fmt| {
+                for format in formats {
+                    fmtln!(fmt, "InstructionData::{} {{", format.name);
+                    fmt.indent(|fmt| {
+                        fmt.line("opcode,");
+                        for field in &format.imm_fields {
+                            fmtln!(fmt, "{},", field.member);
+                        }
+                        fmt.line("..");
+                    });
+                    fmtln!(fmt, "}} => InstructionImms::{} {{", format.name);
+                    fmt.indent(|fmt| {
+                        fmt.line("opcode: *opcode,");
+                        for field in &format.imm_fields {
+                            fmtln!(fmt, "{}: {}.clone(),", field.member, field.member);
+                        }
+                    });
+                    fmt.line("},");
+                }
+            });
+            fmt.line("}");
+        });
+        fmt.line("}");
+    });
+    fmt.line("}");
+    fmt.empty_line();
+}
+
+/// Generate the conversion from `InstructionImms` to `InstructionData`, adding the
+/// `Value`s.
+fn gen_instruction_imms_to_instruction_data(formats: &[&InstructionFormat], fmt: &mut Formatter) {
+    fmt.line("impl  InstructionImms {");
+    fmt.indent(|fmt| {
+        fmt.doc_comment("Convert an `InstructionImms` into an `InstructionData` by adding args.");
+        fmt.line(
+            "pub fn with_args(&self, values: &[Value], value_list: &mut ValueListPool) -> InstructionData {",
+        );
+        fmt.indent(|fmt| {
+            fmt.line("match self {");
+            fmt.indent(|fmt| {
+                for format in formats {
+                    fmtln!(fmt, "InstructionImms::{} {{", format.name);
+                    fmt.indent(|fmt| {
+                        fmt.line("opcode,");
+                        for field in &format.imm_fields {
+                            fmtln!(fmt, "{},", field.member);
+                        }
+                    });
+                    fmt.line("} => {");
+                    if format.has_value_list {
+                        fmtln!(fmt, "let args = ValueList::from_slice(values, value_list);");
+                    }
+                    fmt.indent(|fmt| {
+                        fmtln!(fmt, "InstructionData::{} {{", format.name);
+                        fmt.indent(|fmt| {
+                            fmt.line("opcode: *opcode,");
+                            for field in &format.imm_fields {
+                                fmtln!(fmt, "{}: {}.clone(),", field.member, field.member);
+                            }
+                            if format.has_value_list {
+                                fmtln!(fmt, "args,");
+                            } else if format.num_value_operands == 1 {
+                                fmtln!(fmt, "arg: values[0],");
+                            } else if format.num_value_operands > 0 {
+                                let mut args = vec![];
+                                for i in 0..format.num_value_operands {
+                                    args.push(format!("values[{}]", i));
+                                }
+                                fmtln!(fmt, "args: [{}],", args.join(", "));
+                            }
+                        });
+                        fmt.line("}");
+                    });
+                    fmt.line("},");
+                }
+            });
+            fmt.line("}");
+        });
+        fmt.line("}");
+    });
+    fmt.line("}");
+    fmt.empty_line();
+}
+
+/// Generate the `opcode` method on InstructionImms.
+fn gen_instruction_imms_impl(formats: &[&InstructionFormat], fmt: &mut Formatter) {
+    fmt.line("impl InstructionImms {");
+    fmt.indent(|fmt| {
+        fmt.doc_comment("Get the opcode of this instruction.");
+        fmt.line("pub fn opcode(&self) -> Opcode {");
+        fmt.indent(|fmt| {
+            let mut m = Match::new("*self");
+            for format in formats {
+                m.arm(
+                    format!("Self::{}", format.name),
+                    vec!["opcode", ".."],
+                    "opcode".to_string(),
+                );
+            }
+            fmt.add_match(m);
+        });
+        fmt.line("}");
+    });
+    fmt.line("}");
+    fmt.empty_line();
+}
+
 /// Generate the boring parts of the InstructionData implementation.
 ///
 /// These methods in `impl InstructionData` can be generated automatically from the instruction
@@ -406,7 +537,7 @@ fn gen_opcodes(all_inst: &AllInstructions, fmt: &mut Formatter) {
     "#,
     );
     fmt.line("#[repr(u16)]");
-    fmt.line("#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]");
+    fmt.line("#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]");
     fmt.line(
         r#"#[cfg_attr(
             feature = "enable-serde",
@@ -1412,6 +1543,9 @@ pub(crate) fn generate(
     gen_instruction_data(&formats, &mut fmt);
     fmt.empty_line();
     gen_instruction_data_impl(&formats, &mut fmt);
+    gen_instruction_data_to_instruction_imms(&formats, &mut fmt);
+    gen_instruction_imms_impl(&formats, &mut fmt);
+    gen_instruction_imms_to_instruction_data(&formats, &mut fmt);
     fmt.empty_line();
     gen_opcodes(all_inst, &mut fmt);
     fmt.empty_line();

cranelift/codegen/meta/src/gen_inst.rs

@@ -31,6 +31,17 @@ pub(crate) fn define() -> SettingGroup {
         vec!["none", "speed", "speed_and_size"],
     );
 
+    settings.add_bool(
+        "use_egraphs",
+        "Enable egraph-based optimization.",
+        r#"
+            This enables an optimization phase that converts CLIF to an egraph (equivalence graph)
+            representation, performs various rewrites, and then converts it back. This can result in
+            better optimization, but at the cost of a longer compile time.
+        "#,
+        false,
+    );
+
     settings.add_bool(
         "enable_verifier",
         "Run the Cranelift IR verifier at strategic times during compilation.",