Update downschedule_compensation_code function

Transitively consider the additional dependent instructions. This is
necessary since we delete the originals of duplicated-for-code-motion
instruction since this is the only way it will [currently] work for
memory stores. Before only instructions that directly dependent on a
monved/copied for code motion instruction were considered.

1 files changed, 33 insertions, 38 deletions

diff --git a/scheduling/MyRTLpathScheduleraux.ml b/scheduling/MyRTLpathScheduleraux.ml
index 79a6ae64..ef8d5eb5 100644
--- a/scheduling/MyRTLpathScheduleraux.ml
+++ b/scheduling/MyRTLpathScheduleraux.ml
@@ -778,22 +778,18 @@ let downschedule_compensation_code sb code pm live_renames ~next_free_pc ~next_f
    let liveins' = update_liveins sb.liveins live_renames in
    (* Use the new names to calculate proper dependencies *)
    let path_deps = intra_path_dependencies {sb with liveins = liveins'} code in
-   (* Maps from a pc, to all instructions (pcs) that directly depend on it *)
-   (* TODO: extract into its own function *)
-   let uses = PTree.fold
-    (fun uses pc deps ->
-      let uses = HashedSet.PSet.fold
-        (fun uses dep ->
-          let old = ptree_get_or_default uses dep HashedSet.PSet.empty in
-          let upd = HashedSet.PSet.add pc old in
-          PTree.set dep upd uses )
-          deps
-          uses
-      in
-      uses )
+   let path_deps_without_iconds = PTree.map
+    (fun _pc deps ->
+      HashedSet.PSet.filter
+        (fun dep_pc -> not @@ is_icond @@ get_some @@ PTree.get dep_pc code)
+        deps)
     path_deps
-    PTree.empty
+  in
+   let transitive_path_deps_without_iconds =
+    PTree.map (fun pc _deps -> transitive_dependencies path_deps_without_iconds pc)
+      path_deps
    in
+   let transitive_uses_without_iconds = uses_of_deps transitive_path_deps_without_iconds in
 
   (* For each side-exit, check if all the dependencies are still above it
    * if not, remember the pc and transitively consider its dependencies until
@@ -810,46 +806,45 @@ let downschedule_compensation_code sb code pm live_renames ~next_free_pc ~next_f
         (side_exit_pc, moved_deps_sorted) )
   in
   let (side_exit_pcs, insts_pcs) = List.split side_exit_and_compensation in
-  assert (List.for_all (fun pc -> Array.mem pc orig_sb.instructions) side_exit_pcs);
 
-  let pc_to_orig_idx =
-    Duplicateaux.generate_fwmap
-      (Array.to_list orig_sb.instructions)
-      (List.init (Array.length orig_sb.instructions) (fun i -> i))
-      PTree.empty
-  in
   let to_insert_pcs =
     ListLabels.fold_left2 side_exit_pcs insts_pcs
     ~init:InsertPositionMap.empty
     ~f:(fun acc side_exit_pc pcs ->
       InsertPositionMap.add (InsertPosition.Above side_exit_pc) pcs acc) in
-  let (to_insert_below_side_exit : Camlcoq.P.t list InsertPositionMap.t)=
-    (* XXX: This code is not correct right now, we need to transitively find dependent
-     * instructions (cf. moved_dependencies). *)
+  let (to_insert_as_well : Camlcoq.P.t list InsertPositionMap.t)=
      InsertPositionMap.fold
       (fun side_exit_pc pcs_to_insert acc ->
         let side_exit_pc = InsertPosition.anchor side_exit_pc in
-        let uses = List.map (fun pc -> ptree_get_or_default uses pc HashedSet.PSet.empty) pcs_to_insert in
-        let uses = List.fold_left HashedSet.PSet.union HashedSet.PSet.empty uses in
-        let uses = HashedSet.PSet.filter
-          (fun pc ->
-            not @@ List.mem pc pcs_to_insert
-            && apply_map' pc_to_idx pc < apply_map' pc_to_idx side_exit_pc
-            (* Cannot move Iconds *)
-            && (not @@ is_icond @@ get_some @@ PTree.get pc code))
-          uses
+        let ( let* ) = Option.bind in
+        let collateral_moves =
+          pcs_to_insert
+          |> List.filter_map (fun pc ->
+            let* uses = PTree.get pc transitive_uses_without_iconds in
+            HashedSet.PSet.filter
+              (fun pc ->
+                  not @@ List.mem pc pcs_to_insert
+                && apply_map' pc_to_idx pc < apply_map' pc_to_idx side_exit_pc
+                && (not @@ is_icond @@ get_some @@ PTree.get pc code))
+              uses
+            |> Option.some)
+          |> ListLabels.fold_left ~f:HashedSet.PSet.union ~init:HashedSet.PSet.empty
         in
-        let uses = ListLabels.sort (HashedSet.PSet.elements uses)
+        let collateral_moves = ListLabels.sort (HashedSet.PSet.elements collateral_moves)
           ~cmp:(fun pc1 pc2 -> Int.compare (apply_map' pc_to_idx pc1) (apply_map' pc_to_idx pc2))
         in
-        let side_exit_idx = apply_map' pc_to_orig_idx side_exit_pc in
-        let acc = InsertPositionMap.add (InsertPosition.Below orig_sb.instructions.(side_exit_idx)) uses acc in
-        (* let acc = InsertPositionMap.add (InsertPosition.Below side_exit_pc) uses acc in *)
+        (* TODO?
+         * In principle we should be able to move the instructions Below, unless they are
+         * live (liveins') at the side exit.
+         * But adding them above is simpler and unnecessary instructions should be removed
+         * by DCE. *)
+        let acc = InsertPositionMap.add (InsertPosition.Above side_exit_pc) collateral_moves acc in
         acc)
       to_insert_pcs
       InsertPositionMap.empty
   in
-  let to_insert_pcs = InsertPositionMap.merge merge_append to_insert_pcs to_insert_below_side_exit in
+  let to_insert_pcs = InsertPositionMap.merge merge_append to_insert_pcs to_insert_as_well in
+  (* TODO: Move out of this function since it can be inferred from to_insert *)
   let to_rename =
     InsertPositionMap.to_seq to_insert_pcs
     |> List.of_seq