Accelerated FourIdxStorage

This commit is contained in:
Anthony Scemama 2018-06-01 10:07:17 +02:00
parent 15390759f9
commit 63f0b379de
2 changed files with 74 additions and 33 deletions

View File

@ -57,7 +57,12 @@ let make ?guess:(guess=`Huckel) ?max_scf:(max_scf=64) ?level_shift:(level_shift=
in
x.Fock.fock, x.Fock.core, x.Fock.coulomb, x.Fock.exchange
in
(*
debug_matrix "Fock" m_F;
debug_matrix "Coulomb" m_J;
debug_matrix "Exchange" m_K;
debug_matrix "HCore" m_Hc;
*)
(* Add level shift in AO basis *)
let m_F =
let m_SC =

View File

@ -4,7 +4,7 @@ type index_pair = { first : int ; second : int }
type storage_t =
| Dense of (float, Bigarray.float32_elt, Bigarray.fortran_layout) Bigarray.Genarray.t
| Dense of (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t
| Sparse of (int, float) Hashtbl.t
type t =
@ -24,49 +24,85 @@ let key_of_indices ~r1 ~r2 =
f p q
let dense_index i j size =
(j-1)*size + i
let get_four_index ~r1 ~r2 t =
match t.four_index with
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
Bigarray.Genarray.get t [| i; j; k; l|]
| Sparse t -> let key = key_of_indices ~r1 ~r2 in
try Hashtbl.find t key
| Dense a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
let size = t.size in
assert ( (i lor j lor k lor l) > 0 );
assert ( i <= size && j <= size && k <= size && l <= size );
Bigarray.Array2.unsafe_get a (dense_index i j size) (dense_index k l size)
)
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
try Hashtbl.find a key
with Not_found -> 0.
let set_four_index ~r1 ~r2 ~value t =
match t.four_index with
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
Bigarray.Genarray.set t [| i; j; k; l|] value;
Bigarray.Genarray.set t [| k; j; i; l|] value;
Bigarray.Genarray.set t [| i; l; k; j|] value;
Bigarray.Genarray.set t [| k; l; i; j|] value;
Bigarray.Genarray.set t [| j; i; l; k|] value;
Bigarray.Genarray.set t [| j; k; l; i|] value;
Bigarray.Genarray.set t [| l; i; j; k|] value;
Bigarray.Genarray.set t [| l; k; j; i|] value;
| Dense a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
let size = t.size in
assert ( (i lor j lor k lor l) > 0 );
assert ( i <= size && j <= size && k <= size && l <= size);
let ij = (dense_index i j size)
and kl = (dense_index k l size)
and il = (dense_index i l size)
and kj = (dense_index k j size)
and ji = (dense_index j i size)
and lk = (dense_index l k size)
and li = (dense_index l i size)
and jk = (dense_index j k size)
in
let open Bigarray.Array2 in
unsafe_set a ij kl value;
unsafe_set a kj il value;
unsafe_set a il kj value;
unsafe_set a kl ij value;
unsafe_set a ji lk value;
unsafe_set a li jk value;
unsafe_set a jk li value;
unsafe_set a lk ji value
)
| Sparse t -> let key = key_of_indices ~r1 ~r2 in
Hashtbl.replace t key value
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
Hashtbl.replace a key value
let increment_four_index ~r1 ~r2 ~value x =
match x.four_index with
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
t.{i,j,k,l} <- t.{i,j,k,l} +. value;
t.{k,j,i,l} <- t.{k,j,i,l} +. value;
t.{i,l,k,j} <- t.{i,l,k,j} +. value;
t.{k,l,i,j} <- t.{k,l,i,j} +. value;
t.{j,i,l,k} <- t.{j,i,l,k} +. value;
t.{j,k,l,i} <- t.{j,k,l,i} +. value;
t.{l,i,j,k} <- t.{l,i,j,k} +. value;
t.{l,k,j,i} <- t.{l,k,j,i} +. value
let increment_four_index ~r1 ~r2 ~value t =
match t.four_index with
| Dense a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
let size = t.size in
assert ( (i lor j lor k lor l) > 0 );
assert ( i <= size && j <= size && k <= size && l <= size);
let ij = (dense_index i j size)
and kl = (dense_index k l size)
and il = (dense_index i l size)
and kj = (dense_index k j size)
and ji = (dense_index j i size)
and lk = (dense_index l k size)
and li = (dense_index l i size)
and jk = (dense_index j k size)
in
let open Bigarray.Array2 in
unsafe_set a ij kl (value +. unsafe_get a ij kl) ;
unsafe_set a kj il (value +. unsafe_get a kj il) ;
unsafe_set a il kj (value +. unsafe_get a il kj) ;
unsafe_set a kl ij (value +. unsafe_get a kl ij) ;
unsafe_set a ji lk (value +. unsafe_get a ji lk) ;
unsafe_set a li jk (value +. unsafe_get a li jk) ;
unsafe_set a jk li (value +. unsafe_get a jk li) ;
unsafe_set a lk ji (value +. unsafe_get a lk ji)
)
| Sparse t -> let key = key_of_indices ~r1 ~r2 in
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
let old_value =
try Hashtbl.find t key
try Hashtbl.find a key
with Not_found -> 0.
in
Hashtbl.replace t key (old_value +. value)
Hashtbl.replace a key (old_value +. value)
let get ~r1 ~r2 =
get_four_index ~r1 ~r2
@ -83,9 +119,9 @@ let create ~size sparsity =
match sparsity with
| `Dense ->
let result =
Bigarray.Genarray.create Float32 Bigarray.fortran_layout [| size ; size ; size ; size |]
Bigarray.Array2.create Float64 Bigarray.fortran_layout (size*size) (size*size)
in
Bigarray.Genarray.fill result 0.;
Bigarray.Array2.fill result 0.;
Dense result
| `Sparse ->