QCaml/Parallel/Parallel.ml

(** Module for handling distributed parallelism *)

let size =
  let result = 
    Mpi.comm_size Mpi.comm_world
  in
  assert (result >  0);
  result


let rank = 
  let result = 
    Mpi.comm_rank Mpi.comm_world
  in
  assert (result >= 0);
  result


let barrier () =
  Mpi.barrier Mpi.comm_world


module Vec = struct

  type t =
  {
    global_first : int ; (* Lower  index in the global array *)
    global_last  : int ; (* Higher index in the global array *)
    local_first  : int ; (* Lower  index in the local array *)
    local_last   : int ; (* Higher index in the local array *)
    data         : Lacaml.D.vec ; (* Lacaml vector containing the data *)
  }

  let pp ppf v =
    Format.fprintf ppf "@[<2>";
    Format.fprintf ppf "@[ gf : %d@]@;" v.global_first;
    Format.fprintf ppf "@[ gl : %d@]@;" v.global_last;
    Format.fprintf ppf "@[ lf : %d@]@;" v.local_first;
    Format.fprintf ppf "@[ ll : %d@]@;" v.local_last;
    Format.fprintf ppf "@[ data : %a@]@;" (Lacaml.Io.pp_lfvec ()) v.data;
    Format.fprintf ppf "@]@."; 
    ()

  let create n = 
    let step = n / size + 1 in
    let local_first = step * rank + 1 in
    let local_last = min (local_first + step - 1) n  in
    {
      global_first = 1 ;
      global_last  = n ;
      local_first  ;
      local_last   ;
      data         = Lacaml.D.Vec.create (local_last - local_first + 1)
    }


  let make n x = 
    let result = create n in
    { result with data =
      Lacaml.D.Vec.make
        (Lacaml.D.Vec.dim result.data)
        x
    }


  let make0 n = 
    make n 0.

  let init n f =
    let result = create n in
    { result with data =
        Lacaml.D.Vec.init
          (Lacaml.D.Vec.dim result.data)
          (fun i -> f (i+result.local_first-1))
    }
end