Configurable "default" default precision

lukstafi · lukstafi · commit 1b56fc4c86e3 · 2025-05-30T22:18:25.000+02:00
diff --git a/arrayjit/lib/ops.ml b/arrayjit/lib/ops.ml
@@ -90,6 +90,9 @@ let prec_string = function
   | Single_prec _ -> "single"
   | Double_prec _ -> "double"
 
+let prec_of_string s =
+  prec_of_sexp (Sexp.Atom (String.(capitalize @@ lowercase s) ^ "_prec"))
+
 let equal_prec p1 p2 =
   match (p1, p2) with
   | Void_prec, Void_prec -> true
diff --git a/arrayjit/lib/tnode.ml b/arrayjit/lib/tnode.ml
@@ -562,12 +562,15 @@ let create ?default_prec ~id ~label ~dims init_op =
   Registry.add registry tn;
   tn
 
+let initial_default_prec =
+  Ops.prec_of_string (Utils.get_global_arg ~default:"single" ~arg_name:"default_prec")
+
 let find =
   let mock =
     {
       array = lazy None;
-      prec = lazy Ops.single;
-      delayed_prec_unsafe = Specified Ops.single;
+      prec = lazy initial_default_prec;
+      delayed_prec_unsafe = Specified initial_default_prec;
       dims = lazy [||];
       size_in_bytes = lazy 0;
       id = -1;
diff --git a/arrayjit/test/ocannl_config b/arrayjit/test/ocannl_config
@@ -3,4 +3,5 @@ log_main_domain_to_stdout=true
 backend=multicore_cc
 log_level=0
 print_decimals_precision=2
-prefer_backend_uniformity=true
+prefer_backend_uniformity=true
+default_prec=bfloat16
diff --git a/lib/tensor.ml b/lib/tensor.ml
@@ -99,8 +99,11 @@ let iter_embedded ~f t =
   Set.iter ~f t.forward.embedded_nodes;
   Option.iter t.diff ~f:(fun diff -> Set.iter ~f diff.backprop.embedded_nodes)
 
-let default_value_prec = ref Ir.Ops.single
-let default_grad_prec = ref Ir.Ops.single
+let initial_default_prec =
+  Ir.Ops.prec_of_string (Utils.get_global_arg ~default:"single" ~arg_name:"default_prec")
+
+let default_value_prec = ref initial_default_prec
+let default_grad_prec = ref initial_default_prec
 
 exception Session_error of string * t option [@@deriving sexp]
 
diff --git a/ocannl_config.example b/ocannl_config.example
@@ -168,4 +168,8 @@ log_file_stem=debug
 # It is useful for testing to have outputs more uniform across backends even if that criples
 # some backends. Currently, this setting only affects logging from routines to accomodate Metal's
 # shortcoming.
-prefer_backend_uniformity=false
+prefer_backend_uniformity=false
+
+# The initial value for the default precisions for tensors. The default precisions for values and
+# gradients can be changed separately via the `Tensor` API.
+default_prec=single
diff --git a/test/config/ocannl_config b/test/config/ocannl_config
@@ -3,4 +3,5 @@ log_main_domain_to_stdout=true
 backend=multicore_cc
 log_level=0
 print_decimals_precision=2
-prefer_backend_uniformity=true
+prefer_backend_uniformity=true
+default_prec=bfloat16
