From 5c95109ad16074d43577c61bc49e206fdae0b2b7 Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Mon, 21 Aug 2023 16:25:07 +0100
Subject: [PATCH 1/9] Refactor Lambda.check_attribute: separate constructor for
 Assume

---
 backend/checkmach.ml                          |  8 +++--
 backend/cmm.ml                                |  4 +--
 backend/cmm.mli                               |  3 +-
 backend/cmm_helpers.ml                        |  6 ++--
 backend/printcmm.ml                           |  9 +++--
 .../simplify/simplify_set_of_closures.ml      |  1 +
 middle_end/flambda2/terms/check_attribute.ml  | 34 +++++++++++++------
 middle_end/flambda2/terms/check_attribute.mli |  6 +++-
 .../flambda2/to_cmm/to_cmm_set_of_closures.ml |  6 ++--
 ocaml/lambda/lambda.ml                        |  5 ++-
 ocaml/lambda/lambda.mli                       |  7 ++--
 ocaml/lambda/printlambda.ml                   |  9 +++--
 ocaml/lambda/translattribute.ml               | 34 +++++++++++--------
 13 files changed, 86 insertions(+), 46 deletions(-)

diff --git a/backend/checkmach.ml b/backend/checkmach.ml
index 3336c531262..f7b3e16252e 100644
--- a/backend/checkmach.ml
+++ b/backend/checkmach.ml
@@ -419,12 +419,14 @@ end = struct
       List.filter_map
         (fun (c : Cmm.codegen_option) ->
           match c with
-          | Check { property; strict; assume; loc } when property = spec ->
-            Some { strict; assume; loc }
+          | Check { property; strict; loc } when property = spec ->
+            Some { strict; assume = false; loc }
+          | Assume { property; strict; loc; } when property = spec ->
+            Some { strict; assume = true; loc }
           | Ignore_assert_all property when property = spec ->
             ignore_assert_all := true;
             None
-          | Ignore_assert_all _ | Check _ | Reduce_code_size | No_CSE
+          | Ignore_assert_all _ | Check _ | Assume _ | Reduce_code_size | No_CSE
           | Use_linscan_regalloc ->
             None)
         codegen_options
