quantum_package/plugins/dress_zmq/run_dress_slave.irp.f

use bitmasks

BEGIN_PROVIDER [ integer, fragment_count ]
  implicit none
  BEGIN_DOC
    ! Number of fragments for the deterministic part
  END_DOC
  fragment_count = 1
END_PROVIDER


subroutine run_dress_slave(thread,iproce,energy)
  use f77_zmq
  implicit none

  double precision, intent(in)    :: energy(N_states_diag)
  integer,  intent(in)            :: thread, iproce
  integer                        :: rc, i, subset, i_generator

  integer                        :: worker_id, task_id, ctask, ltask
  character*(5120)                :: task

  integer(ZMQ_PTR),external      :: new_zmq_to_qp_run_socket
  integer(ZMQ_PTR)               :: zmq_to_qp_run_socket

  integer(ZMQ_PTR), external     :: new_zmq_push_socket
  integer(ZMQ_PTR)               :: zmq_socket_push

  logical :: done

  double precision,allocatable :: dress_detail(:)
  integer :: ind
  
  double precision,allocatable :: delta_ij_loc(:,:,:)
  double precision :: div(N_states) 
  integer :: h,p,n,i_state
  logical :: ok

  integer, allocatable :: int_buf(:)
  double precision, allocatable :: double_buf(:)
  integer(bit_kind), allocatable :: det_buf(:,:,:)
  integer :: N_buf(3)
  logical :: last
  integer, external :: omp_get_thread_num 
  double precision, allocatable :: delta_det(:,:,:,:), cp(:,:,:,:)
  integer :: toothMwen
  logical :: fracted
  double precision :: fac
      


  if(iproce /= 0) stop "RUN DRESS SLAVE is OMP"
  
  allocate(delta_det(N_states, N_det, 0:comb_teeth+1, 2))
  allocate(cp(N_states, N_det, N_cp, 2))
  delta_det = 0d9
  cp = 0d0

  
  task(:) = CHAR(0)

  
  zmq_to_qp_run_socket = new_zmq_to_qp_run_socket()
  zmq_socket_push      = new_zmq_push_socket(thread)
  call connect_to_taskserver(zmq_to_qp_run_socket,worker_id,thread)
  if(worker_id == -1) then
    print *, "WORKER -1"
    call end_zmq_to_qp_run_socket(zmq_to_qp_run_socket)
    call end_zmq_push_socket(zmq_socket_push,thread)
    return
  end if
  do i=1,N_states
    div(i) = psi_coef(dressed_column_idx(i), i)
  end do
  
  integer :: iproc, cur_cp, done_for(0:N_cp)
  integer, allocatable :: tasks(:)
  logical :: loop, donedone
  integer :: res_task(Nproc), res_gen(Nproc), res_sub(Nproc)
  res_gen = 0

  donedone = .false.
  allocate(tasks(0:N_det))
  done_for = 0

  do cur_cp=0, N_cp
    if(donedone) exit
    print *, "DOING CP", cur_cp
    tasks(0) = 0
    !$OMP PARALLEL DEFAULT(SHARED) &
    !$OMP PRIVATE(int_buf, double_buf, det_buf, delta_ij_loc, task, task_id) &
    !$OMP PRIVATE(toothMwen, fracted, fac) &
    !$OMP PRIVATE(loop, i_generator, subset, iproc, N_buf)
    iproc = omp_get_thread_num()+1
    loop = .true.
    allocate(int_buf(N_dress_int_buffer)) 
    allocate(double_buf(N_dress_double_buffer)) 
    allocate(det_buf(N_int, 2, N_dress_det_buffer)) 
    allocate(delta_ij_loc(N_states,N_det,2)) 
    do while(loop)
      if(res_gen(iproc) == 0) then
        !$OMP CRITICAL
        call get_task_from_taskserver(zmq_to_qp_run_socket,worker_id, task_id, task)
        !$OMP END CRITICAL
        task = task//"  0"
        if(task_id == 0) then
          donedone = .true.
          print *, "DONEDONE"
          exit   !! LAST MESSAGE ???
        end if
        read (task,*) subset, i_generator
      else
        subset = res_sub(iproc)
        i_generator = res_gen(iproc)
        task_id = res_task(iproc)
        res_gen(iproc) = 0
      end if
      
      !if(done_cp_at_det(i_generator) > cur_cp) loop = .false.
      if(done_cp_at_det(i_generator) > cur_cp) then
        res_gen(iproc) = i_generator
        res_task(iproc) = task_id
        res_sub(iproc) = subset
        exit
      end if
      
      !$OMP ATOMIC
      done_for(done_cp_at_det(i_generator)) += 1

      delta_ij_loc(:,:,:) = 0d0
      call generator_start(i_generator, iproc)
      call alpha_callback(delta_ij_loc, i_generator, subset, iproc)
      call generator_done(i_generator, int_buf, double_buf, det_buf, N_buf, iproc)
     
      !if(.false.) then
      !$OMP CRITICAL
      do i=1,N_cp
        fac = cps(i_generator, i) * dress_weight_inv(i_generator) * comb_step
        if(fac == 0d0) cycle
        cp(:,:,i,1) += (delta_ij_loc(:,:,1) * fac)
        cp(:,:,i,2) += (delta_ij_loc(:,:,2) * fac)
      end do


      toothMwen = tooth_of_det(i_generator)
      fracted = (toothMwen /= 0)
      if(fracted) fracted = (i_generator == first_det_of_teeth(toothMwen))
      if(fracted) then
        delta_det(:,:,toothMwen-1, 1) += delta_ij_loc(:,:,1) * (1d0-fractage(toothMwen))
        delta_det(:,:,toothMwen-1, 2) += delta_ij_loc(:,:,2) * (1d0-fractage(toothMwen))
        delta_det(:,:,toothMwen  , 1) += delta_ij_loc(:,:,1) * (fractage(toothMwen))
        delta_det(:,:,toothMwen  , 2) += delta_ij_loc(:,:,2) * (fractage(toothMwen))
      else
        delta_det(:,:,toothMwen  , 1) += delta_ij_loc(:,:,1)
        delta_det(:,:,toothMwen  , 2) += delta_ij_loc(:,:,2)
      end if

      
      !$OMP END CRITICAL
      !end if

      !$OMP CRITICAL
      call push_dress_results(zmq_socket_push, i_generator, -1, delta_ij_loc, int_buf, double_buf, det_buf, N_buf, task_id)
      call task_done_to_taskserver(zmq_to_qp_run_socket,worker_id,task_id)
      !$OMP END CRITICAL
      tasks(0) += 1
      tasks(tasks(0)) = task_id
      
    end do
    print *, "SLAVE", iproc, "waits"
    deallocate(int_buf,double_buf,det_buf,delta_ij_loc)
    !$OMP END PARALLEL
    
    allocate(delta_ij_loc(N_states,N_det,2))
    allocate(int_buf(1), double_buf(1), det_buf(1,1,1))
    N_buf = (/0,1,0/)
    
    delta_ij_loc = 0d0
    
    if(cur_cp > 0) then
      do i=1,cur_cp
        delta_ij_loc(:,:,:) += cp(:,:,i,:)
        !delta_s2(:,:) += cp(:,:,i,2)
      end do

      delta_ij_loc(:,:,:) = delta_ij_loc(:,:,:) / cps_N(cur_cp)
      do i=cp_first_tooth(cur_cp)-1,0,-1
        delta_ij_loc(:,:,:) = delta_ij_loc(:,:,:) +delta_det(:,:,i,:)
      end do
    end if
    call sleep(1)
    call push_dress_results(zmq_socket_push, done_for(cur_cp), cur_cp, delta_ij_loc, int_buf, double_buf, det_buf, N_buf, -1)
    !do i=1,tasks(0)
    !  call task_done_to_taskserver(zmq_to_qp_run_socket,worker_id,tasks(i))
    !end do
    deallocate(delta_ij_loc, int_buf, double_buf, det_buf)
  end do
  
  call sleep(10) 
  call disconnect_from_taskserver(zmq_to_qp_run_socket,zmq_socket_push,worker_id)
  call end_zmq_to_qp_run_socket(zmq_to_qp_run_socket)
  call end_zmq_push_socket(zmq_socket_push,thread)
