B
    dG                 @   s*  d Z ddlZddlZddlZddlmZ ddlm	Z	m
Z
mZmZ ddlmZmZmZmZmZ ejdkr|ddlmZmZmZ nHyddlmZmZmZ W n. ek
r   e jZd	d
 Zdd ZY nX dd Zdd Z G dd deZe Zdd Z!G dd dZ"e" Z#dd Z$dd Z%dd Z&dd Z'dd  Z(d!d" Z)d#d$ Z*d%d& Z+d'd( Z,dQd)d*Z-d+d, Z.d-d. Z/dRd/d0Z0d1d2 Z1dSd3d4Z2d5d6 Z3d7d8 Z4d9d: Z5d;d< Z6d=d> Z7d?d@ Z8dAdB Z9dCdD Z:dEdF Z;G dGdH dHe<Z=G dIdJ dJe<Z>G dKdL dLej?Z?G dMdN dNe?ej@Z@G dOdP dPe?ejAZAdS )Tz,Pickle-based communication using protocol 5.    N   )MPI)	PROC_NULL
ANY_SOURCEANY_TAGStatus)	_typedict
_comm_lock_commctx_intermemoryPickle)      )dumpsloadsHIGHEST_PROTOCOLc             O   s   t j| S )N)r   pickler   )obj_p_kw r   ]/work/yifan.wang/ringdown/master-ringdown-env/lib/python3.7/site-packages/mpi4py/util/pkl5.py_dumps)   s    r   c             O   s   t j| S )N)r   r   r   )bufr   r   r   r   r   _loads,   s    r   c                sP   g }| d ks| dk rt } | dk r(|d fS tt|j  fdd}||fS )Nr      c                s$   | } | kr  |  dS dS )NFTr   )r   )
buffer_add
buffer_len
buffer_raw	thresholdr   r   buf_cb9   s
    z_buffer_handler.<locals>.buf_cb)	_PROTOCOLlen_memoryappend)protocolr   bufsr    r   )r   r   r   r   r   _buffer_handler0   s    r'   c             C   s   d}t tj|| S )NZMPI4PY_PICKLE_THRESHOLD)int_osenvironget)defaultvarnamer   r   r   _get_thresholdB   s    r.   c                   sH   e Zd ZdZedZeeef fdd	Z	 fddZ
 fddZ  ZS )	r   z9Pickle/unpickle Python objects using out-of-band buffers.i   c                s   t  ||| dS )zInitialize pickle context.N)super__init__)selfr   r   r%   )	__class__r   r   r0   L   s    zPickle.__init__c                s(   t | j| j\}}t ||}||fS )z1Serialize object to data and out-of-band buffers.)r'   PROTOCOL	THRESHOLDr/   r   )r1   r   r&   r    data)r2   r   r   r   Q   s    zPickle.dumpsc                s   t  ||S )z5Deserialize object from data and out-of-band buffers.)r/   r   )r1   r5   r&   )r2   r   r   r   W   s    zPickle.loads)__name__
__module____qualname____doc__r.   r4   r   r   r!   r0   r   r   __classcell__r   r   )r2   r   r   G   s
   r   c       	      C   s^   t ||\}}| |||}| |}| d }tjdd|f||f}|  |  |S )N   )r;   r;   r   )divmodZCreate_vectorZCreate_contiguousZ
