Detect FP16 constant overflow

lukstafi · lukstafi · commit db206b002815 · 2024-09-12T15:19:12.000+02:00
diff --git a/arrayjit/lib/low_level.ml b/arrayjit/lib/low_level.ml
@@ -736,7 +736,47 @@ let simplify_llc llc =
         let result = Unop (op, v) in
         if equal_float_t llv v then result else loop_float result
   in
-  loop_proc llc
+  let check_constant =
+    match Utils.settings.check_half_prec_constants_cutoff with
+    | None -> fun _prec _c -> ()
+    | Some cutoff -> (
+        fun tn c ->
+          match tn.Tn.prec with
+          | Ops.Half_prec _ ->
+              if Float.(abs c >= cutoff) then
+                raise
+                @@ Utils.User_error
+                     ("Constant " ^ Float.to_string c
+                    ^ " is too big for FP16 aka. half precision, risk of overflow; increase \
+                       precision of tensor node " ^ Tn.debug_name tn)
+          | _ -> ())
+  in
+  let rec check_proc llc =
+    let loop = check_proc in
+    match llc with
+    | Seq (c1, c2) ->
+        loop c1;
+        loop c2
+    | For_loop { body; _ } -> loop body
+    | Zero_out _ -> ()
+    | Set { tn; llv; _ } -> check_float tn llv
+    | Set_local (id, llv) -> check_float id.tn llv
+    | Noop | Comment _ | Staged_compilation _ -> ()
+  and check_float tn llv =
+    let loop = check_float tn in
+    match llv with
+    | Constant c -> check_constant tn c
+    | Local_scope { body; _ } -> check_proc body
+    | Binop (_, v1, v2) ->
+        loop v1;
+        loop v2
+    | Unop (_, v) -> loop v
+    | Embed_index (Indexing.Fixed_idx i) -> check_constant tn (Float.of_int i)
+    | Embed_index _ | Get_local _ | Get_global (_, _) | Get (_, _) -> ()
+  in
+  let result = loop_proc llc in
+  if Option.is_some Utils.settings.check_half_prec_constants_cutoff then check_proc result;
+  result
 
 type traced_store = (Tn.t, traced_array) Base.Hashtbl.t [@@deriving sexp_of]
 
diff --git a/arrayjit/lib/utils.ml b/arrayjit/lib/utils.ml
@@ -40,6 +40,9 @@ type settings = {
   mutable fixed_state_for_init : int option;
   mutable print_decimals_precision : int;
       (** When rendering arrays etc., outputs this many decimal digits. *)
+  mutable check_half_prec_constants_cutoff : float option;
+      (** If given, generic code optimization should fail if a half precision FP16 constant exceeds
+          the cutoff. *)
 }
 [@@deriving sexp]
 
@@ -51,6 +54,7 @@ let settings =
     output_debug_files_in_build_directory = false;
     fixed_state_for_init = None;
     print_decimals_precision = 2;
+    check_half_prec_constants_cutoff = Some (2. **. 14.);
   }
 
 let accessed_global_args = Hash_set.create (module String)
@@ -314,7 +318,10 @@ let restore_settings () =
     (let seed = get_global_arg ~arg_name:"fixed_state_for_init" ~default:"" in
      if String.is_empty seed then None else Some (Int.of_string seed));
   settings.print_decimals_precision <-
-    Int.of_string @@ get_global_arg ~arg_name:"print_decimals_precision" ~default:"2"
+    Int.of_string @@ get_global_arg ~arg_name:"print_decimals_precision" ~default:"2";
+  settings.check_half_prec_constants_cutoff <-
+    Float.of_string_opt
+    @@ get_global_arg ~arg_name:"check_half_prec_constants_cutoff" ~default:"16384.0"
 
 let () = restore_settings ()
 let with_runtime_debug () = settings.output_debug_files_in_build_directory && settings.log_level > 1
diff --git a/bin/moons_benchmark.ml b/bin/moons_benchmark.ml
@@ -176,14 +176,14 @@ let _suspend () =
        ~backend_name:"gccjit" ~value_prec:CDSL.single ~grad_prec:CDSL.double ()
 
 let cuda_benchmarks =
-  List.concat_map [ 1; 3; 6; 12 (* ; 16; 32; 64 *) ] ~f:(fun num_devices ->
+  List.concat_map [ (* 1; *) 3 (* ; 6; 12; 16; 32; 64 *) ] ~f:(fun num_devices ->
       List.concat_map
-        [ 3 * 32; 3 * 64; 3 * 128 ]
+        [ 3 * 32(* ; 3 * 64; 3 * 128 *)]
         ~f:(fun batch_size ->
-          List.concat_map [ 0; 1; (* 2; *) 3 ] ~f:(fun inlining_cutoff ->
-              List.concat_map [ 1; 3 (* ; 7 *) ] ~f:(fun seed ->
-                  List.concat_map [ (* "gccjit" ;  "cc"; *) "cuda" ] ~f:(fun backend_name ->
-                      List.concat_map [ (* CDSL.double; CDSL.single; *) CDSL.half ]
+          List.concat_map [  (* 0; 1; 2; *) 3 ] ~f:(fun inlining_cutoff ->
+              List.concat_map [ 1 (* ; 3 ; 7 *) ] ~f:(fun seed ->
+                  List.concat_map [ (* "gccjit" ;   *) "cc";(* "cuda" *) ] ~f:(fun backend_name ->
+                      List.concat_map [ (* CDSL.double;  CDSL.single; *) CDSL.half ]
                         ~f:(fun value_prec ->
                           [
                             classify_moons ~seed ~on_device:true ~inlining_cutoff ~num_devices
