Tweak polling test

.github/workflows/build.yml

@@ -89,6 +89,13 @@ jobs:
             ocamlparam: '_,w=-46,regalloc=irc,regalloc-param=SPLIT_LIVE_RANGES:on,regalloc-param=IRC_SPILLING_HEURISTICS:flat-uses,regalloc-validate=1'
             check_arch: true
 
+          - name: irc_polling
+            config: --enable-middle-end=flambda2 --enable-poll-insertion
+            os: ubuntu-latest
+            build_ocamlparam: '_,w=-46,regalloc=irc,regalloc-param=SPLIT_LIVE_RANGES:on,regalloc-param=IRC_SPILLING_HEURISTICS:flat-uses,regalloc-validate=1'
+            ocamlparam: '_,w=-46,regalloc=irc,regalloc-param=SPLIT_LIVE_RANGES:on,regalloc-param=IRC_SPILLING_HEURISTICS:flat-uses,regalloc-validate=1'
+            check_arch: true
+
           - name: irc_frame_pointers
             config: --enable-middle-end=flambda2 --enable-frame-pointers
             os: ubuntu-latest

backend/cfg/cfg_polling.ml

@@ -7,94 +7,56 @@ module DLL = Flambda_backend_utils.Doubly_linked_list
 let function_is_assumed_to_never_poll func =
   Polling.function_is_assumed_to_never_poll func
 
+(* Compututation of the "safe" map, which is a map from labels to booleans where
+   `true` indicates the block contains a safe point such as a poll or an alloc
+   instruction. *)
+
+(* CR-soon xclerc for xclerc: given how we use the safe map below, it is not
+   clear taking into accounts terminator makes a difference; maybe matching over
+   the terminator to always return `false` would be better. *)
+
+let is_safe_basic : Cfg.basic Cfg.instruction -> bool =
+ fun instr ->
+  match[@ocaml.warning "-4"] instr.desc with
+  | Op (Poll | Alloc _) -> true
+  | Op _ | Reloadretaddr | Pushtrap _ | Poptrap | Prologue | Stack_check _ ->
+    false
+
+let is_safe_terminator : Cfg.terminator Cfg.instruction -> bool =
+ fun term ->
+  match term.desc with
+  | Never -> assert false
+  | Always _ | Parity_test _ | Truth_test _ | Float_test _ | Int_test _
+  | Switch _ ->
+    false
+  | Raise _ -> false
+  | Tailcall_self _ | Tailcall_func _ | Return -> true
+  | Call_no_return _ | Call _ | Prim _ | Specific_can_raise _ -> false
+
+let is_safe_block : Cfg.basic_block -> bool =
+ fun block ->
+  is_safe_terminator block.terminator || DLL.exists block.body ~f:is_safe_basic
+
+let safe_map_of_cfg : Cfg.t -> bool Label.Tbl.t =
+ fun cfg ->
+  Cfg.fold_blocks cfg
+    ~init:(Label.Tbl.create (Label.Tbl.length cfg.blocks))
+    ~f:(fun label block acc ->
+      Label.Tbl.replace acc label (is_safe_block block);
+      acc)
+
 (* These are used for the poll error annotation later on*)
 type polling_point = Polling.polling_point =
   | Alloc
   | Poll
   | Function_call
   | External_call
 
-type error = Polling.error = Poll_error of (polling_point * Debuginfo.t) list
+type error = Polling.error =
+  | Poll_error of Debuginfo.t * (polling_point * Debuginfo.t) list
 
 exception Error = Polling.Error
 