diff --git a/backend/cmm.ml b/backend/cmm.ml
index eb7b62fa3e5..a11f259671b 100644
--- a/backend/cmm.ml
+++ b/backend/cmm.ml
@@ -302,8 +302,8 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Ignore_assert_all of property
-  | Check of { property: property; strict: bool; assume: bool;
-               loc : Location.t; }
+  | Assume of { property: property; strict: bool; loc: Location.t }
+  | Check of { property: property; strict: bool; loc : Location.t; }
 
 type fundecl =
   { fun_name: symbol;
diff --git a/backend/cmm.mli b/backend/cmm.mli
index ada1d10ae80..f23516dfbd1 100644
--- a/backend/cmm.mli
+++ b/backend/cmm.mli
@@ -317,7 +317,8 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Ignore_assert_all of property
-  | Check of { property: property; strict: bool; assume: bool; loc: Location.t }
+  | Assume of { property: property; strict: bool; loc: Location.t }
+  | Check of { property: property; strict: bool; loc: Location.t }
 
 type fundecl =
   { fun_name: symbol;
diff --git a/backend/cmm_helpers.ml b/backend/cmm_helpers.ml
index 13e4760eb6d..28bae4ee40c 100644
--- a/backend/cmm_helpers.ml
+++ b/backend/cmm_helpers.ml
@@ -4512,8 +4512,10 @@ let transl_property : Lambda.property -> Cmm.property = function
 let transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list = function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
-  | Check { property; strict; assume; loc } ->
-    [Check { property = transl_property property; strict; assume; loc }]
+  | Assume { property; strict; loc } ->
+    [Assume { property = transl_property property; strict; loc }]
+  | Check { property; strict; loc } ->
+    [Check { property = transl_property property; strict; loc }]
 
 let kind_of_layout (layout : Lambda.layout) =
   match layout with
diff --git a/backend/printcmm.ml b/backend/printcmm.ml
index 8b8fff0af37..0add192cef7 100644
--- a/backend/printcmm.ml
+++ b/backend/printcmm.ml
@@ -390,9 +390,12 @@ let codegen_option = function
   | Use_linscan_regalloc -> "linscan"
   | Ignore_assert_all property ->
     Printf.sprintf "ignore %s" (property_to_string property)
-  | Check { property; strict; assume; loc = _ } ->
-    Printf.sprintf "%s %s%s"
-      (if assume then "assume" else "assert")
+  | Assume { property; strict; loc = _ } ->
+    Printf.sprintf "assume %s%s"
+      (property_to_string property)
+      (if strict then " strict" else "")
+  | Check { property; strict; loc = _ } ->
+    Printf.sprintf "assert %s%s"
       (property_to_string property)
       (if strict then " strict" else "")
 
diff --git a/middle_end/flambda2/simplify/simplify_set_of_closures.ml b/middle_end/flambda2/simplify/simplify_set_of_closures.ml
index c6ac7431834..240366aa8f2 100644
--- a/middle_end/flambda2/simplify/simplify_set_of_closures.ml
+++ b/middle_end/flambda2/simplify/simplify_set_of_closures.ml
@@ -493,6 +493,7 @@ let simplify_function context ~outer_dacc function_slot code_id
       match Code_metadata.check code_metadata with
       | Default_check -> !Clflags.zero_alloc_check_assert_all
       | Ignore_assert_all Zero_alloc -> false
+      | Assume { property = Zero_alloc; _ } -> false
       | Check { property = Zero_alloc; _ } -> true
     in
     if never_delete then Code_id.Set.singleton code_id else Code_id.Set.empty
diff --git a/middle_end/flambda2/terms/check_attribute.ml b/middle_end/flambda2/terms/check_attribute.ml
index 24021e7e5f1..ac9840dd5c1 100644
--- a/middle_end/flambda2/terms/check_attribute.ml
+++ b/middle_end/flambda2/terms/check_attribute.ml
@@ -23,10 +23,14 @@ end
 type t =
   | Default_check
   | Ignore_assert_all of Property.t
+  | Assume of
+      { property : Property.t;
+        strict : bool;
+        loc : Location.t
+      }
   | Check of
       { property : Property.t;
         strict : bool;
-        assume : bool;
         loc : Location.t
       }
 
@@ -35,27 +39,35 @@ let print ppf t =
   | Default_check -> ()
   | Ignore_assert_all property ->
     Format.fprintf ppf "@[ignore %a@]" Property.print property
-  | Check { property; strict; assume; loc = _ } ->
-    Format.fprintf ppf "@[%s%s %a@]"
-      (if assume then "assume" else "assert")
+  | Assume { property; strict; loc = _ } ->
+    Format.fprintf ppf "@[assume%s %a@]"
+      (if strict then " strict" else "")
+      Property.print property
+  | Check { property; strict; loc = _ } ->
+    Format.fprintf ppf "@[assert%s %a@]"
       (if strict then " strict" else "")
       Property.print property
 
 let from_lambda : Lambda.check_attribute -> t = function
   | Default_check -> Default_check
   | Ignore_assert_all p -> Ignore_assert_all (Property.from_lambda p)
-  | Check { property; strict; assume; loc } ->
-    Check { property = Property.from_lambda property; strict; assume; loc }
+  | Assume { property; strict; loc } ->
+    Assume { property = Property.from_lambda property; strict; loc }
+  | Check { property; strict; loc } ->
+    Check { property = Property.from_lambda property; strict; loc }
 
 let equal x y =
   match x, y with
   | Default_check, Default_check -> true
   | Ignore_assert_all p1, Ignore_assert_all p2 -> Property.equal p1 p2
-  | ( Check { property = p1; strict = s1; assume = a1; loc = loc1 },
-      Check { property = p2; strict = s2; assume = a2; loc = loc2 } ) ->
-    Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal a1 a2 && loc1 = loc2
-  | (Default_check | Ignore_assert_all _ | Check _), _ -> false
+  | ( Check { property = p1; strict = s1; loc = loc1 },
+      Check { property = p2; strict = s2; loc = loc2 } ) ->
+    Property.equal p1 p2 && Bool.equal s1 s2 && loc1 = loc2
+  | ( Assume { property = p1; strict = s1; loc = loc1 },
+      Assume { property = p2; strict = s2; loc = loc2 } ) ->
+    Property.equal p1 p2 && Bool.equal s1 s2 && loc1 = loc2
+  | (Default_check | Ignore_assert_all _ | Check _ | Assume _), _ -> false
 
 let is_default : t -> bool = function
   | Default_check -> true
-  | Ignore_assert_all _ | Check _ -> false
+  | Ignore_assert_all _ | Check _ | Assume _  -> false
diff --git a/middle_end/flambda2/terms/check_attribute.mli b/middle_end/flambda2/terms/check_attribute.mli
index 71d08bc6449..bbd3a805ae4 100644
--- a/middle_end/flambda2/terms/check_attribute.mli
+++ b/middle_end/flambda2/terms/check_attribute.mli
@@ -17,10 +17,14 @@ end
 type t =
   | Default_check
   | Ignore_assert_all of Property.t
+  | Assume of
+      { property : Property.t;
+        strict : bool;
+        loc : Location.t
+      }
   | Check of
       { property : Property.t;
         strict : bool;
-        assume : bool;
         loc : Location.t
       }
 
diff --git a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
index caebc11cbb1..b7a344e1c83 100644
--- a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
+++ b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
@@ -324,8 +324,10 @@ let transl_check_attrib : Check_attribute.t -> Cmm.codegen_option list =
   function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
-  | Check { property; strict; assume; loc } ->
-    [Check { property = transl_property property; strict; assume; loc }]
+  | Assume { property; strict; loc } ->
+    [Assume { property = transl_property property; strict; loc }]
+  | Check { property; strict; loc } ->
+    [Check { property = transl_property property; strict; loc }]
 
 (* Translation of the bodies of functions. *)
 
diff --git a/ocaml/lambda/lambda.ml b/ocaml/lambda/lambda.ml
index fe11f14e9f6..6330767ce91 100644
--- a/ocaml/lambda/lambda.ml
+++ b/ocaml/lambda/lambda.ml
@@ -496,9 +496,12 @@ type check_attribute =
   | Ignore_assert_all of property
   | Check of { property: property;
                strict: bool;
-               assume: bool;
                loc: Location.t;
              }
+  | Assume of { property: property;
+                strict: bool;
+                loc: Location.t;
+              }
 
 type loop_attribute =
   | Always_loop (* [@loop] or [@loop always] *)
diff --git a/ocaml/lambda/lambda.mli b/ocaml/lambda/lambda.mli
index 24d29bdda42..05f99fb0f48 100644
--- a/ocaml/lambda/lambda.mli
+++ b/ocaml/lambda/lambda.mli
@@ -391,11 +391,12 @@ type check_attribute =
                   then the property holds (but property violations on
                   exceptional returns or divering loops are ignored).
                   This definition may not be applicable to new properties. *)
-               assume: bool;
-               (* [assume=true] assume without checking that the
-                  property holds *)
                loc: Location.t;
              }
+  | Assume of { property: property;
+                strict: bool;
+                loc: Location.t;
+              }
 
 type loop_attribute =
   | Always_loop (* [@loop] or [@loop always] *)
diff --git a/ocaml/lambda/printlambda.ml b/ocaml/lambda/printlambda.ml
index 32b7b988c23..75ac6ef61c9 100644
--- a/ocaml/lambda/printlambda.ml
+++ b/ocaml/lambda/printlambda.ml
@@ -634,9 +634,12 @@ let check_attribute ppf check =
   | Default_check -> ()
   | Ignore_assert_all p ->
     fprintf ppf "ignore assert all %s@ " (check_property p)
-  | Check {property=p; assume; strict; loc = _} ->
-    fprintf ppf "%s %s%s@ "
-      (if assume then "assume" else "assert")
+  | Assume {property=p; strict; loc = _} ->
+    fprintf ppf "assume %s%s@ "
+      (check_property p)
+      (if strict then " strict" else "")
+  | Check {property=p; strict; loc = _} ->
+    fprintf ppf "assert %s%s@ "
       (check_property p)
       (if strict then " strict" else "")
 
diff --git a/ocaml/lambda/translattribute.ml b/ocaml/lambda/translattribute.ml
index d0ab18182de..61ee8eb4ed0 100644
--- a/ocaml/lambda/translattribute.ml
+++ b/ocaml/lambda/translattribute.ml
@@ -251,11 +251,11 @@ let parse_property_attribute attr property =
   | Some {Parsetree.attr_name = {txt; loc}; attr_payload = payload}->
       parse_ids_payload txt loc
         ~default:Default_check