Get_extentr   ZDatatypeZCreate_structFree)	Zbasetypecount	blocksizeqsizeZrsizeZqtypeZrtypeZrdispZbigtyper   r   r   _bigmpi_create_type`   s    

rA   c               @   s4   e Zd ZdZdZdd Zdd Zdd Zd	d
 ZdS )_BigMPIz!Support for large message counts.i   @c             C   s
   i | _ d S )N)cache)r1   r   r   r   r0   r   s    z_BigMPI.__init__c             C   s   | S )Nr   )r1   r   r   r   	__enter__u   s    z_BigMPI.__enter__c             G   s,   | j }x| D ]}|  qW |  d S )N)rC   valuesr=   clear)r1   excrC   dtyper   r   r   __exit__x   s    z_BigMPI.__exit__c             C   sp   t |}t|}| j}||k r*||tjfS | j}||}|d k	rL|d|fS ttj||}| ||< |d|fS )Nr;   )	r#   r"   r?   r   BYTErC   r+   rA   ZCommit)r1   r   r>   r?   rC   rH   r   r   r   __call__~   s    

z_BigMPI.__call__N)	r6   r7   r8   r9   r?   r0   rD   rI   rK   r   r   r   r   rB   m   s   rB   c               C   s   dS )Nqr   r   r   r   r   _info_typecode   s    rM   c              C   s   t  } t|  S )N)rM   r   )coder   r   r   _info_datatype   s    rO   c             C   s*   t  }t| }d||}tj|f|  S )Nz{0}{1})rM   r"   format_structpack)inforN   sizesfmtr   r   r   
_info_pack   s    rV   c             C   s   t  }t|}t| | S )N)rM   rQ   calcsize	bytearray)rT   rN   itemsizer   r   r   _info_alloc   s    
rZ   c             C   s4   t  }t|}t| | }d||}t|| S )Nz{0}{1})rM   rQ   rW   r"   rP   unpack)rS   rN   rY   rT   rU   r   r   r   _info_unpack   s
    
r\   c             C   s   t j| S )N)r   r   allocate)rT   r   r   r   _new_buffer   s    r^   c       
   	   C   s   t |g}|dd |D  t }t|}|| ||f|| t8}|| |||| x|D ]}	|| ||	|| q\W W d Q R X d S )Nc             s   s   | ]}t t|V  qd S )N)r"   r#   ).0sbufr   r   r   	<genexpr>   s    z_send_raw.<locals>.<genexpr>)r"   extendrO   rV   _bigmpi)
commsendr5   r&   desttagrS   infotypebigmpir`   r   r   r   	_send_raw   s    

rj   c          	   C   s\   |t kr"|| d dtjf|| d S t|\}}t| d t| ||||| W d Q R X d S )Nr   re   )r   r   rJ   r   r   r	   rj   )rd   re   r   rf   rg   r5   r&   r   r   r   _send   s    rk   c                s.   g  fdd}t | |||| t}|S )Nc                s     | ||| d S )N)r$   )rd   r   rf   rg   )isendsreqsr   r   re      s    z_isend.<locals>.send)rk   Request)rd   rl   r   rf   rg   re   requestr   )rl   rm   r   _isend   s
    rp   c          	   C   s8  |d krt  }tj| ||| | }| }t }||}t|}tj	| ||f||| t
