Optimisation des acces aux bigarrays

git-svn-id: http://caml.inria.fr/svn/ocaml/trunk@2873 f963ae5c-01c2-4b8c-9fe0-0dff7051ff02

asmcomp/cmmgen.ml

@@ -120,6 +120,14 @@ let sub_int c1 c2 =
   | (c1, c2) ->
       Cop(Csubi, [c1; c2])
 
+let mul_int c1 c2 =
+  match (c1, c2) with
+    (Cconst_int 0, _) -> c1
+  | (Cconst_int 1, _) -> c2
+  | (_, Cconst_int 0) -> c2
+  | (_, Cconst_int 1) -> c1
+  | (_, _) -> Cop(Cmuli, [c1; c2])
+
 let tag_int = function
     Cconst_int n -> int_const n
   | c -> Cop(Caddi, [Cop(Clsl, [c; Cconst_int 1]); Cconst_int 1])
@@ -475,15 +483,80 @@ let unbox_unsigned_int bi arg =
       Cop(Cload(if bi = Pint32 then Thirtytwo_unsigned else Word),
           [Cop(Cadda, [arg; Cconst_int size_addr])])
 
+(* Big arrays *)
+
+let bigarray_indexing elt_kind layout b args =
+  let rec ba_indexing dim_ofs delta_ofs = function
+    [] -> assert false
+  | [arg] ->
+      bind "idx" (untag_int arg)
+        (fun idx ->
+          Csequence(
+            Cop(Ccheckbound, [Cop(Cload Word,[field_address b dim_ofs]); idx]),
+            idx))
+  | arg1 :: argl ->
+      let rem = ba_indexing (dim_ofs + delta_ofs) delta_ofs argl in
+      bind "idx" (untag_int arg1)
+        (fun idx ->
+          bind "bound" (Cop(Cload Word, [field_address b dim_ofs]))
+          (fun bound ->
+            Csequence(Cop(Ccheckbound, [bound; idx]),
+                      add_int (mul_int rem bound) idx))) in
+  let offset =
+    match layout with
+      Pbigarray_unknown_layout ->
+        assert false
+    | Pbigarray_c_layout ->
+        ba_indexing (4 + List.length args) (-1) (List.rev args)
+    | Pbigarray_fortran_layout ->
+        ba_indexing 5 1 (List.map (fun idx -> sub_int idx (Cconst_int 2)) args)
+  and elt_size =
+    match elt_kind with
+      Pbigarray_unknown -> assert false
+    | Pbigarray_float32 -> 4
+    | Pbigarray_float64 -> 8
+    | Pbigarray_sint8 -> 1
+    | Pbigarray_uint8 -> 1
+    | Pbigarray_sint16 -> 2
+    | Pbigarray_uint16 -> 2
+    | Pbigarray_int32 -> 4
+    | Pbigarray_int64 -> 8
+    | Pbigarray_caml_int -> size_int
+    | Pbigarray_native_int -> size_int in
+  let byte_offset =
+    if elt_size = 1
+    then offset
+    else Cop(Clsl, [offset; Cconst_int(log2 elt_size)]) in
+  Cop(Cadda, [Cop(Cload Word, [field_address b 1]); byte_offset])
+
+let bigarray_word_kind = function
+    Pbigarray_unknown -> assert false
+  | Pbigarray_float32 -> Single
+  | Pbigarray_float64 -> Double
+  | Pbigarray_sint8 -> Byte_signed
+  | Pbigarray_uint8 -> Byte_unsigned
+  | Pbigarray_sint16 -> Sixteen_signed
+  | Pbigarray_uint16 -> Sixteen_unsigned
+  | Pbigarray_int32 -> Thirtytwo_signed
+  | Pbigarray_int64 -> Word
+  | Pbigarray_caml_int -> Word
+  | Pbigarray_native_int -> Word
+
+let bigarray_get elt_kind layout b args =
+  Cop(Cload (bigarray_word_kind elt_kind),
+      [bigarray_indexing elt_kind layout b args])
+
+let bigarray_set elt_kind layout b args newval =
+  Cop(Cstore (bigarray_word_kind elt_kind),
+      [bigarray_indexing elt_kind layout b args; newval])
+
 (* Simplification of some primitives into C calls *)
 
 let default_prim name =
   { prim_name = name; prim_arity = 0 (*ignored*);
     prim_alloc = true; prim_native_name = ""; prim_native_float = false }
 