-(* Detection of recursive handlers that are not guaranteed to poll at every loop
-   iteration. *)
-
-(* We use a backwards dataflow analysis to compute a mapping from handlers H (=
-   loop heads) to either "safe" or "unsafe".
-
-   H is "safe" if every path starting from H goes through an Alloc, Poll,
-   Return, Tailcall_self or Tailcall_func instruction.
-
-   H is "unsafe", therefore, if starting from H we can loop infinitely without
-   crossing an Alloc or Poll instruction. *)
-
-module Unsafe_or_safe_domain = struct
-  type t = Polling.Unsafe_or_safe.t
-
-  let bot = Polling.Unsafe_or_safe.bot
-
-  let join = Polling.Unsafe_or_safe.join
-
-  let less_equal = Polling.Unsafe_or_safe.lessequal
-end
-
-module Unsafe_or_safe_transfer = struct
-  type domain = Unsafe_or_safe_domain.t
-
-  type context = unit
-
-  type error = |
-
-  let basic :
-      domain -> Cfg.basic Cfg.instruction -> context -> (domain, error) result =
-   fun dom instr () ->
-    match[@ocaml.warning "-4"] instr.desc with
-    | Op (Poll | Alloc _) -> Ok Safe
-    | Op _ | Reloadretaddr | Pushtrap _ | Poptrap | Prologue | Stack_check _ ->
-      Ok dom
-
-  let terminator :
-      domain ->
-      exn:domain ->
-      Cfg.terminator Cfg.instruction ->
-      context ->
-      (domain, error) result =
-   fun dom ~exn term () ->
-    match term.desc with
-    | Never -> assert false
-    | Always _ | Parity_test _ | Truth_test _ | Float_test _ | Int_test _
-    | Switch _ ->
-      Ok dom
-    | Raise _ -> Ok exn
-    | Tailcall_self _ | Tailcall_func _ | Return -> Ok Safe
-    | Call_no_return _ | Call _ | Prim _ | Specific_can_raise _ ->
-      if Cfg.can_raise_terminator term.desc
-      then Ok (Unsafe_or_safe_domain.join dom exn)
-      else Ok dom
-
-  let exception_ : domain -> context -> (domain, error) result =
-   fun dom () -> Ok dom
-end
-
-module PolledLoopsAnalysis =
-  Cfg_dataflow.Backward (Unsafe_or_safe_domain) (Unsafe_or_safe_transfer)
-
-let polled_loops_analysis :
-    Cfg_with_layout.t -> PolledLoopsAnalysis.domain Label.Tbl.t =
- fun cfg_with_layout ->
-  let init : Unsafe_or_safe_domain.t = Unsafe_or_safe_domain.bot in
-  let cfg = Cfg_with_layout.cfg cfg_with_layout in
-  match
-    PolledLoopsAnalysis.run ~init ~map:PolledLoopsAnalysis.Block ~exnescape:Safe
-      cfg ()
-  with
-  | Ok res -> res
-  | Aborted _ -> .
-  | Max_iterations_reached ->
-    Misc.fatal_error "Cfg_polling.polled_loops_analysis has been interrupted"
-
 (* Detection of functions that can loop via a tail-call without going through a
    poll point. *)
 
@@ -188,21 +150,54 @@ let potentially_recursive_tailcall :
     Misc.fatal_error
       "Cfg_polling.potentially_recursive_tailcall has been interrupted"
 
+(* Insertion of poll instruction onto back edges: a poll instruction is needed
+   if there is an unsafe path from the header of the loop to the back edge. An
+   unsafe path is simply a path such that all blocks are unsafe. *)
+
+let exists_unsafe_path :
+    Cfg.t -> safe_map:bool Label.Tbl.t -> from:Label.t -> to_:Label.t -> bool =
+ fun cfg ~safe_map ~from ~to_ ->
+  let exception Found in
+  try
+    let open_ = ref (Label.Set.singleton from) in
+    let closed = ref Label.Set.empty in
+    while not (Label.Set.is_empty !open_) do
+      let label = Label.Set.choose !open_ in
+      if Label.equal label to_ then raise Found;
+      open_ := Label.Set.remove label !open_;
+      closed := Label.Set.add label !closed;
+      match Label.Tbl.find_opt safe_map label with
+      | None ->
+        Misc.fatal_errorf
+          "Cfg_polling.exists_unsafe_path: missing safety information for \
+           block %d"
+          label
+      | Some true -> ()
+      | Some false ->
+        let block = Cfg.get_block_exn cfg label in
+        let successor_labels =
+          Cfg.successor_labels ~normal:true ~exn:true block
+        in
+        Label.Set.iter
+          (fun successor_label ->
+            match Label.Set.mem successor_label !closed with
+            | true -> ()
+            | false -> open_ := Label.Set.add successor_label !open_)
+          successor_labels
+    done;
+    false
+  with Found -> true
+
 let instr_cfg_with_layout :
     Cfg_with_layout.t ->
-    block_needs_poll:PolledLoopsAnalysis.domain Label.Tbl.t ->
+    safe_map:bool Label.Tbl.t ->
     back_edges:Cfg_loop_infos.EdgeSet.t ->
     bool =