|}|d k	rt|}t||d kr|d |d  }t||d k rd }|d krt|d n|}	ttt|dd  }
t8}|| ||	|| x|
D ]}|| |||| qW W d Q R X |tjt| |	|
fS )Nr   r;   )r   r   CommZProbe
Get_sourceGet_tagrO   Get_elementsrZ   Recvr\   r#   Z
frombufferr"   r^   listmaprc   Set_elementsrJ   sum)rd   recvr   sourcerg   statusrh   infosizerS   r5   r&   ri   rbufr   r   r   	_recv_raw   s0    


 r   c          	   C   s\   |t kr$|| d dtjf||| d S t| d t| |||||\}}W d Q R X t||S )Nr   rz   )r   r   rJ   r	   r   r   r   )rd   rz   r   r{   rg   r|   r5   r&   r   r   r   _recv   s     r   c          	   C   s   |t kr$tj| |||}t|gS |d kr2t }t| d g }d}|| |||}|d kr`d S || | }|	 }t
 }||}	x<t|	D ]0}
tj| |||}|| ||tj7 }qW |tj| t|S Q R X d S )Nrz   r   )r   r   rq   MprobeMessager   r	   r$   rr   rs   rO   rt   rangerJ   rx   )rd   mprober{   rg   r|   rmsgmessagenumbytesrh   r}   _r   r   r   _mprobe   s,    



r   c             C   s2   t jj}t }t|}|| ||f| t|}|S )N)r   r   ru   rO   rZ   r\   )r   rT   r|   mrecvrh   rS   r   r   r   _mrecv_info  s    r   c             C   s>   t | d| tj}||d dtjf td \}}t||fS )Nr   )r   r   MESSAGE_NO_PROCrJ   r   r   rX   )r   r   r|   Znoprocr5   r&   r   r   r   _mrecv_none  s
    r   c       
   	   C   s   | d t jkr"| d }t|||S t| }t| d }tt|||}t|d }tt	t|dd  }t
8}|t||| x|D ]}	|t|||	 qW W d Q R X |d k	r|t jt| ||fS )Nr   r;   )r   r   r   iterr"   r   nextr^   rv   rw   rc   rx   rJ   ry   )
r   r   r|   r   ZicntrS   r5   r&   ri   r~   r   r   r   _mrecv_data!  s    
 r   c             C   s$   dd }t | ||\}}t||S )Nc             S   s   t j| | d S )N)r   r   ru   )r   r   r   r   r   r   4  s    z_mrecv.<locals>.mrecv)r   r   r   )r   r|   r   r5   r&   r   r   r   _mrecv3  s    r   c                s:   g   fdd}t | |\}}t }t|d||f |S )Nc                s     tj| | d S )N)r$   r   r   ZIrecv)r   r   )rreqsr   r   r   <  s    z_imrecv.<locals>.mrecv
_data_bufs)r   rn   setattr)r   r   r5   r&   ro   r   )r   r   _imrecv:  s    r   c             C   sL   t | dd }| tjkr(|d k	r(t| d |d k	rH|\}}t||}|S d S )Nr   )getattrr   REQUEST_NULLdelattrr   r   )ro   Z	data_bufsr5   r&   r   r   r   r   	_req_loadD  s    
r   c             C   s8   |d krd n|g}|| |}|r0t | }||fS |d fS )N)r   )ro   testr|   statusesflagr   r   r   r   _testO  s    
r   c       
      C   s   t |tr4x(tt| t| D ]}|t  q W g }d }x| D ]}|| qBW |d k	rg }x*t| |D ]\}}||gt|  qlW |||}|rdd | D }	||	fS |d fS )Nc             S   s   g | ]}t |qS r   )r   )r_   reqr   r   r   
<listcomp>f  s    z_testall.<locals>.<listcomp>)
isinstancerv   r   r"   r$   r   rb   zip)
requeststestallr   r   ZreqarrayZstsarrayr   stsr   Zobjsr   r   r   _testallX  s     


r   c          	   C   s  |   }||krlt|g}|dd |D  t }tt|g}|| ||f| t|}|| ||f| nlt }td}|| ||f| t|d }t|}|| ||f| t|}t|d }tt	t|dd  }t
4}	|| |	|| x|D ]}
|| |	|
| qW W d Q R X ||fS )Nc             s   s   | ]}t t|V  qd S )N)r"   r#   )r_   r`   r   r   r   ra   o  s    z#_bcast_intra_raw.<locals>.<genexpr>r;   r   )Get_rankr"   rb   rO   rV   rZ   r\   r^   rv   rw   rc   )rd   bcastr5   r&   rootrankrS   rh   r}   ri   r~   r   r   r   _bcast_intra_rawk  s.    

r   c          	   C   sd   |   }||kr t|\}}ntd \}}t| d t| ||||\}}W d Q R X t||S )Nr   )r   r   r   r	   r   r   )rd   r   r   r   r   r5   r&   r   r   r   _bcast_intra  s    r   c          	   C   s  |   }|  }t| \} }}}|tjkr.d S |tjkrdtjj}	t	|\}
}t
| |	|
|d| d S d|  krx|k rn nj|dkrtjj}t| |d ||\}
}nt	d \}
}t|d t|||
|d\}
}W d Q R X t|
|S | tj ttjd S )Nr   r   )r   ZGet_remote_sizer
   r   r   ROOTrq   Sendr   r   rj   ru   r   r	   r   r   ZCall_errhandlerZERR_ROOT	Exception)rd   r   r   r   r   rT   rg   Z	localcommr   re   r5   r&   rz   r   r   r   _bcast_inter  s(    

r   c             C   s(   |   rt| |||S t| |||S d S )N)ZIs_interr   r   )rd   r   r   r   r   r   r   _bcast  s    r   c                   s   e Zd ZdZd fdd	Zdd Zdd Zd	d
 ZddddZdd Z	dddZ
