Add Float16 pattern matches

owlbarn · May 17, 2024 · ef69ebe · ef69ebe
1 parent c3230eb
commit ef69ebe
Show file tree

Hide file tree

Showing 7 changed files with 166 additions and 0 deletions.
diff --git a/owl-base.opam b/owl-base.opam
@@ -17,4 +17,5 @@ depends: [
   "ocaml" {>= "4.10.0"}
   "base-bigarray"
   "dune" {>= "2.0.0"}
+  "cppo"
 ]
diff --git a/src/base/core/owl_const.ml → src/base/core/owl_const.cppo.ml b/src/base/core/owl_const.ml → src/base/core/owl_const.cppo.ml
@@ -46,6 +46,9 @@ let eps = 1e-15
 (** Functions that return constants using Bigarray kind *)
 
 let zero : type a b. (a, b) kind -> a = function
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16        -> 0.0
+#endif
   | Float32        -> 0.0
   | Complex32      -> Complex.zero
   | Float64        -> 0.0
@@ -62,6 +65,9 @@ let zero : type a b. (a, b) kind -> a = function
 
 
 let one : type a b. (a, b) kind -> a = function
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16        -> 1.0
+#endif
   | Float32        -> 1.0
   | Complex32      -> Complex.one
   | Float64        -> 1.0
@@ -78,6 +84,9 @@ let one : type a b. (a, b) kind -> a = function
 
 
 let neg_one : type a b. (a, b) kind -> a = function
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16        -> -1.0
+#endif
   | Float32        -> -1.0
   | Float64        -> -1.0
   | Complex32      -> Complex.{ re = -1.; im = 0. }
@@ -94,6 +103,9 @@ let neg_one : type a b. (a, b) kind -> a = function
 
 
 let pos_inf : type a b. (a, b) kind -> a = function
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16   -> infinity
+#endif
   | Float32   -> infinity
   | Float64   -> infinity
   | Complex32 -> Complex.{ re = infinity; im = infinity }
@@ -102,6 +114,9 @@ let pos_inf : type a b. (a, b) kind -> a = function
 
 
 let neg_inf : type a b. (a, b) kind -> a = function
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16   -> neg_infinity
+#endif
   | Float32   -> neg_infinity
   | Float64   -> neg_infinity
   | Complex32 -> Complex.{ re = neg_infinity; im = neg_infinity }

