
    J/Ph                          d dl Zd dlmZmZmZmZ d dlmZ d dl	m
Z
 d dlmZmZmZ ddlmZmZ d Zd	 Zd
 Z ed           G d de                      Zedk    r ej                     dS dS )    N)cudaint32
complex128void)types)TypingError)unittestCUDATestCaseskip_on_cudasim   )test_struct_model_type
TestStructc                     t           j                            dt                    }t	          |j        d                   D ]}| |         ||<   t	          |j        d                   D ]}||         ||<   d S )N  dtyper   r   localarrayr   rangeshapeABCis       e/var/www/html/test/jupyter/venv/lib/python3.11/site-packages/numba/cuda/tests/cudapy/test_localmem.pyculocalr   
   ~    
U++A171:  t!171:  t!     c                     t           j                            dt                    }t	          |j        d                   D ]}| |         ||<   t	          |j        d                   D ]}||         ||<   d S )Nd   r   r   )r   r   r   r   r   r   r   s       r   culocalcomplexr#      s~    
J//A171:  t!171:  t! r    c                     t           j                            dt                    }t	          |j        d                   D ]}| |         ||<   t	          |j        d                   D ]}||         ||<   d S )N)   r   r   r   r   s       r   culocal1tupler&      r   r    z'PTX inspection not available in cudasimc                       e Zd Zd Zd Zd Zd Z ed          d             Z ed          d             Z	 ed          d             Z
 ed          d	             Zd
 Zd Zd Zd ZdS )TestCudaLocalMemc                    t           d d          t           d d          f} t          j        |          t                    }|                     d|                    |          v            t          j        dd          }t          j        |          } |d         ||           |                     t          j	        ||k                         d S )Nz.localr   r   r   r   r   )
r   r   jitr   
assertTrueinspect_asmnparange
zeros_likeall)selfsigjculocalr   r   s        r   test_local_arrayz!TestCudaLocalMem.test_local_array$   s    QQQxqqq" 48C==))H$8$8$=$==>>>Id'***M!q!qAv'''''r    c                     t          j        d          t                    }t          j        dd          }t          j        |          } |d         ||           |                     t          j        ||k                         dS )zGEnsure that local arrays can be constructed with 1-tuple shape
        zvoid(int32[:], int32[:])r%   r   r   r*   N)r   r+   r&   r.   r/   r0   r,   r1   )r2   r4   r   r   s       r   test_local_array_1_tuplez)TestCudaLocalMem.test_local_array_1_tuple-   s}     848677FF Iaw'''M!q!qAv'''''r    c                 (   d} t          j        |          t                    }t          j        dd          dz
  dz  }t          j        |          } |d         ||           |                     t          j        ||k                         d S )Nz"void(complex128[:], complex128[:])r"   r   r   r   y               @r*   )r   r+   r#   r.   r/   r0   r,   r1   )r2   r3   jculocalcomplexr   r   s        r   test_local_array_complexz)TestCudaLocalMem.test_local_array_complex8   s    2'$(3--77Ys,///!3r9M!a###qAv'''''r    c                     t          t          |j                                                            j        }|j        d         j        }|                     ||           d S )Nl)nextiter	overloadsvalues_type_annotationtypemapr   assertEqual)r2   fr   