dddZdddZedddZedddZ  ZS ) rn   zRequest.Nc                s0   |dkrt jf}t|t jr"|f}t | |S )zCreate and return a new object.N)r   r   r   rn   r/   __new__)clsro   )r2   r   r   r     s
    zRequest.__new__c                s@   t  trt| t kS t  tjr<t fdd| D S tS )zReturn ``self==other``.c             3   s   | ]}| kV  qd S )Nr   )r_   r   )otherr   r   ra     s    z!Request.__eq__.<locals>.<genexpr>)r   rn   tupler   allNotImplemented)r1   r   r   )r   r   __eq__  s
    
zRequest.__eq__c                s@   t  trt| t kS t  tjr<t fdd| D S tS )zReturn ``self!=other``.c             3   s   | ]}| kV  qd S )Nr   )r_   r   )r   r   r   ra     s    z!Request.__ne__.<locals>.<genexpr>)r   rn   r   r   anyr   )r1   r   r   )r   r   __ne__  s
    
zRequest.__ne__c             C   s   t dd | D S )zReturn ``bool(self)``.c             s   s   | ]
}|V  qd S )Nr   )r_   r   r   r   r   ra     s    z#Request.__bool__.<locals>.<genexpr>)r   )r1   r   r   r   __bool__  s    zRequest.__bool__)returnc             C   s   x| D ]}|   qW dS )zFree a communication request.N)r=   )r1   r   r   r   r   r=     s    
zRequest.Freec             C   s   x| D ]}|   qW dS )zCancel a communication request.N)ZCancel)r1   r   r   r   r   cancel  s    
zRequest.cancelc             C   s2   |gdgt t| d d  }tttjj| |S )z5Non-destructive test for the completion of a request.Nr;   r   )maxr"   r   rw   r   rn   Z
Get_status)r1   r|   r   r   r   r   
get_status  s    zRequest.get_statusc             C   s   t | tjj|S )z%Test for the completion of a request.)r   r   rn   Testall)r1   r|   r   r   r   r     s    zRequest.testc             C   s   t | tjj|d S )zWait for a request to complete.r;   )r   r   rn   Waitall)r1   r|   r   r   r   wait  s    zRequest.waitc             C   s   t |tjj|S )z(Test for the completion of all requests.)r   r   rn   r   )r   r   r   r   r   r   r     s    zRequest.testallc             C   s   t |tjj|d S )z"Wait for all requests to complete.r;   )r   r   rn   r   )r   r   r   r   r   r   waitall  s    zRequest.waitall)N)N)N)N)N)N)r6   r7   r8   r9   r   r   r   r   r=   r   r   r   r   classmethodr   r   r:   r   r   )r2   r   rn     s   


rn   c                   st   e Zd ZdZd fdd	Zdd Zdd Zd	d
 ZdddZdd Z	e
eedfddZe
eedfddZ  ZS )r   zMessage.Nc                s0   |dkrt jf}t|t jr"|f}t | |S )zCreate and return a new object.N)r   ZMESSAGE_NULLr   r   r/   r   )r   r   )r2   r   r   r     s
    zMessage.__new__c                s@   t  trt| t kS t  tjr<t fdd| D S tS )zReturn ``self==other``.c             3   s   | ]}| kV  qd S )Nr   )r_   msg)r   r   r   ra     s    z!Message.__eq__.<locals>.<genexpr>)r   r   r   r   r   r   )r1   r   r   )r   r   r     s
    
zMessage.__eq__c                s@   t  trt| t kS t  tjr<t fdd| D S tS )zReturn ``self!=other``.c             3   s   | ]}| kV  qd S )Nr   )r_   r   )r   r   r   ra     s    z!Message.__ne__.<locals>.<genexpr>)r   r   r   r   r   r   )r1   r   r   )r   r   r     s
    