- fun cfg_with_layout ~block_needs_poll ~back_edges ->
+ fun cfg_with_layout ~safe_map ~back_edges ->
   let cfg = Cfg_with_layout.cfg cfg_with_layout in
   Cfg_loop_infos.EdgeSet.fold
     (fun { Cfg_loop_infos.Edge.src; dst } added_poll ->
-      let needs_poll =
-        match Label.Tbl.find_opt block_needs_poll dst with
-        | None -> assert false
-        | Some Safe -> false
-        | Some Unsafe -> true
-      in
+      let needs_poll = exists_unsafe_path cfg ~safe_map ~from:dst ~to_:src in
       if needs_poll
       then (
         let after = Cfg.get_block_exn cfg src in
@@ -230,11 +225,14 @@ let instr_cfg_with_layout :
           let instrs = DLL.of_list [poll] in
           (* CR-soon xclerc: that kind of indicates the `insert_block` function
              should be moved outside of "regalloc/" *)
-          let (_ : Cfg.basic_block list) =
+          let inserted_blocks =
             Regalloc_utils.insert_block cfg_with_layout instrs ~after ~before
               ~next_instruction_id
           in
-          ());
+          (* All the inserted blocks are safe since they contain a poll
+             instruction *)
+          List.iter inserted_blocks ~f:(fun block ->
+              Label.Tbl.replace safe_map block.Cfg.start true));
         true)
       else added_poll)
     back_edges false
@@ -332,13 +330,13 @@ let instrument_fundecl :
   if is_disabled cfg.fun_name
   then cfg_with_layout
   else
-    let block_needs_poll = polled_loops_analysis cfg_with_layout in
+    let safe_map = safe_map_of_cfg cfg in
     (* CR-soon xclerc for xclerc: consider using `Cfg_with_infos` to cache the
        computations *)
     let doms = Cfg_dominators.build cfg in
     let back_edges = Cfg_loop_infos.compute_back_edges cfg doms in
     let added_poll =