-        ~empty:(Check { property; strict = false; assume = false; loc; } )
+        ~empty:(Check { property; strict = false; loc; } )
         [
-          ["assume"], Check { property; strict = false; assume = true; loc; };
-          ["strict"], Check { property; strict = true; assume = false; loc; };
-          ["assume"; "strict"], Check { property; strict = true; assume = true; loc; };
+          ["assume"], Assume { property; strict = false; loc; };
+          ["strict"], Check { property; strict = true; loc; };
+          ["assume"; "strict"], Assume { property; strict = true; loc; };
           ["ignore"], Ignore_assert_all property
         ]
         payload
@@ -303,15 +303,15 @@ let get_property_attribute l p =
   (match attr, res with
    | None, Default_check -> ()
    | _, Default_check -> ()
-   | None, (Check _ | Ignore_assert_all _ ) -> assert false
+   | None, (Check _ | Assume _ | Ignore_assert_all _) -> assert false
    | Some _, Ignore_assert_all _ -> ()
-   | Some attr, Check { assume; _ } ->
+   | Some _, Assume _ -> ()
+   | Some attr, Check _ ->
      if !Clflags.zero_alloc_check && !Clflags.native_code then
        (* The warning for unchecked functions will not trigger if the check is requested
           through the [@@@zero_alloc all] top-level annotation rather than through the
           function annotation [@zero_alloc]. *)
-       if not assume then
-         Builtin_attributes.register_property attr.attr_name);
+       Builtin_attributes.register_property attr.attr_name);
    res
 
 let get_poll_attribute l =
@@ -414,7 +414,8 @@ let assume_zero_alloc attributes =
   match parse_property_attribute attr p with
   | Default_check -> false
   | Ignore_assert_all _ -> false