diff --git a/src/base/dune b/src/base/dune
@@ -20,6 +20,16 @@
 
 (copy_files# misc/*)
 
+(rule
+  (targets owl_const.ml)
+  (deps (:out owl_const.cppo.ml))
+  (action (run cppo -V OCAML:%{ocaml_version} %{out} -o %{targets})))
+
+(rule
+  (targets owl_utils_ndarray.ml)
+  (deps (:out owl_utils_ndarray.cppo.ml))
+  (action (run cppo -V OCAML:%{ocaml_version} %{out} -o %{targets})))
+
 (library
  (name owl_base)
  (public_name owl-base)

diff --git a/src/base/misc/owl_utils_ndarray.ml → src/base/misc/owl_utils_ndarray.cppo.ml b/src/base/misc/owl_utils_ndarray.ml → src/base/misc/owl_utils_ndarray.cppo.ml
@@ -16,6 +16,9 @@ let elt_to_str : type a b. (a, b) kind -> a -> string = function
   | Int            -> fun v -> Printf.sprintf "%i" v
   | Int32          -> fun v -> Printf.sprintf "%ld" v
   | Int64          -> fun v -> Printf.sprintf "%Ld" v
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16        -> fun v -> Printf.sprintf "%G" v
+#endif
   | Float32        -> fun v -> Printf.sprintf "%G" v
   | Float64        -> fun v -> Printf.sprintf "%G" v
   | Complex32      -> fun v -> Printf.sprintf "(%G, %Gi)" Complex.(v.re) Complex.(v.im)
@@ -33,6 +36,9 @@ let elt_of_str : type a b. (a, b) kind -> string -> a = function
   | Int            -> fun v -> int_of_string v
   | Int32          -> fun v -> Int32.of_string v
   | Int64          -> fun v -> Int64.of_string v
+#if OCAML_VERSION >= (5, 2, 0)
+  | Float16        -> fun v -> float_of_string v
+#endif
   | Float32        -> fun v -> float_of_string v
   | Float64        -> fun v -> float_of_string v
   | Complex32      -> fun v ->

diff --git a/src/owl/cblas/owl_cblas_basic.ml → src/owl/cblas/owl_cblas_basic.cppo.ml b/src/owl/cblas/owl_cblas_basic.ml → src/owl/cblas/owl_cblas_basic.cppo.ml
@@ -856,6 +856,10 @@ let symv
   let _y = bigarray_start Ctypes_static.Array1 y in
   let _a = bigarray_start Ctypes_static.Array1 a in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 ->
+     C.ssymv ~order:_layout ~uplo:_uplo ~n ~alpha ~a:_a ~lda ~x:_x ~incx ~beta ~y:_y ~incy
+#endif
   | Bigarray.Float32 ->
     C.ssymv ~order:_layout ~uplo:_uplo ~n ~alpha ~a:_a ~lda ~x:_x ~incx ~beta ~y:_y ~incy
   | Bigarray.Float64 ->
@@ -887,6 +891,22 @@ let sbmv
   let _y = bigarray_start Ctypes_static.Array1 y in
   let _a = bigarray_start Ctypes_static.Array1 a in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 ->
+     C.ssbmv
+       ~order:_layout
+       ~uplo:_uplo
+       ~n
+       ~k
+       ~alpha
+       ~a:_a
+       ~lda
+       ~x:_x
+       ~incx
+       ~beta
+       ~y:_y
+       ~incy
+#endif
   | Bigarray.Float32 ->
     C.ssbmv
       ~order:_layout
@@ -941,6 +961,10 @@ let spmv
   let _y = bigarray_start Ctypes_static.Array1 y in
   let _ap = bigarray_start Ctypes_static.Array1 ap in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 ->
+     C.sspmv ~order:_layout ~uplo:_uplo ~n ~alpha ~ap:_ap ~x:_x ~incx ~beta ~y:_y ~incy
+#endif
   | Bigarray.Float32 ->
     C.sspmv ~order:_layout ~uplo:_uplo ~n ~alpha ~ap:_ap ~x:_x ~incx ~beta ~y:_y ~incy
   | Bigarray.Float64 ->
@@ -1049,6 +1073,9 @@ let syr
   let _x = bigarray_start Ctypes_static.Array1 x in
   let _a = bigarray_start Ctypes_static.Array1 a in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 -> C.ssyr ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~a:_a ~lda
+#endif
   | Bigarray.Float32 -> C.ssyr ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~a:_a ~lda
   | Bigarray.Float64 -> C.dsyr ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~a:_a ~lda
 
@@ -1072,6 +1099,9 @@ let spr
   let _x = bigarray_start Ctypes_static.Array1 x in
   let _ap = bigarray_start Ctypes_static.Array1 ap in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 -> C.sspr ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~ap:_ap
+#endif
   | Bigarray.Float32 -> C.sspr ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~ap:_ap
   | Bigarray.Float64 -> C.dspr ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~ap:_ap
 
@@ -1099,6 +1129,10 @@ let syr2
   let _y = bigarray_start Ctypes_static.Array1 y in
   let _a = bigarray_start Ctypes_static.Array1 a in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 ->
+     C.ssyr2 ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~y:_y ~incy ~a:_a ~lda
+#endif
   | Bigarray.Float32 ->
     C.ssyr2 ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~y:_y ~incy ~a:_a ~lda
   | Bigarray.Float64 ->
@@ -1127,6 +1161,10 @@ let spr2
   let _y = bigarray_start Ctypes_static.Array1 y in
   let _a = bigarray_start Ctypes_static.Array1 a in
   match Bigarray.Array1.kind x with
+#if OCAML_VERSION >= (5, 2, 0)
+  | Bigarray.Float16 ->
+     C.sspr2 ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~y:_y ~incy ~a:_a
+#endif
   | Bigarray.Float32 ->
     C.sspr2 ~order:_layout ~uplo:_uplo ~n ~alpha ~x:_x ~incx ~y:_y ~incy ~a:_a
   | Bigarray.Float64 ->

diff --git a/src/owl/dune b/src/owl/dune
@@ -38,6 +38,16 @@
 
 (copy_files# signal/*)
 
+(rule
+  (targets owl_cblas_basic.ml)
+  (deps (:out owl_cblas_basic.cppo.ml))
+  (action (run cppo -V OCAML:%{ocaml_version} %{out} -o %{targets})))
+
+(rule
+  (targets owl_lapacke.ml)
+  (deps (:out owl_lapacke.cppo.ml))
+  (action (run cppo -V OCAML:%{ocaml_version} %{out} -o %{targets})))
+
 (library
  (name owl)
  (public_name owl)