-      instr_cfg_with_layout cfg_with_layout ~block_needs_poll ~back_edges
+      instr_cfg_with_layout cfg_with_layout ~safe_map ~back_edges
     in
     (match cfg.fun_poll with
     | Error_poll -> (
@@ -350,7 +348,7 @@ let instrument_fundecl :
             ~cmp:(fun left right -> Debuginfo.compare (snd left) (snd right))
             poll_error_instrs
         in
-        raise (Error (Poll_error poll_error_instrs)))
+        raise (Error (Poll_error (cfg.fun_dbg, poll_error_instrs))))
     | Default_poll -> ());
     let new_contains_calls =
       (* `added_poll` is used to avoid iterating over the CFG if we have added a

backend/polling.ml

@@ -29,7 +29,7 @@ let function_is_assumed_to_never_poll func =
 
 (* These are used for the poll error annotation later on*)
 type polling_point = Alloc | Poll | Function_call | External_call
-type error = Poll_error of (polling_point * Debuginfo.t) list
+type error = Poll_error of Debuginfo.t * (polling_point * Debuginfo.t) list
 
 exception Error of error
 
@@ -276,7 +276,7 @@ let instrument_fundecl ~future_funcnames:_ (f : Mach.fundecl) : Mach.fundecl =
     | Error_poll -> begin
         match find_poll_alloc_or_calls new_body with
         | [] -> ()
-        | poll_error_instrs -> raise (Error(Poll_error poll_error_instrs))
+        | poll_error_instrs -> raise (Error(Poll_error (f.fun_dbg, poll_error_instrs)))
       end
     | Default_poll -> () end;
     let new_contains_calls = f.fun_contains_calls || !contains_polls in
@@ -300,7 +300,7 @@ let instr_type p =
   | External_call -> "external call that allocates"
 
 let report_error ppf = function
-| Poll_error instrs ->
+| Poll_error (_fun_dbg, instrs) ->
   begin
     let num_inserted_polls =
       List.fold_left
@@ -311,27 +311,28 @@ let report_error ppf = function
       if num_user_polls = 0 then
         fprintf ppf "Function with poll-error attribute contains polling \
         points (inserted by the compiler)\n"
-      else begin
+      else
         fprintf ppf
         "Function with poll-error attribute contains polling points:\n";
-        List.iter (fun (p,dbg) ->
-          begin match p with
-          | Poll -> ()
-          | Alloc | Function_call | External_call ->
-            fprintf ppf "\t%s at " (instr_type p);
+      List.iter (fun (p,dbg) ->
+        begin match p with
+        | Poll
+        | Alloc | Function_call | External_call ->
+          fprintf ppf "\t%s" (instr_type p);
+          if not (Debuginfo.is_none dbg) then begin
+            fprintf ppf " at ";
             Location.print_loc ppf (Debuginfo.to_location dbg);
-            fprintf ppf "\n"
-          end
-        ) instrs;
-        if num_inserted_polls > 0 then
-          fprintf ppf "\t(plus compiler-inserted polling point(s) in prologue \
-          and/or loop back edges)\n"
-      end
+          end;
+          fprintf ppf "\n"
+        end
+      ) (List.sort (fun (_, left) (_, right) -> Debuginfo.compare left right) instrs)
   end
 
 let () =
   Location.register_error_of_exn
     (function
-      | Error err -> Some (Location.error_of_printer_file report_error err)
+      | Error (Poll_error (fun_dbg, _instrs) as err) ->
+        let loc = Debuginfo.to_location fun_dbg in
+        Some (Location.error_of_printer ~loc report_error err)
       | _ -> None
     )

backend/polling.mli

@@ -24,7 +24,7 @@ val is_disabled : string -> bool
 
 type polling_point = Alloc | Poll | Function_call | External_call
 
-type error = Poll_error of (polling_point * Debuginfo.t) list
+type error = Poll_error of Debuginfo.t * (polling_point * Debuginfo.t) list
 
 exception Error of error
 

flambda-backend/testsuite/tests/asmcomp/poll_attr_both.compilers.reference

@@ -1,6 +1,13 @@
-File "poll_attr_both.ml", line 1:
+File "poll_attr_both.ml", lines 16-22, characters 19-13:
+16 | ...................x =
+17 |   let y = Sys.opaque_identity(ref 42) in
+18 |     let x2 = v x in
+19 |       for c = 0 to x2 do
+20 |         ignore(Sys.opaque_identity(42))
+21 |       done;
+22 |       x2 + !y
 Error: Function with poll-error attribute contains polling points:
+	inserted poll
 	allocation at File "poll_attr_both.ml", line 17, characters 29-37
 	function call at File "poll_attr_both.ml", line 18, characters 13-16
-	(plus compiler-inserted polling point(s) in prologue and/or loop back edges)
 

flambda-backend/testsuite/tests/asmcomp/poll_attr_inserted.compilers.reference

@@ -1,3 +1,8 @@
-File "poll_attr_inserted.ml", line 1:
+File "poll_attr_inserted.ml", lines 16-19, characters 19-6:
+16 | ...................x =
+17 |   for c = 0 to 2 do
+18 |     ignore(Sys.opaque_identity(42))
+19 |   done
 Error: Function with poll-error attribute contains polling points (inserted by the compiler)
+	inserted poll
 

flambda-backend/testsuite/tests/asmcomp/poll_attr_prologue.compilers.reference

@@ -1,5 +1,9 @@
-File "poll_attr_prologue.ml", line 1:
+File "poll_attr_prologue.ml", lines 14-17, characters 36-38:
+14 | ....................................x l =
+15 |   match l with
+16 |   | [] -> 0
+17 |   | _ :: tl -> (d[@tailcall]) (x+1) tl
 Error: Function with poll-error attribute contains polling points:
+	inserted poll at File "poll_attr_prologue.ml", lines 14-17, characters 36-38
 	function call at File "poll_attr_prologue.ml", line 17, characters 15-38
-	(plus compiler-inserted polling point(s) in prologue and/or loop back edges)
 

flambda-backend/testsuite/tests/asmcomp/poll_attr_prologue.ml

@@ -12,6 +12,11 @@
 *)
 
 let[@poll error][@loop never] rec c x l =
+  match l with
+  | [] -> 0
+  | _ :: tl -> (d[@tailcall]) (x+1) tl
+
+and d x l =
   match l with
   | [] -> 0
   | _ :: tl -> (c[@tailcall]) (x+1) tl

flambda-backend/testsuite/tests/asmcomp/poll_attr_user.compilers.reference

@@ -1,4 +1,11 @@
-File "poll_attr_user.ml", line 1:
+File "poll_attr_user.ml", lines 16-22, characters 19-13:
+16 | ...................x =
+17 |   let y = Sys.opaque_identity(ref 42) in
+18 |     let x2 = v x in
+19 |       for c = 0 to x2 do
+20 |         ignore(Sys.opaque_identity(ref 42))
+21 |       done;
+22 |       x2 + !y
 Error: Function with poll-error attribute contains polling points:
 	allocation at File "poll_attr_user.ml", line 17, characters 29-37
 	function call at File "poll_attr_user.ml", line 18, characters 13-16