-  | Check { property = Zero_alloc; assume } -> assume
+  | Assume { property = Zero_alloc; _ } -> true
+  | Check { property = Zero_alloc; _ } -> false
 
 let assume_zero_alloc attributes =
   (* This function is used for "look-ahead" to find attributes
@@ -428,9 +429,12 @@ let add_check_attribute expr loc attributes =
     | Zero_alloc -> "zero_alloc"
   in
   let to_string = function
-    | Check { property; strict; assume; loc = _} ->
-      Printf.sprintf "%s %s%s"
-        (if assume then "assume" else "assert")
+    | Check { property; strict; loc = _} ->
+      Printf.sprintf "assert %s%s"
+        (to_string property)
+        (if strict then " strict" else "")
+    | Assume { property; strict; loc = _} ->
+      Printf.sprintf "assume %s%s"
         (to_string property)
         (if strict then " strict" else "")
     | Ignore_assert_all property ->
@@ -441,11 +445,13 @@ let add_check_attribute expr loc attributes =
   | Lfunction({ attr = { stub = false } as attr; } as funct) ->
     begin match get_property_attribute attributes Zero_alloc with
     | Default_check -> expr
-    | (Ignore_assert_all p | Check { property = p; _ }) as check ->
+    | (Ignore_assert_all p | Check { property = p; _ } | Assume { property = p; _ })
+      as check ->
       begin match attr.check with
       | Default_check -> ()
       | Ignore_assert_all p'
-      | Check { property = p'; strict = _; assume = _; loc = _; } ->
+      | Assume { property = p'; strict = _; loc = _; }
+      | Check { property = p'; strict = _; loc = _; } ->
         if p = p' then
           Location.prerr_warning loc
             (Warnings.Duplicated_attribute (to_string check));

From bcd7bfd392b04ff1ad7f51065b16c720f2cbe8f0 Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Mon, 21 Aug 2023 16:30:19 +0100
Subject: [PATCH 2/9] Slightyly improve debug printing of zero_alloc
 annotations

---
 backend/printcmm.ml                          | 8 ++++----
 middle_end/flambda2/terms/check_attribute.ml | 8 ++++----
 middle_end/printclambda.ml                   | 2 +-
 ocaml/lambda/printlambda.ml                  | 8 ++++----
 4 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/backend/printcmm.ml b/backend/printcmm.ml
index 0add192cef7..514b24318a9 100644
--- a/backend/printcmm.ml
+++ b/backend/printcmm.ml
@@ -391,13 +391,13 @@ let codegen_option = function
   | Ignore_assert_all property ->
     Printf.sprintf "ignore %s" (property_to_string property)
   | Assume { property; strict; loc = _ } ->
-    Printf.sprintf "assume %s%s"
+    Printf.sprintf "assume_%s%s"
       (property_to_string property)
-      (if strict then " strict" else "")
+      (if strict then "_strict" else "")
   | Check { property; strict; loc = _ } ->
-    Printf.sprintf "assert %s%s"
+    Printf.sprintf "assert_%s%s"
       (property_to_string property)
-      (if strict then " strict" else "")
+      (if strict then "_strict" else "")
 
 let print_codegen_options ppf l =
   List.iter (fun c -> fprintf ppf " %s" (codegen_option c)) l
diff --git a/middle_end/flambda2/terms/check_attribute.ml b/middle_end/flambda2/terms/check_attribute.ml
index ac9840dd5c1..5d7841b83ac 100644
--- a/middle_end/flambda2/terms/check_attribute.ml
+++ b/middle_end/flambda2/terms/check_attribute.ml
@@ -40,13 +40,13 @@ let print ppf t =
   | Ignore_assert_all property ->
     Format.fprintf ppf "@[ignore %a@]" Property.print property
   | Assume { property; strict; loc = _ } ->
-    Format.fprintf ppf "@[assume%s %a@]"
-      (if strict then " strict" else "")
+    Format.fprintf ppf "@[assume_%a%s@]"
       Property.print property
+      (if strict then "_strict" else "")
   | Check { property; strict; loc = _ } ->
-    Format.fprintf ppf "@[assert%s %a@]"
-      (if strict then " strict" else "")
+    Format.fprintf ppf "@[assert_%a%s@]"
       Property.print property
+      (if strict then "_strict" else "")
 
 let from_lambda : Lambda.check_attribute -> t = function
   | Default_check -> Default_check
diff --git a/middle_end/printclambda.ml b/middle_end/printclambda.ml
index 4791123896e..a4ce14c0c58 100644
--- a/middle_end/printclambda.ml
+++ b/middle_end/printclambda.ml
@@ -105,7 +105,7 @@ and one_fun ppf f =
     in
     iter f.params f.arity.params_layout
   in
-  fprintf ppf "(fun@ %s%s%a@ %d@ @[<2>%t@]@ @[<2>%a@])"
+  fprintf ppf "(fun@ %s%s@ %a@ %d@ @[<2>%t@]@ @[<2>%a@])"
     f.label (layout f.arity.return_layout) Printlambda.check_attribute f.check
     (List.length f.arity.params_layout) idents lam f.body
 
diff --git a/ocaml/lambda/printlambda.ml b/ocaml/lambda/printlambda.ml
index 75ac6ef61c9..8da6d532483 100644
--- a/ocaml/lambda/printlambda.ml
+++ b/ocaml/lambda/printlambda.ml
@@ -635,13 +635,13 @@ let check_attribute ppf check =
   | Ignore_assert_all p ->
     fprintf ppf "ignore assert all %s@ " (check_property p)
   | Assume {property=p; strict; loc = _} ->
-    fprintf ppf "assume %s%s@ "
+    fprintf ppf "assume_%s%s@ "
       (check_property p)
-      (if strict then " strict" else "")
+      (if strict then "_strict" else "")
   | Check {property=p; strict; loc = _} ->
-    fprintf ppf "assert %s%s@ "
+    fprintf ppf "assert_%s%s@ "
       (check_property p)
-      (if strict then " strict" else "")
+      (if strict then "_strict" else "")
 
 let function_attribute ppf t =
   if t.is_a_functor then

From 0b7c848c8655bd17106ee4d8052836a11860e7be Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Mon, 21 Aug 2023 18:20:34 +0100
Subject: [PATCH 3/9] Add field [never_returns_normally] to
 [Lambda.Check_attribute.Assume]

---
 backend/checkmach.ml                            | 17 ++++++++++++-----
 backend/cmm.ml                                  |  3 ++-
 backend/cmm.mli                                 |  3 ++-
 backend/cmm_helpers.ml                          |  4 ++--
 backend/printcmm.ml                             |  5 +++--
 middle_end/flambda2/terms/check_attribute.ml    | 16 +++++++++-------
 middle_end/flambda2/terms/check_attribute.mli   |  1 +
 .../flambda2/to_cmm/to_cmm_set_of_closures.ml   |  4 ++--
 ocaml/lambda/lambda.ml                          |  1 +
 ocaml/lambda/lambda.mli                         |  1 +
 ocaml/lambda/printlambda.ml                     |  5 +++--
 ocaml/lambda/translattribute.ml                 | 10 ++++++++--
 12 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/backend/checkmach.ml b/backend/checkmach.ml
index f7b3e16252e..05547f548c5 100644
--- a/backend/checkmach.ml
+++ b/backend/checkmach.ml
@@ -403,13 +403,19 @@ end = struct
   type t =
     { strict : bool;  (** strict or relaxed? *)
       assume : bool;
+      never_returns_normally : bool;
       loc : Location.t
           (** Source location of the annotation, used for error messages. *)
     }
 
   let get_loc t = t.loc
 
-  let expected_value t = if t.strict then Value.safe else Value.relaxed
+  let expected_value t =
+    let res = if t.strict then Value.safe else Value.relaxed in
+    if t.never_returns_normally then
+      { res with nor = V.Bot }
+    else
+      res
 
   let is_assume t = t.assume
 
@@ -420,9 +426,9 @@ end = struct
         (fun (c : Cmm.codegen_option) ->
           match c with
           | Check { property; strict; loc } when property = spec ->
-            Some { strict; assume = false; loc }
-          | Assume { property; strict; loc; } when property = spec ->
-            Some { strict; assume = true; loc }
+            Some { strict; assume = false; never_returns_normally = false; loc }
+          | Assume { property; strict; never_returns_normally; loc; } when property = spec ->
+            Some { strict; assume = true; never_returns_normally; loc }
           | Ignore_assert_all property when property = spec ->
             ignore_assert_all := true;
             None
@@ -435,7 +441,8 @@ end = struct
     | [] ->
       if !Clflags.zero_alloc_check_assert_all && not !ignore_assert_all
       then
-        Some { strict = false; assume = false; loc = Debuginfo.to_location dbg }
+        Some { strict = false; assume = false; never_returns_normally = false;
+               loc = Debuginfo.to_location dbg }
       else None
     | [p] -> Some p
     | _ :: _ ->
diff --git a/backend/cmm.ml b/backend/cmm.ml
index a11f259671b..84a86f2682d 100644
--- a/backend/cmm.ml
+++ b/backend/cmm.ml
@@ -302,7 +302,8 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Ignore_assert_all of property
-  | Assume of { property: property; strict: bool; loc: Location.t }
+  | Assume of { property: property; strict: bool; never_returns_normally: bool;
+                loc: Location.t }
   | Check of { property: property; strict: bool; loc : Location.t; }
 
 type fundecl =
diff --git a/backend/cmm.mli b/backend/cmm.mli
index f23516dfbd1..15a2cd17a2b 100644
--- a/backend/cmm.mli
+++ b/backend/cmm.mli
@@ -317,7 +317,8 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Ignore_assert_all of property
-  | Assume of { property: property; strict: bool; loc: Location.t }
+  | Assume of { property: property; strict: bool; never_returns_normally: bool;
+                loc: Location.t }
   | Check of { property: property; strict: bool; loc: Location.t }
 
 type fundecl =
diff --git a/backend/cmm_helpers.ml b/backend/cmm_helpers.ml
index 28bae4ee40c..ec9e391695e 100644
--- a/backend/cmm_helpers.ml
+++ b/backend/cmm_helpers.ml
@@ -4512,8 +4512,8 @@ let transl_property : Lambda.property -> Cmm.property = function
 let transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list = function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
-  | Assume { property; strict; loc } ->
-    [Assume { property = transl_property property; strict; loc }]
+  | Assume { property; strict; never_returns_normally; loc } ->
+    [Assume { property = transl_property property; strict; never_returns_normally; loc }]
   | Check { property; strict; loc } ->
     [Check { property = transl_property property; strict; loc }]
 
diff --git a/backend/printcmm.ml b/backend/printcmm.ml
index 514b24318a9..031691af092 100644
--- a/backend/printcmm.ml
+++ b/backend/printcmm.ml
@@ -390,10 +390,11 @@ let codegen_option = function
   | Use_linscan_regalloc -> "linscan"
   | Ignore_assert_all property ->
     Printf.sprintf "ignore %s" (property_to_string property)
-  | Assume { property; strict; loc = _ } ->
-    Printf.sprintf "assume_%s%s"
+  | Assume { property; strict; never_returns_normally; loc = _ } ->
+    Printf.sprintf "assume_%s%s%s"
       (property_to_string property)
       (if strict then "_strict" else "")
+      (if strict then "_never_returns_normally" else "")
   | Check { property; strict; loc = _ } ->
     Printf.sprintf "assert_%s%s"
       (property_to_string property)
diff --git a/middle_end/flambda2/terms/check_attribute.ml b/middle_end/flambda2/terms/check_attribute.ml
index 5d7841b83ac..7b413d4c223 100644
--- a/middle_end/flambda2/terms/check_attribute.ml
+++ b/middle_end/flambda2/terms/check_attribute.ml
@@ -26,6 +26,7 @@ type t =
   | Assume of
       { property : Property.t;
         strict : bool;
+        never_returns_normally : bool;
         loc : Location.t
       }
   | Check of
@@ -39,10 +40,11 @@ let print ppf t =
   | Default_check -> ()
   | Ignore_assert_all property ->
     Format.fprintf ppf "@[ignore %a@]" Property.print property
-  | Assume { property; strict; loc = _ } ->
-    Format.fprintf ppf "@[assume_%a%s@]"
+  | Assume { property; strict; never_returns_normally; loc = _ } ->
+    Format.fprintf ppf "@[assume_%a%s%s@]"
       Property.print property
       (if strict then "_strict" else "")
+      (if never_returns_normally then "_never_returns_normally" else "")
   | Check { property; strict; loc = _ } ->
     Format.fprintf ppf "@[assert_%a%s@]"
       Property.print property
@@ -51,8 +53,8 @@ let print ppf t =
 let from_lambda : Lambda.check_attribute -> t = function
   | Default_check -> Default_check
   | Ignore_assert_all p -> Ignore_assert_all (Property.from_lambda p)
-  | Assume { property; strict; loc } ->
-    Assume { property = Property.from_lambda property; strict; loc }
+  | Assume { property; strict; never_returns_normally; loc } ->
+    Assume { property = Property.from_lambda property; strict; never_returns_normally; loc }
   | Check { property; strict; loc } ->
     Check { property = Property.from_lambda property; strict; loc }
 
@@ -63,9 +65,9 @@ let equal x y =
   | ( Check { property = p1; strict = s1; loc = loc1 },
       Check { property = p2; strict = s2; loc = loc2 } ) ->
     Property.equal p1 p2 && Bool.equal s1 s2 && loc1 = loc2
-  | ( Assume { property = p1; strict = s1; loc = loc1 },
-      Assume { property = p2; strict = s2; loc = loc2 } ) ->
-    Property.equal p1 p2 && Bool.equal s1 s2 && loc1 = loc2
+  | ( Assume { property = p1; strict = s1; never_returns_normally = n1; loc = loc1 },
+      Assume { property = p2; strict = s2; never_returns_normally = n2; loc = loc2 } ) ->
+    Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal n1 n2 && loc1 = loc2
   | (Default_check | Ignore_assert_all _ | Check _ | Assume _), _ -> false
 
 let is_default : t -> bool = function
diff --git a/middle_end/flambda2/terms/check_attribute.mli b/middle_end/flambda2/terms/check_attribute.mli
index bbd3a805ae4..470b41c6c66 100644
--- a/middle_end/flambda2/terms/check_attribute.mli
+++ b/middle_end/flambda2/terms/check_attribute.mli
@@ -20,6 +20,7 @@ type t =
   | Assume of
       { property : Property.t;
         strict : bool;
+        never_returns_normally : bool;
         loc : Location.t
       }
   | Check of
diff --git a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
index b7a344e1c83..98f4ccf8ba1 100644
--- a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
+++ b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
@@ -324,8 +324,8 @@ let transl_check_attrib : Check_attribute.t -> Cmm.codegen_option list =
   function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
-  | Assume { property; strict; loc } ->
-    [Assume { property = transl_property property; strict; loc }]
+  | Assume { property; strict; never_returns_normally; loc } ->
+    [Assume { property = transl_property property; strict; never_returns_normally; loc }]
   | Check { property; strict; loc } ->
     [Check { property = transl_property property; strict; loc }]
 
diff --git a/ocaml/lambda/lambda.ml b/ocaml/lambda/lambda.ml
index 6330767ce91..70e34694007 100644
--- a/ocaml/lambda/lambda.ml
+++ b/ocaml/lambda/lambda.ml
@@ -501,6 +501,7 @@ type check_attribute =
   | Assume of { property: property;
                 strict: bool;
                 loc: Location.t;
+                never_returns_normally: bool;
               }
 
 type loop_attribute =
diff --git a/ocaml/lambda/lambda.mli b/ocaml/lambda/lambda.mli
index 05f99fb0f48..1dfc2eb72a9 100644
--- a/ocaml/lambda/lambda.mli
+++ b/ocaml/lambda/lambda.mli
@@ -396,6 +396,7 @@ type check_attribute =
   | Assume of { property: property;
                 strict: bool;
                 loc: Location.t;
+                never_returns_normally: bool;
               }
 
 type loop_attribute =
diff --git a/ocaml/lambda/printlambda.ml b/ocaml/lambda/printlambda.ml
index 8da6d532483..f7eb96c08a5 100644
--- a/ocaml/lambda/printlambda.ml
+++ b/ocaml/lambda/printlambda.ml
@@ -634,10 +634,11 @@ let check_attribute ppf check =
   | Default_check -> ()
   | Ignore_assert_all p ->
     fprintf ppf "ignore assert all %s@ " (check_property p)
-  | Assume {property=p; strict; loc = _} ->
-    fprintf ppf "assume_%s%s@ "
+  | Assume {property=p; strict; never_returns_normally; loc = _} ->
+    fprintf ppf "assume_%s%s%s@ "
       (check_property p)
       (if strict then "_strict" else "")
+      (if never_returns_normally then "_never_returns_normally" else "")
   | Check {property=p; strict; loc = _} ->
     fprintf ppf "assert_%s%s@ "
       (check_property p)
diff --git a/ocaml/lambda/translattribute.ml b/ocaml/lambda/translattribute.ml
index 61ee8eb4ed0..cdc6da76e12 100644
--- a/ocaml/lambda/translattribute.ml
+++ b/ocaml/lambda/translattribute.ml
@@ -253,9 +253,15 @@ let parse_property_attribute attr property =
         ~default:Default_check
         ~empty:(Check { property; strict = false; loc; } )
         [
-          ["assume"], Assume { property; strict = false; loc; };
+          ["assume"],
+          Assume { property; strict = false; never_returns_normally = false; loc; };
           ["strict"], Check { property; strict = true; loc; };
-          ["assume"; "strict"], Assume { property; strict = true; loc; };
+          ["assume"; "strict"],
+          Assume { property; strict = true; never_returns_normally = false; loc; };
+          ["assume"; "never_returns_normally"],
+          Assume { property; strict = false; never_returns_normally = true; loc; };
+          ["assume"; "strict"; "never_returns_normally"],
+          Assume { property; strict = true; never_returns_normally = true; loc; };
           ["ignore"], Ignore_assert_all property
         ]
         payload

From 6a8e312b4bfe85b49d2fb1560fa57a601ef284d7 Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Thu, 31 Aug 2023 18:09:48 +0100
Subject: [PATCH 4/9] Fix tests

---
 .../backend/checkmach/test_attribute_error_duplicate.output | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/backend/checkmach/test_attribute_error_duplicate.output b/tests/backend/checkmach/test_attribute_error_duplicate.output
index f7ac2d0cbe5..26e43e02361 100644
--- a/tests/backend/checkmach/test_attribute_error_duplicate.output
+++ b/tests/backend/checkmach/test_attribute_error_duplicate.output
@@ -4,12 +4,12 @@ File "test_attribute_error_duplicate.ml", line 2, characters 66-76:
 Warning 54 [duplicated-attribute]: the "zero_alloc" attribute is used more than once on this expression
 File "test_attribute_error_duplicate.ml", line 3, characters 5-15:
 Warning 47 [attribute-payload]: illegal payload for attribute 'zero_alloc'.
-It must be either 'assume', 'strict', 'assume strict', 'ignore' or empty
+It must be either 'assume', 'strict', 'assume strict', 'assume never_returns_normally', 'assume strict never_returns_normally', 'ignore' or empty
 File "test_attribute_error_duplicate.ml", line 4, characters 5-15:
 Warning 47 [attribute-payload]: illegal payload for attribute 'zero_alloc'.
-It must be either 'assume', 'strict', 'assume strict', 'ignore' or empty
+It must be either 'assume', 'strict', 'assume strict', 'assume never_returns_normally', 'assume strict never_returns_normally', 'ignore' or empty
 File "test_attribute_error_duplicate.ml", line 5, characters 5-15:
 Warning 47 [attribute-payload]: illegal payload for attribute 'zero_alloc'.
-It must be either 'assume', 'strict', 'assume strict', 'ignore' or empty
+It must be either 'assume', 'strict', 'assume strict', 'assume never_returns_normally', 'assume strict never_returns_normally', 'ignore' or empty
 File "test_attribute_error_duplicate.ml", line 1, characters 5-15:
 Error: Annotation check for zero_alloc strict failed on function Test_attribute_error_duplicate.test1 (camlTest_attribute_error_duplicate__test1_HIDE_STAMP)

From a96734f8447cc7c2252f317f0ee0812cbd67db06 Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Fri, 15 Sep 2023 10:35:12 +0100
Subject: [PATCH 5/9] Format

---
 backend/checkmach.ml                          | 16 ++++++++------
 backend/cmm_helpers.ml                        |  7 +++++-
 middle_end/flambda2/terms/check_attribute.ml  | 22 ++++++++++++-------
 .../flambda2/to_cmm/to_cmm_set_of_closures.ml |  7 +++++-
 4 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/backend/checkmach.ml b/backend/checkmach.ml
index 05547f548c5..47a1423d392 100644
--- a/backend/checkmach.ml
+++ b/backend/checkmach.ml
@@ -412,10 +412,7 @@ end = struct
 
   let expected_value t =
     let res = if t.strict then Value.safe else Value.relaxed in
-    if t.never_returns_normally then
-      { res with nor = V.Bot }
-    else
-      res
+    if t.never_returns_normally then { res with nor = V.Bot } else res
 
   let is_assume t = t.assume
 
@@ -427,7 +424,8 @@ end = struct
           match c with
           | Check { property; strict; loc } when property = spec ->
             Some { strict; assume = false; never_returns_normally = false; loc }
-          | Assume { property; strict; never_returns_normally; loc; } when property = spec ->
+          | Assume { property; strict; never_returns_normally; loc }
+            when property = spec ->
             Some { strict; assume = true; never_returns_normally; loc }
           | Ignore_assert_all property when property = spec ->
             ignore_assert_all := true;
@@ -441,8 +439,12 @@ end = struct
     | [] ->
       if !Clflags.zero_alloc_check_assert_all && not !ignore_assert_all
       then
-        Some { strict = false; assume = false; never_returns_normally = false;
-               loc = Debuginfo.to_location dbg }
+        Some
+          { strict = false;
+            assume = false;
+            never_returns_normally = false;
+            loc = Debuginfo.to_location dbg
+          }
       else None
     | [p] -> Some p
     | _ :: _ ->
diff --git a/backend/cmm_helpers.ml b/backend/cmm_helpers.ml
index ec9e391695e..6ee8ad70b1b 100644
--- a/backend/cmm_helpers.ml
+++ b/backend/cmm_helpers.ml
@@ -4513,7 +4513,12 @@ let transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list = function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
   | Assume { property; strict; never_returns_normally; loc } ->
-    [Assume { property = transl_property property; strict; never_returns_normally; loc }]
+    [ Assume
+        { property = transl_property property;
+          strict;
+          never_returns_normally;
+          loc
+        } ]
   | Check { property; strict; loc } ->
     [Check { property = transl_property property; strict; loc }]
 
diff --git a/middle_end/flambda2/terms/check_attribute.ml b/middle_end/flambda2/terms/check_attribute.ml
index 7b413d4c223..19a4f48fba2 100644
--- a/middle_end/flambda2/terms/check_attribute.ml
+++ b/middle_end/flambda2/terms/check_attribute.ml
@@ -41,20 +41,23 @@ let print ppf t =
   | Ignore_assert_all property ->
     Format.fprintf ppf "@[ignore %a@]" Property.print property
   | Assume { property; strict; never_returns_normally; loc = _ } ->
-    Format.fprintf ppf "@[assume_%a%s%s@]"
-      Property.print property
+    Format.fprintf ppf "@[assume_%a%s%s@]" Property.print property
       (if strict then "_strict" else "")
       (if never_returns_normally then "_never_returns_normally" else "")
   | Check { property; strict; loc = _ } ->
-    Format.fprintf ppf "@[assert_%a%s@]"
-      Property.print property
+    Format.fprintf ppf "@[assert_%a%s@]" Property.print property
       (if strict then "_strict" else "")
 
 let from_lambda : Lambda.check_attribute -> t = function
   | Default_check -> Default_check
   | Ignore_assert_all p -> Ignore_assert_all (Property.from_lambda p)
   | Assume { property; strict; never_returns_normally; loc } ->
-    Assume { property = Property.from_lambda property; strict; never_returns_normally; loc }
+    Assume
+      { property = Property.from_lambda property;
+        strict;
+        never_returns_normally;
+        loc
+      }
   | Check { property; strict; loc } ->
     Check { property = Property.from_lambda property; strict; loc }
 
@@ -65,11 +68,14 @@ let equal x y =
   | ( Check { property = p1; strict = s1; loc = loc1 },
       Check { property = p2; strict = s2; loc = loc2 } ) ->
     Property.equal p1 p2 && Bool.equal s1 s2 && loc1 = loc2
-  | ( Assume { property = p1; strict = s1; never_returns_normally = n1; loc = loc1 },
-      Assume { property = p2; strict = s2; never_returns_normally = n2; loc = loc2 } ) ->
+  | ( Assume
+        { property = p1; strict = s1; never_returns_normally = n1; loc = loc1 },
+      Assume
+        { property = p2; strict = s2; never_returns_normally = n2; loc = loc2 }
+    ) ->
     Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal n1 n2 && loc1 = loc2
   | (Default_check | Ignore_assert_all _ | Check _ | Assume _), _ -> false
 
 let is_default : t -> bool = function
   | Default_check -> true
-  | Ignore_assert_all _ | Check _ | Assume _  -> false
+  | Ignore_assert_all _ | Check _ | Assume _ -> false
diff --git a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
index 98f4ccf8ba1..6618f4a7f7a 100644
--- a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
+++ b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
@@ -325,7 +325,12 @@ let transl_check_attrib : Check_attribute.t -> Cmm.codegen_option list =
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
   | Assume { property; strict; never_returns_normally; loc } ->
-    [Assume { property = transl_property property; strict; never_returns_normally; loc }]
+    [ Assume
+        { property = transl_property property;
+          strict;
+          never_returns_normally;
+          loc
+        } ]
   | Check { property; strict; loc } ->
     [Check { property = transl_property property; strict; loc }]
 

From 70cedd728ec2591cf9db3bf0d17e827403006613 Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Mon, 18 Sep 2023 13:09:45 +0100
Subject: [PATCH 6/9] Add a test for never_returns_normally using failwithf,
 invalid_argf

---
 tests/backend/checkmach/t.ml | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/backend/checkmach/t.ml b/tests/backend/checkmach/t.ml
index e66d884f886..329c5392d34 100644
--- a/tests/backend/checkmach/t.ml
+++ b/tests/backend/checkmach/t.ml
@@ -239,3 +239,15 @@ let[@zero_alloc assume][@inline always] test40 x = (x,x)
 
 let[@zero_alloc] test41 b x =
   if b then test40 (x+1) else test40 (x*2)
+
+
+module Never_returns_normally = struct
+  let[@inline never][@zero_alloc assume never_returns_normally] failwithf fmt =
+   Printf.ksprintf (fun s () -> failwith s) fmt
+
+  let[@inline never][@zero_alloc assume never_returns_normally] invalid_argf fmt =
+    Printf.ksprintf (fun s () -> invalid_arg s) fmt
+
+ let[@zero_alloc] foo x = failwithf "%d" x
+ let[@zero_alloc] bar x y = invalid_argf "%d" (x+y)
+end

From dff24a3b2b4e4865e0503271b94c0ec434d2965b Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Mon, 18 Sep 2023 13:13:55 +0100
Subject: [PATCH 7/9] Add a test that currently fails but will be fixed in a
 subsequent PR.

never_returns_normally is not propagated on Debuginfo.t
---
 tests/backend/checkmach/dune.inc              | 19 +++++++++++++++++++
 tests/backend/checkmach/gen/gen_dune.ml       |  1 +
 .../checkmach/test_never_returns_normally.ml  | 11 +++++++++++
 .../test_never_returns_normally.output        |  4 ++++
 4 files changed, 35 insertions(+)
 create mode 100644 tests/backend/checkmach/test_never_returns_normally.ml
 create mode 100644 tests/backend/checkmach/test_never_returns_normally.output

diff --git a/tests/backend/checkmach/dune.inc b/tests/backend/checkmach/dune.inc
index 03beb6b0fb2..8b42234c071 100644
--- a/tests/backend/checkmach/dune.inc
+++ b/tests/backend/checkmach/dune.inc
@@ -509,3 +509,22 @@
  (enabled_if (and (= %{context_name} "main") %{ocaml-config:flambda}))
  (deps test_warning199.output test_warning199.output.corrected)
  (action (diff test_warning199.output test_warning199.output.corrected)))
+
+(rule
+ (enabled_if (= %{context_name} "main"))
+ (targets test_never_returns_normally.output.corrected)
+ (deps (:ml test_never_returns_normally.ml) filter.sh)
+ (action
+   (with-outputs-to test_never_returns_normally.output.corrected
+    (pipe-outputs
+    (with-accepted-exit-codes 2
+     (run %{bin:ocamlopt.opt} %{ml} -g -color never -error-style short -c
+          -zero-alloc-check -checkmach-details-cutoff 20 -O3))
+    (run "./filter.sh")
+   ))))
+
+(rule
+ (alias   runtest)
+ (enabled_if (= %{context_name} "main"))
+ (deps test_never_returns_normally.output test_never_returns_normally.output.corrected)
+ (action (diff test_never_returns_normally.output test_never_returns_normally.output.corrected)))
diff --git a/tests/backend/checkmach/gen/gen_dune.ml b/tests/backend/checkmach/gen/gen_dune.ml
index 84bf706007a..04b353342e0 100644
--- a/tests/backend/checkmach/gen/gen_dune.ml
+++ b/tests/backend/checkmach/gen/gen_dune.ml
@@ -117,4 +117,5 @@ let () =
   print_test ~flambda_only:true ~deps:"t1.ml";
   (* closure does not delete dead functions *)
   print_test_expected_output ~cutoff:default_cutoff ~flambda_only:true ~extra_dep:(Some "test_warning199.mli") ~exit_code:0 "test_warning199";
+  print_test_expected_output ~cutoff:default_cutoff ~flambda_only:false ~extra_dep:None ~exit_code:2 "test_never_returns_normally";
   ()
diff --git a/tests/backend/checkmach/test_never_returns_normally.ml b/tests/backend/checkmach/test_never_returns_normally.ml
new file mode 100644
index 00000000000..0535fc034c4
--- /dev/null
+++ b/tests/backend/checkmach/test_never_returns_normally.ml
@@ -0,0 +1,11 @@
+(* CR gyorsh: Does not work without [@inline]. Fix by propagating never_returns_normally
+   to check_mach as an extra annotation on [Debuginfo.t]. *)
+let[@zero_alloc assume never_returns_normally] failwithf fmt =
+  Printf.ksprintf (fun s () -> failwith s) fmt
+
+let[@zero_alloc assume never_returns_normally] invalid_argf fmt =
+  Printf.ksprintf (fun s () -> invalid_arg s) fmt
+
+let[@zero_alloc] foo x = failwithf "%d" x
+let[@zero_alloc] bar x y = invalid_argf "%d" (x+y)
+
diff --git a/tests/backend/checkmach/test_never_returns_normally.output b/tests/backend/checkmach/test_never_returns_normally.output
new file mode 100644
index 00000000000..0e4438f95e1
--- /dev/null
+++ b/tests/backend/checkmach/test_never_returns_normally.output
@@ -0,0 +1,4 @@
+File "test_never_returns_normally.ml", line 9, characters 5-15:
+Error: Annotation check for zero_alloc failed on function Test_never_returns_normally.foo (camlTest_never_returns_normally__foo_HIDE_STAMP)
+File "test_never_returns_normally.ml", line 9, characters 25-41:
+  indirect tailcall

From d41002dde7bba12c0c5e64112115ad43b2968a6c Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Mon, 18 Sep 2023 13:49:32 +0100
Subject: [PATCH 8/9] Fix the new test failure with closure: only run the test
 w/flambda

---
 tests/backend/checkmach/dune.inc        | 4 ++--
 tests/backend/checkmach/gen/gen_dune.ml | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/backend/checkmach/dune.inc b/tests/backend/checkmach/dune.inc
index 8b42234c071..a67e85bdc28 100644
--- a/tests/backend/checkmach/dune.inc
+++ b/tests/backend/checkmach/dune.inc
@@ -511,7 +511,7 @@
  (action (diff test_warning199.output test_warning199.output.corrected)))
 
 (rule
- (enabled_if (= %{context_name} "main"))
+ (enabled_if (and (= %{context_name} "main") %{ocaml-config:flambda}))
  (targets test_never_returns_normally.output.corrected)
  (deps (:ml test_never_returns_normally.ml) filter.sh)
  (action
@@ -525,6 +525,6 @@
 
 (rule
  (alias   runtest)
- (enabled_if (= %{context_name} "main"))
+ (enabled_if (and (= %{context_name} "main") %{ocaml-config:flambda}))
  (deps test_never_returns_normally.output test_never_returns_normally.output.corrected)
  (action (diff test_never_returns_normally.output test_never_returns_normally.output.corrected)))
diff --git a/tests/backend/checkmach/gen/gen_dune.ml b/tests/backend/checkmach/gen/gen_dune.ml
index 04b353342e0..8b9a56344cf 100644
--- a/tests/backend/checkmach/gen/gen_dune.ml
+++ b/tests/backend/checkmach/gen/gen_dune.ml
@@ -117,5 +117,5 @@ let () =
   print_test ~flambda_only:true ~deps:"t1.ml";
   (* closure does not delete dead functions *)
   print_test_expected_output ~cutoff:default_cutoff ~flambda_only:true ~extra_dep:(Some "test_warning199.mli") ~exit_code:0 "test_warning199";
-  print_test_expected_output ~cutoff:default_cutoff ~flambda_only:false ~extra_dep:None ~exit_code:2 "test_never_returns_normally";
+  print_test_expected_output ~cutoff:default_cutoff ~flambda_only:true ~extra_dep:None ~exit_code:2 "test_never_returns_normally";
   ()

From 6a90d16cba41f1bee0fe6fd242c9216f109efe25 Mon Sep 17 00:00:00 2001
From: Greta Yorsh <gyorsh@janestreet.com>
Date: Wed, 20 Sep 2023 09:55:34 +0100
Subject: [PATCH 9/9] Use Location.compare instead of Poly compare

---
 middle_end/flambda2/terms/check_attribute.ml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/middle_end/flambda2/terms/check_attribute.ml b/middle_end/flambda2/terms/check_attribute.ml
index 19a4f48fba2..2c931cb4285 100644
--- a/middle_end/flambda2/terms/check_attribute.ml
+++ b/middle_end/flambda2/terms/check_attribute.ml
@@ -67,13 +67,14 @@ let equal x y =
   | Ignore_assert_all p1, Ignore_assert_all p2 -> Property.equal p1 p2
   | ( Check { property = p1; strict = s1; loc = loc1 },
       Check { property = p2; strict = s2; loc = loc2 } ) ->
-    Property.equal p1 p2 && Bool.equal s1 s2 && loc1 = loc2
+    Property.equal p1 p2 && Bool.equal s1 s2 && Location.compare loc1 loc2 = 0
   | ( Assume
         { property = p1; strict = s1; never_returns_normally = n1; loc = loc1 },
       Assume
         { property = p2; strict = s2; never_returns_normally = n2; loc = loc2 }
     ) ->
-    Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal n1 n2 && loc1 = loc2
+    Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal n1 n2
+    && Location.compare loc1 loc2 = 0
   | (Default_check | Ignore_assert_all _ | Check _ | Assume _), _ -> false
 
 let is_default : t -> bool = function