zMessage.__ne__c             C   s   t dd | D S )zReturn ``bool(self)``.c             s   s   | ]
}|V  qd S )Nr   )r_   r   r   r   r   ra     s    z#Message.__bool__.<locals>.<genexpr>)r   )r1   r   r   r   r     s    zMessage.__bool__c             C   s
   t | |S )z$Blocking receive of matched message.)r   )r1   r|   r   r   r   rz     s    zMessage.recvc             C   s   t | S )z'Nonblocking receive of matched message.)r   )r1   r   r   r   irecv  s    zMessage.irecvc             C   s   t |tjj|||S )z$Blocking test for a matched message.)r   r   rq   r   )r   rd   r{   rg   r|   r   r   r   probe  s    zMessage.probec             C   s   t |tjj|||S )z'Nonblocking test for a matched message.)r   r   rq   Improbe)r   rd   r{   rg   r|   r   r   r   iprobe"  s    zMessage.iprobe)N)N)r6   r7   r8   r9   r   r   r   r   rz   r   r   r   r   r   r   r:   r   r   )r2   r   r     s   
r   c               @   s   e Zd ZdZdddZdddZdddZdd	d
Zd ddZd!ddZ	de
edfddZde
efddZdde
edfddZe
edfddZe
edfddZd"ddZdS )#rq   zCommunicator.r   c             C   s   t | tjj||| dS )zBlocking send in standard mode.N)rk   r   rq   r   )r1   r   rf   rg   r   r   r   re   -  s    z	Comm.sendc             C   s   t | tjj||| dS )zBlocking send in buffered mode.N)rk   r   rq   ZBsend)r1   r   rf   rg   r   r   r   bsend1  s    z
Comm.bsendc             C   s$   t | tjj|||}tj| dS )z"Blocking send in synchronous mode.N)rp   r   rq   Issendrn   r   )r1   r   rf   rg   sreqr   r   r   ssend5  s    z
Comm.ssendc             C   s   t | tjj|||S )z"Nonblocking send in standard mode.)rp   r   rq   Isend)r1   r   rf   rg   r   r   r   rl   :  s    z
Comm.isendc             C   s   t | tjj|||S )z"Nonblocking send in buffered mode.)rp   r   rq   ZIbsend)r1   r   rf   rg   r   r   r   ibsend>  s    zComm.ibsendc             C   s   t | tjj|||S )z%Nonblocking send in synchronous mode.)rp   r   rq   r   )r1   r   rf   rg   r   r   r   issendB  s    zComm.issendNc             C   s   t | tjj||||S )zBlocking receive.)r   r   rq   ru   )r1   r   r{   rg   r|   r   r   r   rz   F  s    z	Comm.recvc             C   s   t ddS )zNonblocking receive.unsupportedN)RuntimeError)r1   r   r{   rg   r   r   r   r   L  s    z
Comm.irecvc       
      C   s:   t | tjj|||}t| tjj||||}	tj| |	S )zSend and receive.)rp   r   rq   r   r   ru   rn   r   )
r1   Zsendobjrf   ZsendtagZrecvbufr{   Zrecvtagr|   r   Zrobjr   r   r   sendrecvQ  s    zComm.sendrecvc             C   s   t | tjj|||S )z$Blocking test for a matched message.)r   r   rq   r   )r1   r{   rg   r|   r   r   r   r   \  s    zComm.mprobec             C   s   t | tjj|||S )z'Nonblocking test for a matched message.)r   r   rq   r   )r1   r{   rg   r|   r   r   r   improbeb  s    zComm.improbec             C   s   t | tjj||S )z
Broadcast.)r   r   rq   ZBcast)r1   r   r   r   r   r   r   h  s    z
Comm.bcast)r   )r   )r   )r   )r   )r   )r   )r6   r7   r8   r9   re   r   r   rl   r   r   r   r   rz   r   r   r   r   r   r   r   r   r   rq   *  s$   





	rq   c               @   s   e Zd ZdZdS )	IntracommzIntracommunicator.N)r6   r7   r8   r9   r   r   r   r   r   m  s   r   c               @   s   e Zd ZdZdS )	IntercommzIntercommunicator.N)r6   r7   r8   r9   r   r   r   r   r   q  s   r   )N)N)N)Br9   osr)   sys_sysstructrQ    r   r   r   r   r   r   r	   r
   r   r#   r   Z_Pickleversion_infor   r   r   r   r   r   r!   Zpickle5ImportErrorr3   r'   r.   rA   rB   rc   rM   rO   rV   rZ   r\   r^   rj   rk   rp   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rn   r   rq   r   r   r   r   r   r   <module>   sd   

 		


	

	D6C