-let simplif_primitive p =
-  if size_int = 8 then p else
-  match p with
+let simplif_primitive_32bits = function
     Pbintofint Pint64 -> Pccall (default_prim "int64_of_int")
   | Pintofbint Pint64 -> Pccall (default_prim "int64_to_int")
   | Pnegbint Pint64 -> Pccall (default_prim "int64_neg")
@@ -504,8 +577,25 @@ let simplif_primitive p =
   | Pbintcomp(Pint64, Lambda.Cgt) -> Pccall (default_prim "greaterthan")
   | Pbintcomp(Pint64, Lambda.Cle) -> Pccall (default_prim "lessequal")
   | Pbintcomp(Pint64, Lambda.Cge) -> Pccall (default_prim "greaterequal")
+  | Pbigarrayref(n, Pbigarray_int64, layout) ->
+      Pccall (default_prim ("bigarray_get_" ^ string_of_int n))
+  | Pbigarrayset(n, Pbigarray_int64, layout) ->
+      Pccall (default_prim ("bigarray_set_" ^ string_of_int n))
   | p -> p
 
+let simplif_primitive p =
+  match p with
+    Pbigarrayref(n, Pbigarray_unknown, layout) ->
+      Pccall (default_prim ("bigarray_get_" ^ string_of_int n))
+  | Pbigarrayset(n, Pbigarray_unknown, layout) ->
+      Pccall (default_prim ("bigarray_set_" ^ string_of_int n))
+  | Pbigarrayref(n, kind, Pbigarray_unknown_layout) ->
+      Pccall (default_prim ("bigarray_get_" ^ string_of_int n))
+  | Pbigarrayset(n, kind, Pbigarray_unknown_layout) ->
+      Pccall (default_prim ("bigarray_set_" ^ string_of_int n))
+  | p ->
+      if size_int = 8 then p else simplif_primitive_32bits p
+
 (* Translate an expression *)
 
 let functions = (Queue.create() : (string * Ident.t list * ulambda) Queue.t)
@@ -625,6 +715,29 @@ let rec transl = function
               Cop(Calloc, alloc_floatarray_header (List.length args) ::
                           List.map transl_unbox_float args)
           end
+      | (Pbigarrayref(num_dims, elt_kind, layout), arg1 :: argl) ->
+          let elt =
+            bigarray_get elt_kind layout
+                         (transl arg1) (List.map transl argl) in
+          begin match elt_kind with
+            Pbigarray_float32 | Pbigarray_float64 -> box_float elt
+          | Pbigarray_int32 -> box_int Pint32 elt
+          | Pbigarray_int64 -> box_int Pint64 elt
+          | Pbigarray_native_int -> box_int Pnativeint elt
+          | _ -> tag_int elt
+          end
+      | (Pbigarrayset(num_dims, elt_kind, layout), arg1 :: argl) ->
+          let (argidx, argnewval) = split_last argl in
+          bigarray_set elt_kind layout
+            (transl arg1)
+            (List.map transl argidx)
+            (match elt_kind with
+                Pbigarray_float32 | Pbigarray_float64 ->
+                  transl_unbox_float argnewval
+              | Pbigarray_int32 -> transl_unbox_int Pint32 argnewval
+              | Pbigarray_int64 -> transl_unbox_int Pint64 argnewval
+              | Pbigarray_native_int -> transl_unbox_int Pnativeint argnewval
+              | _ -> untag_int (transl argnewval))
       | (p, [arg]) ->
           transl_prim_1 p arg
       | (p, [arg1; arg2]) ->

bytecomp/bytegen.ml

@@ -308,6 +308,8 @@ let comp_primitive p args =
   | Pbintcomp(bi, Cgt) -> Kccall("greaterthan", 2)
   | Pbintcomp(bi, Cle) -> Kccall("lessequal", 2)
   | Pbintcomp(bi, Cge) -> Kccall("greaterequal", 2)