end subroutine



subroutine push_dress_results(zmq_socket_push, ind, cur_cp, delta_loc, int_buf, double_buf, det_buf, N_bufi, task_id)
  use f77_zmq
  implicit none

  integer(ZMQ_PTR), intent(in)   :: zmq_socket_push
  double precision, intent(inout)   :: delta_loc(N_states, N_det, 2)
  real(kind=4), allocatable   :: delta_loc4(:,:,:)
  double precision, intent(in) :: double_buf(*)
  integer, intent(in) :: int_buf(*)
  integer(bit_kind), intent(in) :: det_buf(N_int, 2, *)
  integer, intent(in) :: N_bufi(3)
  integer :: N_buf(3)
  integer, intent(in) :: ind, cur_cp, task_id
  integer :: rc, i, j, k, l
  double precision :: contrib(N_states)
  

  
  rc = f77_zmq_send( zmq_socket_push, ind, 4, ZMQ_SNDMORE)
  if(rc /= 4) stop "push"
  
  rc = f77_zmq_send( zmq_socket_push, cur_cp, 4, ZMQ_SNDMORE)
  if(rc /= 4) stop "push"
  

  if(cur_cp /= -1) then
    rc = f77_zmq_send( zmq_socket_push, delta_loc(1,1,1), 8*N_states*N_det, ZMQ_SNDMORE)
    if(rc /= 8*N_states*N_det) stop "push"

    rc = f77_zmq_send( zmq_socket_push, delta_loc(1,1,2), 8*N_states*N_det, ZMQ_SNDMORE)
    if(rc /= 8*N_states*N_det) stop "push"
  else
    contrib = 0d0
    
    do i=1,N_det
      contrib(:) += delta_loc(:,i, 1) * psi_coef(i, :)
    end do

    rc = f77_zmq_send( zmq_socket_push, contrib, 8*N_states, ZMQ_SNDMORE)
    if(rc /= 8*N_states) stop "push"
    
    N_buf = N_bufi
    N_buf = (/0,1,0/)

    rc = f77_zmq_send( zmq_socket_push, N_buf, 4*3, ZMQ_SNDMORE)
    if(rc /= 4*3) stop "push5" 
    
    if(N_buf(1) > N_dress_int_buffer) stop "run_dress_slave N_buf bad size?"
    if(N_buf(2) > N_dress_double_buffer) stop "run_dress_slave N_buf bad size?"
    if(N_buf(3) > N_dress_det_buffer) stop "run_dress_slave N_buf bad size?"

    
    if(N_buf(1) > 0) then
      rc = f77_zmq_send( zmq_socket_push, int_buf, 4*N_buf(1), ZMQ_SNDMORE)
      if(rc /= 4*N_buf(1)) stop "push6"
    end if
    
    if(N_buf(2) > 0) then
      rc = f77_zmq_send( zmq_socket_push, double_buf, 8*N_buf(2), ZMQ_SNDMORE)
      if(rc /= 8*N_buf(2)) stop "push8"
    end if

    if(N_buf(3) > 0) then
      rc = f77_zmq_send( zmq_socket_push, det_buf, 2*N_int*bit_kind*N_buf(3), ZMQ_SNDMORE)
      if(rc /= 2*N_int*bit_kind*N_buf(3)) stop "push10"
    end if

    rc = f77_zmq_send( zmq_socket_push, task_id, 4, 0)
    if(rc /= 4) stop "push11"
  end if

