From ce333b30701af496ce5da2420fffc0c3293197b3 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Wed, 1 Mar 2023 17:23:09 -0800
Subject: [PATCH] Rematerialize constants in backend codegen.

---
 src/backend/localify.rs |  2 +-
 src/backend/mod.rs      | 20 ++++++++++++--------
 src/backend/treeify.rs  | 37 +++++++++++++++++++++++++++++++------
 3 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/src/backend/localify.rs b/src/backend/localify.rs
index 2770575..8b1008a 100644
--- a/src/backend/localify.rs
+++ b/src/backend/localify.rs
@@ -74,7 +74,7 @@ impl<'a, V: Visitor> BlockVisitor<'a, V> {
         self.visitor.pre_term();
 
         for &inst in self.body.blocks[block].insts.iter().rev() {
-            if self.trees.owner.contains_key(&inst) {
+            if self.trees.owner.contains_key(&inst) || self.trees.remat.contains(&inst) {
                 continue;
             }
             self.visitor.post_inst(inst);
diff --git a/src/backend/mod.rs b/src/backend/mod.rs
index 8a649a5..9b42cce 100644
--- a/src/backend/mod.rs
+++ b/src/backend/mod.rs
@@ -142,7 +142,7 @@ impl<'a> WasmFuncBackend<'a> {
                 for &inst in &self.body.blocks[*block].insts {
                     // If this value is "owned", do nothing: it will be lowered in
                     // the one place it's used.
-                    if self.trees.owner.contains_key(&inst) {
+                    if self.trees.owner.contains_key(&inst) || self.trees.remat.contains(&inst) {
                         continue;
                     }
                     if let &ValueDef::Operator(..) = &self.body.values[inst] {
@@ -180,12 +180,16 @@ impl<'a> WasmFuncBackend<'a> {
     fn lower_value(&self, value: Value, func: &mut wasm_encoder::Function) {
         log::trace!("lower_value: value {}", value);
         let value = self.body.resolve_alias(value);
-        let local = match &self.body.values[value] {
-            &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => self.locals.values[value][0],
-            &ValueDef::PickOutput(orig_value, idx, _) => self.locals.values[orig_value][idx],
-            _ => unreachable!(),
-        };
-        func.instruction(&wasm_encoder::Instruction::LocalGet(local.index() as u32));
+        if self.trees.remat.contains(&value) {
+            self.lower_inst(value, /* root = */ false, func);
+        } else {
+            let local = match &self.body.values[value] {
+                &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => self.locals.values[value][0],
+                &ValueDef::PickOutput(orig_value, idx, _) => self.locals.values[orig_value][idx],
+                _ => unreachable!(),
+            };
+            func.instruction(&wasm_encoder::Instruction::LocalGet(local.index() as u32));
+        }
     }
 
     fn lower_set_value(&self, value: Value, func: &mut wasm_encoder::Function) {
@@ -205,7 +209,7 @@ impl<'a> WasmFuncBackend<'a> {
             &ValueDef::Operator(ref op, ref args, ref tys) => {
                 for &arg in &args[..] {
                     let arg = self.body.resolve_alias(arg);
-                    if self.trees.owner.contains_key(&arg) {
+                    if self.trees.owner.contains_key(&arg) || self.trees.remat.contains(&arg) {
                         log::trace!(" -> arg {} is owned", arg);
                         self.lower_inst(arg, /* root = */ false, func);
                     } else {
diff --git a/src/backend/treeify.rs b/src/backend/treeify.rs
index 4693ce0..550bc8b 100644
--- a/src/backend/treeify.rs
+++ b/src/backend/treeify.rs
@@ -4,7 +4,7 @@
 use crate::entity::EntityRef;
 use crate::ir::{FunctionBody, Value, ValueDef};
 use crate::Operator;
-use std::collections::{HashMap, HashSet};
+use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
 use std::convert::TryFrom;
 
 /// One "argument slot" of an operator defining a value.
@@ -19,22 +19,43 @@ pub struct Trees {
     /// For a given value that is defined by an operator, which
     /// Values, if any, live at each slot?
     pub owned: HashMap<ValueArg, Value>,
+    /// Values that are regenerated every time they are used.
+    pub remat: HashSet<Value>,
+}
+
+fn is_remat(op: &Operator) -> bool {
+    // Only ops with no args can be always-rematerialized.
+    match op {
+        Operator::I32Const { .. }
+        | Operator::I64Const { .. }
+        | Operator::F32Const { .. }
+        | Operator::F64Const { .. } => true,
+        _ => false,
+    }
 }
 
 impl Trees {
     pub fn compute(body: &FunctionBody) -> Trees {
-        let mut owner = HashMap::new();
-        let mut owned = HashMap::new();
-        let mut multi_use = HashSet::new();
+        let mut owner = HashMap::default();
+        let mut owned = HashMap::default();
+        let mut remat = HashSet::default();
+        let mut multi_use = HashSet::default();
 
         for (value, def) in body.values.entries() {
             match def {
-                &ValueDef::Operator(_, ref args, _) => {
+                &ValueDef::Operator(op, ref args, _) => {
                     // Ignore operators with invalid args: these must
                     // always be unreachable.
                     if args.iter().any(|arg| arg.is_invalid()) {
                         continue;
                     }
+                    // If this is an always-rematerialized operator,
+                    // mark it as such and continue.
+                    if is_remat(&op) {
+                        remat.insert(value);
+                        continue;
+                    }
+
                     // For each of the args, if the value is produced
                     // by a single-output op and is movable, and is
                     // not already recorded in `multi_use`, place it
@@ -75,7 +96,11 @@ impl Trees {
             });
         }
 
-        Trees { owner, owned }
+        Trees {
+            owner,
+            owned,
+            remat,
+        }
     }
 
     fn is_single_output_op(body: &FunctionBody, value: Value) -> Option<Operator> {