+  | Pbigarrayref(n, _, _) -> Kccall("bigarray_get_" ^ string_of_int n, n + 1)
+  | Pbigarrayset(n, _, _) -> Kccall("bigarray_set_" ^ string_of_int n, n + 2)
   | _ -> fatal_error "Bytegen.comp_primitive"
 
 (* Compile an expression.

bytecomp/lambda.ml

@@ -75,6 +75,9 @@ type primitive =
   | Plsrbint of boxed_integer
   | Pasrbint of boxed_integer
   | Pbintcomp of boxed_integer * comparison
+  (* Operations on big arrays *)
+  | Pbigarrayref of int * bigarray_kind * bigarray_layout
+  | Pbigarrayset of int * bigarray_kind * bigarray_layout
 
 and comparison =
     Ceq | Cneq | Clt | Cgt | Cle | Cge
@@ -85,6 +88,19 @@ and array_kind =
 and boxed_integer =
     Pnativeint | Pint32 | Pint64
 
+and bigarray_kind =
+    Pbigarray_unknown
+  | Pbigarray_float32 | Pbigarray_float64
+  | Pbigarray_sint8 | Pbigarray_uint8
+  | Pbigarray_sint16 | Pbigarray_uint16
+  | Pbigarray_int32 | Pbigarray_int64 
+  | Pbigarray_caml_int | Pbigarray_native_int
+
+and bigarray_layout =
+    Pbigarray_unknown_layout
+  | Pbigarray_c_layout
+  | Pbigarray_fortran_layout
+
 type structured_constant =
     Const_base of constant
   | Const_pointer of int

bytecomp/lambda.mli

@@ -75,6 +75,9 @@ type primitive =
   | Plsrbint of boxed_integer
   | Pasrbint of boxed_integer
   | Pbintcomp of boxed_integer * comparison
+  (* Operations on big arrays *)
+  | Pbigarrayref of int * bigarray_kind * bigarray_layout
+  | Pbigarrayset of int * bigarray_kind * bigarray_layout
 
 and comparison =
     Ceq | Cneq | Clt | Cgt | Cle | Cge
@@ -85,6 +88,19 @@ and array_kind =
 and boxed_integer =
     Pnativeint | Pint32 | Pint64
 
+and bigarray_kind =
+    Pbigarray_unknown
+  | Pbigarray_float32 | Pbigarray_float64
+  | Pbigarray_sint8 | Pbigarray_uint8
+  | Pbigarray_sint16 | Pbigarray_uint16
+  | Pbigarray_int32 | Pbigarray_int64 
+  | Pbigarray_caml_int | Pbigarray_native_int
+
+and bigarray_layout =
+    Pbigarray_unknown_layout
+  | Pbigarray_c_layout
+  | Pbigarray_fortran_layout
+
 type structured_constant =
     Const_base of constant
   | Const_pointer of int

bytecomp/printlambda.ml

@@ -49,6 +49,26 @@ let print_boxed_integer name bi =
   | Pint32 -> printf "Int32.%s" name
   | Pint64 -> printf "Int64.%s" name
 
+let print_bigarray name kind layout =
+  printf "Bigarray.%s[%s,%s]"
+         name
+         (match kind with
+            Pbigarray_unknown -> "generic"
+          | Pbigarray_float32 -> "float32"
+          | Pbigarray_float64 -> "float64"
+          | Pbigarray_sint8 -> "sint8"
+          | Pbigarray_uint8 -> "uint8"
+          | Pbigarray_sint16 -> "sint16"
+          | Pbigarray_uint16 -> "uint16"
+          | Pbigarray_int32 -> "int32"
+          | Pbigarray_int64 -> "int64"
+          | Pbigarray_caml_int -> "camlint"
+          | Pbigarray_native_int -> "nativeint")
+         (match layout with
+            Pbigarray_unknown_layout -> "unknown"
+          | Pbigarray_c_layout -> "C"
+          | Pbigarray_fortran_layout -> "Fortran")
+
 let primitive ppf = function
     Pidentity -> print_string "id"
   | Pignore -> print_string "ignore"
@@ -134,6 +154,8 @@ let primitive ppf = function
   | Pbintcomp(bi, Cgt) -> print_boxed_integer ">" bi
   | Pbintcomp(bi, Cle) -> print_boxed_integer "<=" bi
   | Pbintcomp(bi, Cge) -> print_boxed_integer ">=" bi
+  | Pbigarrayref(n, kind, layout) -> print_bigarray "get" kind layout
+  | Pbigarrayset(n, kind, layout) -> print_bigarray "set" kind layout
 
 let rec lam ppf = function
     Lvar id ->

bytecomp/translcore.ml

@@ -205,7 +205,13 @@ let primitives_table = create_hashtable 57 [
   "%int64_xor", Pxorbint Pint64;
   "%int64_lsl", Plslbint Pint64;
   "%int64_lsr", Plsrbint Pint64;
-  "%int64_asr", Pasrbint Pint64
+  "%int64_asr", Pasrbint Pint64;
+  "%bigarray_ref_1", Pbigarrayref(1, Pbigarray_unknown, Pbigarray_c_layout);
+  "%bigarray_ref_2", Pbigarrayref(2, Pbigarray_unknown, Pbigarray_c_layout);
+  "%bigarray_ref_3", Pbigarrayref(3, Pbigarray_unknown, Pbigarray_c_layout);
+  "%bigarray_set_1", Pbigarrayset(1, Pbigarray_unknown, Pbigarray_c_layout);
+  "%bigarray_set_2", Pbigarrayset(2, Pbigarray_unknown, Pbigarray_c_layout);
+  "%bigarray_set_3", Pbigarrayset(3, Pbigarray_unknown, Pbigarray_c_layout)
 ]
 
 let prim_makearray =
@@ -253,6 +259,12 @@ let transl_prim prim args =
       | (Parraysetu Pgenarray, arg1 :: _) -> Parraysetu(array_kind arg1)
       | (Parrayrefs Pgenarray, arg1 :: _) -> Parrayrefs(array_kind arg1)
       | (Parraysets Pgenarray, arg1 :: _) -> Parraysets(array_kind arg1)
+      | (Pbigarrayref(n, Pbigarray_unknown, _), arg1 :: _) ->
+            let (k, l) = bigarray_kind_and_layout arg1 in
+            Pbigarrayref(n, k, l)
+      | (Pbigarrayset(n, Pbigarray_unknown, _), arg1 :: _) ->
+            let (k, l) = bigarray_kind_and_layout arg1 in
+            Pbigarrayset(n, k, l)
       | _ -> p
     end
   with Not_found ->

bytecomp/typeopt.ml

@@ -93,3 +93,35 @@ let array_kind exp = array_kind_gen exp.exp_type exp.exp_env
 
 let array_pattern_kind pat = array_kind_gen pat.pat_type pat.pat_env
 
+let bigarray_decode_type ty tbl dfl =
+  match (Ctype.repr ty).desc with
+    Tconstr(Pdot(Pident mod_id, type_name, _), [], _)
+    when Ident.name mod_id = "Bigarray" ->
+      begin try List.assoc type_name tbl with Not_found -> dfl end
+  | _ ->
+      dfl
+
+let kind_table =
+  ["float32_elt", Pbigarray_float32;
+   "float64_elt", Pbigarray_float64;
+   "int8_signed_elt", Pbigarray_sint8;
+   "int8_unsigned_elt", Pbigarray_uint8;
+   "int16_signed_elt", Pbigarray_sint16;
+   "int16_unsigned_elt", Pbigarray_uint16;
+   "int32_elt", Pbigarray_int32;
+   "int64_elt", Pbigarray_int64;
+   "int_elt", Pbigarray_caml_int;
+   "nativeint_elt", Pbigarray_native_int]
+
+let layout_table =
+  ["c_layout", Pbigarray_c_layout;
+   "fortran_layout", Pbigarray_fortran_layout]
+
+let bigarray_kind_and_layout exp =
+  let ty = Ctype.repr (Ctype.expand_head exp.exp_env exp.exp_type) in
+  match ty.desc with
+    Tconstr(p, [caml_type; elt_type; layout_type], abbrev) ->
+      (bigarray_decode_type elt_type kind_table Pbigarray_unknown,
+       bigarray_decode_type layout_type layout_table Pbigarray_unknown_layout)
+  | _ ->
+      (Pbigarray_unknown, Pbigarray_unknown_layout)

bytecomp/typeopt.mli

@@ -18,3 +18,5 @@ val has_base_type : Typedtree.expression -> Path.t -> bool
 val maybe_pointer : Typedtree.expression -> bool
 val array_kind : Typedtree.expression -> Lambda.array_kind
 val array_pattern_kind : Typedtree.pattern -> Lambda.array_kind
+val bigarray_kind_and_layout :
+      Typedtree.expression -> Lambda.bigarray_kind * Lambda.bigarray_layout