annotationl_dtypes        r   check_dtypezTestCudaLocalMem.check_dtype@   sT    $q{11334455F
$S)/%(((((r    zCan't check typing in simulatorc                     t          j        t          t          d d d                             d             }|                     |t                     d S )Nr   c                 ~    t           j                            dt                    }| d         |d<   |d         | d<   d S N
   r   r   )r   r   r   r   xr<   s     r   rD   z,TestCudaLocalMem.test_numba_dtype.<locals>.fJ   s:    
  5 11AQ4AaDQ4AaDDDr    r   r+   r   r   rG   r2   rD   s     r   test_numba_dtypez!TestCudaLocalMem.test_numba_dtypeG   X     
$uSSqSz""	#	#	 	 
$	#	
 	E"""""r    c                     t          j        t          t          d d d                             d             }|                     |t                     d S )Nr   c                     t           j                            dt          j                  }| d         |d<   |d         | d<   d S rJ   )r   r   r   r.   r   rL   s     r   rD   z,TestCudaLocalMem.test_numpy_dtype.<locals>.fU   s<    
  28 44AQ4AaDQ4AaDDDr    rN   rO   s     r   test_numpy_dtypez!TestCudaLocalMem.test_numpy_dtypeR   rQ   r    c                     t          j        t          t          d d d                             d             }|                     |t                     d S )Nr   c                 t    t           j                            dd          }| d         |d<   |d         | d<   d S )NrK   r   r   r   r   r   r   rL   s     r   rD   z-TestCudaLocalMem.test_string_dtype.<locals>.f`   s:    
  7 33AQ4AaDQ4AaDDDr    rN   rO   s     r   test_string_dtypez"TestCudaLocalMem.test_string_dtype]   rQ   r    c           	          d}|                      t          |          5  t          j        t	          t
          d d d                             d             }d d d            d S # 1 swxY w Y   d S )Nz*.*Invalid NumPy dtype specified: 'int33'.*r   c                 t    t           j                            dd          }| d         |d<   |d         | d<   d S )NrK   int33r   r   rW   rL   s     r   rD   z5TestCudaLocalMem.test_invalid_string_dtype.<locals>.fm   s:    J$$Rw$77t!t!r    )assertRaisesRegexr   r   r+   r   r   )r2   rerD   s      r   test_invalid_string_dtypez*TestCudaLocalMem.test_invalid_string_dtypeh   s     :##K44 	 	Xd51:&&''  ('	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   9A$$A(+A(c                     t          j        t          t          d d d                             d             }|                     |t                     d S )Nr   c                 ~    t           j                            dt                    }| d         |d<   |d         | d<   d S rJ   )r   r   r   r   rL   s     r   rD   z<TestCudaLocalMem.test_type_with_struct_data_model.<locals>.ft   s;    
  +A BBAQ4AaDQ4AaDDDr    )r   r+   r   r   rG   rO   s     r    test_type_with_struct_data_modelz1TestCudaLocalMem.test_type_with_struct_data_models   sY    	$-ccc233	4	4	 	 
5	4	
 	233333r    c           	         t          j        t          t          d d d         t          d d d                             d             }t	          j        dd          }t	          j        dd          } |d         ||           t          |          D ]\  }}|                     ||           t          |          D ]\  }}|                     ||dz             d S )Nr   c                 r   t           j                            dt                    }t	          t          |                    D ]4}t          t          |          t          |dz                      }|||<   5t	          t          |                    D ]"}||         j        | |<   ||         j	        ||<   #d S )NrK   r      )
r   r   r   r   r   lenr   r   rM   y)outxoutyarrr   objs        r   rD   z6TestCudaLocalMem.test_struct_model_type_arr.<locals>.f}   s     *""2-C"DDC3s88__   q5Q<<88A3s88__ # #a&(Qa&(Q# #r    )rK   r   r   r*   rd   )r   r+   r   r   r.   r   	enumeraterC   )r2   rD   arrxarryr   rM   rf   s          r   test_struct_model_type_arrz+TestCudaLocalMem.test_struct_model_type_arr|   s    	$uSSqSz51:..	/	/
	# 
	# 
0	/
	# xW---xW---$ddOO 	# 	#DAqQ""""dOO 	' 	'DAqQA&&&&	' 	'r    c                     t           j        fd            }t          j        dt          j                  } |d         |           |                     |d         |           d S )Nc                 ^    t           j                                      }|j        | d<   d S )Nr   r   )r   r   r   size)ari   r   tys     r   sz8TestCudaLocalMem._check_local_array_size_fp16.<locals>.s   s,    *""5"33C8AaDDDr    r   r   r*   r   )r   r+   r.   zerosfloat16rC   )r2   r   expectedrs   rt   results    ` `  r   _check_local_array_size_fp16z-TestCudaLocalMem._check_local_array_size_fp16   sw    		 	 	 	 	 
	 !2:...$H-----r    c                     |                      ddt          j                   |                      ddt          j                   d S )Nrd   )ry   r   rv   r.   )r2   s    r   test_issue_fp16_supportz(TestCudaLocalMem.test_issue_fp16_support   s<    ))!Q>>>))!Q
;;;;;r    N)__name__
__module____qualname__r5   r7   r:   rG   r   rP   rT   rX   r^   ra   rn   ry   r{    r    r   r(   r(   "   s(       ( ( (	( 	( 	(( ( () ) ) _677# # 87# _677# # 87# _677# # 87# _677  874 4 4' ' '0. . .< < < < <r    r(   __main__)numpyr.   numbar   r   r   r   
numba.corer   numba.core.errorsr   numba.cuda.testingr	   r
   r   extensions_usecasesr   r   r   r#   r&   r(   r|   mainr   r    r   <module>r      s-       / / / / / / / / / / / /       ) ) ) ) ) ) F F F F F F F F F F C C C C C C C C       :;;}< }< }< }< }<| }< }< <;}<@ zHMOOOOO r    