From 63f0b379de9ce0d1443efd6da06df53584d169c3 Mon Sep 17 00:00:00 2001 From: Anthony Scemama Date: Fri, 1 Jun 2018 10:07:17 +0200 Subject: [PATCH] Accelerated FourIdxStorage --- SCF/RHF.ml | 7 ++- Utils/FourIdxStorage.ml | 100 +++++++++++++++++++++++++++------------- 2 files changed, 74 insertions(+), 33 deletions(-) diff --git a/SCF/RHF.ml b/SCF/RHF.ml index e121805..bb492de 100644 --- a/SCF/RHF.ml +++ b/SCF/RHF.ml @@ -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 = diff --git a/Utils/FourIdxStorage.ml b/Utils/FourIdxStorage.ml index 97e7394..2c8b57e 100644 --- a/Utils/FourIdxStorage.ml +++ b/Utils/FourIdxStorage.ml @@ -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 ->