! Activate is zmq_socket_push is a REQ
IRP_IF ZMQ_PUSH
IRP_ELSE
  character*(2) :: ok
  rc = f77_zmq_recv( zmq_socket_push, ok, 2, 0)
IRP_ENDIF

end subroutine


subroutine pull_dress_results(zmq_socket_pull, ind, cur_cp, delta_loc, int_buf, double_buf, det_buf, N_buf, task_id, contrib)
  use f77_zmq
  implicit none
  integer(ZMQ_PTR), intent(in)   :: zmq_socket_pull
  integer, intent(out) :: cur_cp
  double precision, intent(inout) :: delta_loc(N_states, N_det, 2)
  double precision, intent(out) :: double_buf(*), contrib(N_states)
  integer, intent(out) :: int_buf(*)
  integer(bit_kind), intent(out) :: det_buf(N_int, 2, *)
  integer, intent(out) :: ind
  integer, intent(out) :: task_id
  integer :: rc, i, j, k
  integer, intent(out) :: N_buf(3)


  
  rc = f77_zmq_recv( zmq_socket_pull, ind, 4, 0)
  if(rc /= 4) stop "pulla"
 
  rc = f77_zmq_recv( zmq_socket_pull, cur_cp, 4, 0)
  if(rc /= 4) stop "pulla"
  
  
  

  if(cur_cp /= -1) then
    rc = f77_zmq_recv( zmq_socket_pull, delta_loc(1,1,1), N_states*8*N_det, 0)
    if(rc /= 8*N_states*N_det) stop "pullc"
  
    rc = f77_zmq_recv( zmq_socket_pull, delta_loc(1,1,2), N_states*8*N_det, 0)
    if(rc /= 8*N_states*N_det) stop "pulld"
  else
    rc = f77_zmq_recv( zmq_socket_pull, contrib, 8*N_states, 0)
    if(rc /= 8*N_states) stop "pullc"
  
    rc = f77_zmq_recv( zmq_socket_pull, N_buf, 4*3, 0)
    if(rc /= 4*3) stop "pull" 
    if(N_buf(1) > N_dress_int_buffer) stop "run_dress_slave N_buf bad size?"
    if(N_buf(2) > N_dress_double_buffer) stop "run_dress_slave N_buf bad size?"
    if(N_buf(3) > N_dress_det_buffer) stop "run_dress_slave N_buf bad size?"

    
    if(N_buf(1) > 0) then
      rc = f77_zmq_recv( zmq_socket_pull, int_buf, 4*N_buf(1), 0)
      if(rc /= 4*N_buf(1)) stop "pull1"
    end if
    
    if(N_buf(2) > 0) then
      rc = f77_zmq_recv( zmq_socket_pull, double_buf, 8*N_buf(2), 0)
      if(rc /= 8*N_buf(2)) stop "pull2"
    end if
    
    if(N_buf(3) > 0) then
      rc = f77_zmq_recv( zmq_socket_pull, det_buf, 2*N_int*bit_kind*N_buf(3), 0)
      if(rc /= 2*N_int*bit_kind*N_buf(3)) stop "pull3"
    end if

    rc = f77_zmq_recv( zmq_socket_pull, task_id, 4, 0)
    if(rc /= 4) stop "pull4"
  end if
! Activate is zmq_socket_pull is a REP
IRP_IF ZMQ_PUSH
IRP_ELSE
  rc = f77_zmq_send( zmq_socket_pull, 'ok', 2, 0)
IRP_ENDIF

end subroutine