
    sh_n                        S r SSKrSSKJr  SSKrSSKrSSKJr  SSKr	SSK
JrJr  SSKJrJrJrJr  SSKJrJrJrJrJrJrJr  SSKJrJrJr  SS	KJr  SS
K J!r!  SSKJ"r"J#r#J$r$  SSK%J&r&  SSK'J(r(  SSKJ)r)  SSK*J+r+J,r,  SSK-J.r.  SSK/J0r0J1r1  SSK2J3r3  SSK
J4r4  SSK5J6r7  SSK8J9r:  SSK;J<r<  SSKJ=r=  SSK>J?r?J@r@  SSKAJBrBJCrC  SSKDJErE   " S S\5      rF\F\S'    " S S\+5      rG\1" 5       rH " S S \&5      rI " S! S"\!5      rJ\J" S5      rK " S# S$\5      rL\L\\S   '    " S% S&\5      rM\M\\S   '   \HR                  \R                  5      S' 5       rP\HR                  \R                  \R                  5      S( 5       rT\HR                  \R                  5      S) 5       rU\HR                  \R                  5      S* 5       rW\" SS+9S, 5       rX\" \R                  SS+9S- 5       rZ " S. S/\5      r[ " S0 S1\5      r\\]S2:X  a  \R                  " 5         gg)3a  This tests the target extension API to ensure that rudimentary expected
behaviours are present and correct. It uses a piece of fake hardware as a
target, the Dummy Processing Unit (DPU), to do this. The DPU borrows a lot from
the CPU but is part of the GPU class of target. The DPU target has deliberately
strange implementations of fundamental operations so as to make it identifiable
in testing.    N)TestCase)cached_property)njittypes)overloadoverload_attributeoverload_classmethod	intrinsic)JitDecoratortarget_registrydispatcher_registryjit_registrytarget_overrideGPUresolve_dispatcher_from_str)utilsfastmathpasserrors)
Dispatcher)TargetDescriptor)cputypingcgutils)BaseContext)global_compiler_lock)callconv)
CPUCodegenJITCodeLibrary)PyCallWrapper)RegistryLoaderRegistry)typeof)_dynfunc)ir)rtsys)compiler)CompilerBaseDefaultPassBuilder)FunctionPassregister_pass)PreLowerStripPhisc                       \ rS rSrSrg)DPU3    N)__name__
__module____qualname____firstlineno____static_attributes__r/       u/Users/tiagomarins/Projetos/claudeai/copy_bank/venv/lib/python3.13/site-packages/numba/tests/test_target_extension.pyr-   r-   3   s    r5   r-   dpuc                   0    \ rS rSr\rS rS rS rS r	Sr
g)JITDPUCodegen>   c                    U R                  5       US'   [        R                  R                  5       R                  nUR                  S5      (       a  SnOUR                  S5      (       a  SnOSnX1S'   SUS	'   U R                  US
'   [        R                  " [        R                  R                  5      nSUR                  ;   a  SUS'   g g )Nr   x86staticppcpicdefaultreloc
jitdefault	codemodelfeaturesjitT)_get_host_cpu_namellTargetfrom_default_triplename
startswith_tm_featuresr   pysignaturecreate_target_machine
parameters)selfoptionsarchreloc_modelsigs        r6   _customize_tm_options#JITDPUCodegen._customize_tm_optionsC   s    002yy,,.33??5!!"K__U##K#K&+ #//
 		 ? ?@CNN"!GEN #r5   c                 "    U R                  5       $ N)_get_host_cpu_featuresrP   s    r6   _customize_tm_features$JITDPUCodegen._customize_tm_featuresZ   s    **,,r5   c                 :    U R                   R                  U5        g rX   )_engine
add_module)rP   modules     r6   _add_moduleJITDPUCodegen._add_module^   s    'r5   c                     U R                   R                  U5      n[        R                  S-  R	                  U5      n[        R                  " [        U5      5      US'   g)zbSet the environment address.

Update the GlobalVariable named *env_name* to the address of *env*.
   r   N)r^   get_global_value_addressctypesc_void_pfrom_addressid)rP   env_nameenvgvaddrenvptrs        r6   set_envJITDPUCodegen.set_enva   sG    
 66x@//A%33F;OOBsG,q	r5   r/   N)r0   r1   r2   r3   r   _library_classrU   r[   ra   rn   r4   r/   r5   r6   r9   r9   >   s     $N".-(-r5   r9   c                       \ rS rSrSrS r\S 5       rS r\	S 5       r
S r\S 5       rS	 rS
 rS rS r SS jrS rS rSrg)
DPUContextq   Tc                 8    U R                   R                  U5      $ rX   )_internal_codegen_create_empty_module)rP   rJ   s     r6   create_moduleDPUContext.create_moduleu   s    %%::4@@r5   c                 p    [        S5      U l        [        R                  " U 5        U R	                  5         g )Nz
numba.exec)r9   ru   r%   
initializerefreshrZ   s    r6   initDPUContext.initx   s&    !.|!<r5   c                     [         n U R                  U   nU R	                  U5        U R
                  R                  5         g ! [         a    [        U5      nX R                  U'    NQf = frX   )dpu_function_registry_registriesKeyErrorr    install_registrytyping_contextr{   )rP   registryloaders      r6   r{   DPUContext.refresh   si    (	0%%h/F 	h' 	##%  	0#H-F)/X&	0s   A #A)(A)c                 .    U R                   R                  $ rX   )ru   target_datarZ   s    r6   r   DPUContext.target_data   s    %%111r5   c                     U R                   $ rX   ru   rZ   s    r6   codegenDPUContext.codegen   s    %%%r5   c                 .    [         R                  " U 5      $ rX   )r   CPUCallConvrZ   s    r6   	call_convDPUContext.call_conv   s    ##D))r5   c                 ~    [         R                  " X[        R                  S   5      n[        R
                  " XUSS9$ )z
From the given *envptr* (a pointer to a _dynfunc.Environment object),
get a EnvBody allowing structured access to environment fields.
offsetof_env_bodyT)refcast_ref)r   pointer_addr#   
_impl_infor   EnvBody)rP   builderrm   body_ptrs       r6   get_env_bodyDPUContext.get_env_body   s:    
 &&X001DE
 {{4hFFr5   c                 P   U R                  UR                  U R                  U R                  5      5      nUR	                  U5      nU R                  U5      nUR                  X0R                  R                  S9  U R                  X5      nUR                  U R                  XS5      $ )N)	debug_msg)declare_env_globalr`   get_env_namefndescloadget_python_apiemit_environment_sentryrj   r   get_env_managerenvironment)rP   r   envgvenvargpyapienv_bodys         r6   r   DPUContext.get_env_manager   s    ''NND--dkk:
 e$##G,%%kk22 	& 	
 $$W5$$T%5%5xHHr5   c                 R    [         R                  " UU[        R                  S   US9$ )zf
From the given *genptr* (a pointer to a _dynfunc.Generator object),
get a pointer to its state area.
offsetof_generator_state)return_type)r   r   r#   r   )rP   r   genptrr   s       r6   get_generator_stateDPUContext.get_generator_state   s/    
 "" :;#	
 	
r5   c                     U R                   (       a   [        R                  " XR                   5        UR                  [        R
                  5        g rX   )fastmathr   rewrite_moduleadd_linking_libraryr%   library)rP   modr   s      r6   post_loweringDPUContext.post_lowering   s-    =='']];##EMM2r5   c           
      (   U R                  S5      nU R                  R                  UR                  UR                  5      n[
        R                  " XgUR                  5      n[        U UUUUUUS9n	U	R                  5         UR                  U5        g )Nwrapper)call_helperrelease_gil)rw   r   get_function_typerestypeargtypesllirFunctionllvm_func_namer   buildadd_ir_module)
rP   r   r   rk   r   r   wrapper_modulefntywrapper_calleer   s
             r6   create_cpython_wrapper!DPUContext.create_cpython_wrapper   s     ++I6~~//P&"7"7
  ##
 	n-r5   c                     g rX   r/   )rP   r   r   rk   r   s        r6   create_cfunc_wrapperDPUContext.create_cfunc_wrapper   s    r5   c                 >   UR                  UR                  5      nSUR                  < 3n[        R                  " UR                  5       UR                  R                  S5      S   UUUU45      nUR                  R                  U R                  U5      U5        U$ )z
Returns
-------
(cfunc, fnptr)

- cfunc
    callable function (Can be None)
- fnptr
    callable function address
- env
    an execution environment (from _dynfunc)
zcompiled wrapper for .)
get_pointer_to_functionllvm_cpython_wrapper_namequalnamer#   make_functionlookup_modulesplitr   rn   r   )rP   r   r   rk   fnptrdoccfuncs          r6   get_executableDPUContext.get_executable   s     //,,
 ,2??<&&  "OO!!#&r*J
 	 1 1& 93?r5   r   N)F)r0   r1   r2   r3   allow_dynamic_globalsrw   r   r|   r{   propertyr   r   r   r   r   r   r   r   r   r   r   r4   r/   r5   r6   rr   rr   q   s     A  
& 2 2& * *G
I

3 >C.(r5   rr   c                   l    \ rS rSr\R
                  r\S 5       r\S 5       r	\
S 5       r\
S 5       rSrg)	DPUTarget   c                 B    [        U R                  U R                  5      $ rX   )rr   r   _target_namerZ   s    r6   _toplevel_target_context"DPUTarget._toplevel_target_context   s     $--t/@/@AAr5   c                 ,    [         R                  " 5       $ rX   )r   ContextrZ   s    r6   _toplevel_typing_context"DPUTarget._toplevel_typing_context  s     ~~r5   c                     U R                   $ )z%
The target context for DPU targets.
)r   rZ   s    r6   target_contextDPUTarget.target_context      
 ,,,r5   c                     U R                   $ )z%
The typing context for CPU targets.
)r   rZ   s    r6   r   DPUTarget.typing_context  r   r5   r/   N)r0   r1   r2   r3   r   CPUTargetOptionsrQ   r   r   r   r   r   r   r4   r/   r5   r6   r   r      s]    ""GB B     - - - -r5   r   c                   ,   ^  \ rS rSr\rU 4S jrSrU =r$ )DPUDispatcheri  c                 n   > [        S5         [        TU ]	  U5      sS S S 5        $ ! , (       d  f       g = f)Nr7   )r   supercompile)rP   rT   	__class__s     r6   r   DPUDispatcher.compile  s#    U#7?3' $##s   &
4r/   )	r0   r1   r2   r3   
dpu_targettargetdescrr   r4   __classcell__)r   s   @r6   r   r     s    K( (r5   r   c                   ,    \ rS rSrS rS rS rS rSrg)djiti*  c                     Xl         X l        g rX   )_args_kwargs)rP   argskwargss      r6   __init__djit.__init__+  s    
r5   c                     [        U5      S:  d   eU(       a  US   nOU R                  S   nX l        U R                  5       $ )N   r   )lenr   py_funcdispatcher_wrapper)rP   r   funcs      r6   __call__djit.__call__/  s?    4y1}}7D::a=D&&((r5   c                 "    [         [        S      $ )z
Returns the dispatcher
r7   )r   r   rZ   s    r6   get_dispatcherdjit.get_dispatcher9  s     #?5#9::r5   c                     U R                  5       n0 nSU R                  ;   a  SUS'   [        R                  nSU R                  ;   a  U R                  S   nU" U R                  UUS9$ )NnopythonTpipeline_class)r  targetoptionsr  )r  r   r&   Compilerr  )rP   disptoptr  s       r6   r  djit.dispatcher_wrapper?  sm    ""$%#D "**t||+!\\*:;NLL)
 	
r5   )r   r   r  N)	r0   r1   r2   r3   r  r	  r  r  r4   r/   r5   r6   r   r   *  s    );
r5   r   c                 "    U R                  5       $ rX   )get_dummy_value)contextr   typyvals       r6   constant_dummyr  [  s    ""$$r5   c                     U R                  XR                  UR                  5      nU R                  XUR                  U5      $ rX   )get_constant_genericliteral_typeliteral_valuecast)r  r   fromtytotyvallits         r6   literal_int_to_numberr%  a  s>    

&
&$$f&:&:C <<f&9&94@@r5   c                 4    U R                  U5      nU" U5      $ rX   get_value_typer  r   r  r  ltys        r6   	const_intr+  j      

 
 
$Cu:r5   c                 4    U R                  U5      nU" U5      $ rX   r'  r)  s        r6   const_floatr.  q  r,  r5   targetc                     U" X5      nS nX44$ )Nc                      UR                   " U6 $ rX   subcgctxr   tyargsllargss       r6   r   intrin_add.<locals>.codegen|  s    {{F##r5   r/   tyctxxyrT   r   s        r6   
intrin_addr>  x  s    
A'C$ <r5   c                     [        U [        R                  5      (       a%  [        U[        R                  5      (       a  S nU$ g g )Nc                     [        X5      $ rX   )r>  )r<  r=  s     r6   implol_add.<locals>.impl  s    a##r5   )
isinstancer   Integer)r<  r=  rA  s      r6   ol_addrE    s8    !U]]##
1emm(D(D	$  )E#r5   c                   N    \ rS rSrSrS rS rS rS rS r	S r
S	 rS
 rS rSrg)TestTargetHierarchySelectioni  zThis tests that the target hierarchy is scanned in the right order,
that appropriate functions are selected based on what's available and that
the DPU target is distinctly different to the CPUc                 $   U R                  [        R                  5        U R                  [        R                  5        U R	                  [        [        R                  5      S5        U R	                  [        [        R                  5      S5        g)zChecks that the DPU registry only contains the things added

This test must be first to execute among all tests in this file to
ensure the no lazily loaded entries are added yet.
rd      N)assertFalser   	functionsgetattrsassertEqualr  casts	constantsrZ   s    r6   test_0_dpu_registry0TestTargetHierarchySelection.test_0_dpu_registry  sf     	.889.77828891=2<<=qAr5   c                    ^ S m[        TSS9S 5       n[        TSS9S 5       n[        5       U4S j5       n[        5       U4S j5       nU R                  U" 5       S	5        U R                  U" 5       S
5        g )Nc                     g rX   r/   r<  s    r6   my_funcATestTargetHierarchySelection.test_specialise_gpu.<locals>.my_func      r5   genericr/  c                     S nU$ )Nc                     SU -   $ Nrd   r/   rT  s    r6   rA  STestTargetHierarchySelection.test_specialise_gpu.<locals>.ol_my_func1.<locals>.impl      1ur5   r/   r<  rA  s     r6   ol_my_func1ETestTargetHierarchySelection.test_specialise_gpu.<locals>.ol_my_func1       Kr5   gpuc                     S nU$ )Nc                     SU -   $ N
   r/   rT  s    r6   rA  STestTargetHierarchySelection.test_specialise_gpu.<locals>.ol_my_func2.<locals>.impl      Avr5   r/   r^  s     r6   ol_my_func2ETestTargetHierarchySelection.test_specialise_gpu.<locals>.ol_my_func2       Kr5   c                     > T " S5      $ N   r/   rU  s   r6   dpu_fooATestTargetHierarchySelection.test_specialise_gpu.<locals>.dpu_foo      1:r5   c                     > T " S5      $ rm  r/   ro  s   r6   cpu_fooATestTargetHierarchySelection.test_specialise_gpu.<locals>.cpu_foo  rr  r5   rI     r   r   r   assertPreciseEqual)rP   r_  ri  rp  rt  rU  s        @r6   test_specialise_gpu0TestTargetHierarchySelection.test_specialise_gpu  s    	 
')	,	 
-	 
'%	(	 
)	 
	 
	 
	 
	
 		1-	1-r5   c                   ^ S m[        TSS9S 5       n[        TSS9S 5       n[        TSS9S 5       n[        5       U4S	 j5       n[        5       U4S
 j5       nU R                  U" 5       S5        U R                  U" 5       S5        g )Nc                     g rX   r/   rT  s    r6   rU  ATestTargetHierarchySelection.test_specialise_dpu.<locals>.my_func  rW  r5   rX  r/  c                     S nU$ )Nc                     SU -   $ r[  r/   rT  s    r6   rA  STestTargetHierarchySelection.test_specialise_dpu.<locals>.ol_my_func1.<locals>.impl  r]  r5   r/   r^  s     r6   r_  ETestTargetHierarchySelection.test_specialise_dpu.<locals>.ol_my_func1  ra  r5   rb  c                     S nU$ )Nc                     SU -   $ re  r/   rT  s    r6   rA  STestTargetHierarchySelection.test_specialise_dpu.<locals>.ol_my_func2.<locals>.impl  rh  r5   r/   r^  s     r6   ri  ETestTargetHierarchySelection.test_specialise_dpu.<locals>.ol_my_func2  rk  r5   r7   c                     S nU$ )Nc                     SU -   $ )Nd   r/   rT  s    r6   rA  STestTargetHierarchySelection.test_specialise_dpu.<locals>.ol_my_func3.<locals>.impl  s    Qwr5   r/   r^  s     r6   ol_my_func3ETestTargetHierarchySelection.test_specialise_dpu.<locals>.ol_my_func3  s     Kr5   c                     > T " S5      $ rm  r/   ro  s   r6   rp  ATestTargetHierarchySelection.test_specialise_dpu.<locals>.dpu_foo  rr  r5   c                     > T " S5      $ rm  r/   ro  s   r6   rt  ATestTargetHierarchySelection.test_specialise_dpu.<locals>.cpu_foo  rr  r5   ]   rv  rw  )rP   r_  ri  r  rp  rt  rU  s         @r6   test_specialise_dpu0TestTargetHierarchySelection.test_specialise_dpu  s    	 
')	,	 
-	 
'%	(	 
)	 
'%	(	 
)	 
	 
	 
	 
	
 		2.	1-r5   c                 X  ^ S m[        TSS9S 5       n[        SS9U4S j5       n[        R                  [        R                  4nU R                  U5       nU" 5         S S S 5        / SQnU H(  nU R                  U[        WR                  5      5        M*     g ! , (       d  f       NA= f)	Nc                     g rX   r/   rT  s    r6   rU  JTestTargetHierarchySelection.test_no_specialisation_found.<locals>.my_func  rW  r5   cudar/  c                     S $ )Nc                     g rX   r/   rT  s    r6   <lambda>dTestTargetHierarchySelection.test_no_specialisation_found.<locals>.ol_my_func_cuda.<locals>.<lambda>  s    Tr5   r/   rT  s    r6   ol_my_func_cudaRTestTargetHierarchySelection.test_no_specialisation_found.<locals>.ol_my_func_cuda  s    !!r5   Tr  c                     > T " S5        g r[  r/   ro  s   r6   rp  JTestTargetHierarchySelection.test_no_specialisation_found.<locals>.dpu_foo  s
    AJr5   )8Function resolution cannot find any matches for functionz-test_no_specialisation_found.<locals>.my_funczfor the current target:z''numba.tests.test_target_extension.DPU')	r   r   r   UnsupportedErrorTypingErrorassertRaisesassertInstr	exception)rP   r  rp  acceptraisesmsgsmsgrU  s          @r6   test_no_specialisation_found9TestTargetHierarchySelection.test_no_specialisation_found  s    	 
'&	)	" 
*	" 
t		 
	
 ))6+=+=>v&&I ';
 CMM#s6#3#345  '&s   B
B)c                     U R                  [        R                  5       n[        SS9S 5       nU" 5         S S S 5        SnU R	                  U[        WR                  5      5        g ! , (       d  f       N6= f)Ninvalid_silicon)_targetc                      g rX   r/   r/   r5   r6   fooATestTargetHierarchySelection.test_invalid_target_jit.<locals>.foo
  s    r5   1No target is registered against 'invalid_silicon')r  r   NonexistentTargetErrorr   r  r  r  )rP   r  r  r  s       r6   test_invalid_target_jit4TestTargetHierarchySelection.test_invalid_target_jit  sf    v<<=+, - E > Bc3v//01 >=s   A''
A5c                   ^ S mU R                  [        R                  5       n[        TSS9S 5       n[        U4S j5       nU" 5         S S S 5        SnU R                  U[        WR                  5      5        g ! , (       d  f       N6= f)Nc                      g rX   r/   r/   r5   r6   barFTestTargetHierarchySelection.test_invalid_target_overload.<locals>.bar  rW  r5   r  r/  c                      S $ )Nc                      g rX   r/   r/   r5   r6   r  [TestTargetHierarchySelection.test_invalid_target_overload.<locals>.ol_bar.<locals>.<lambda>  s    r5   r/   r/   r5   r6   ol_barITestTargetHierarchySelection.test_invalid_target_overload.<locals>.ol_bar  s    $$r5   c                     > T " 5         g rX   r/   )r  s   r6   r  FTestTargetHierarchySelection.test_invalid_target_overload.<locals>.foo  s    r5   r  )r  r   r  r   r   r  r  r  )rP   r  r  r  r  r  s        @r6   test_invalid_target_overload9TestTargetHierarchySelection.test_invalid_target_overload  s    	
 v<<=c"34% 5%   E > Bc3v//01 >=s   'A;;
B	c                   ^^^ [        SS9S 5       m[        SS9S 5       m[        SS9S 5       m[        U4S j5       nU R                  U" 5       S	5        [        U4S
 j5       nU R                  U" 5       S5        [        U4S j5       n[        R                  [        R
                  4nU R                  U5       nU" 5         SSS5        / SQnU H(  nU R                  U[        WR                  5      5        M*     [        SS9U4S j5       nU R                  U" 5       S5        [        SS9U4S j5       n	U R                  U	" 5       S5        [        SS9U4S j5       n
[        R                  [        R
                  4nU R                  U5       nU
" 5         SSS5        / SQnU H(  nU R                  U[        WR                  5      5        M*     g! , (       d  f       GN= f! , (       d  f       NS= f)z
Test to make sure that targets can share generic implementations and
cannot reach implementations that are not in their target hierarchy.
rX  r/  c                     U" X5      nS nX44$ )Nc                      UR                   " U6 $ rX   )mulr5  s       r6   r   cTestTargetHierarchySelection.test_intrinsic_selection.<locals>.intrin_math_generic.<locals>.codegen3      {{F++r5   r/   r:  s        r6   intrin_math_genericRTestTargetHierarchySelection.test_intrinsic_selection.<locals>.intrin_math_generic/      A'C, <r5   r7   c                     U" X5      nS nX44$ )Nc                      UR                   " U6 $ rX   r3  r5  s       r6   r   _TestTargetHierarchySelection.test_intrinsic_selection.<locals>.intrin_math_dpu.<locals>.codegen<  r  r5   r/   r:  s        r6   intrin_math_dpuNTestTargetHierarchySelection.test_intrinsic_selection.<locals>.intrin_math_dpu8  r  r5   r   c                     U" X5      nS nX44$ )Nc                      UR                   " U6 $ rX   )addr5  s       r6   r   _TestTargetHierarchySelection.test_intrinsic_selection.<locals>.intrin_math_cpu.<locals>.codegenE  r  r5   r/   r:  s        r6   intrin_math_cpuNTestTargetHierarchySelection.test_intrinsic_selection.<locals>.intrin_math_cpuA  r  r5   c                     > T " SS5      $ NrI     r/   r  s   r6   cpu_foo_specificOTestTargetHierarchySelection.test_intrinsic_selection.<locals>.cpu_foo_specificK      "1a((r5   rn  c                     > T " SS5      $ r  r/   r  s   r6   cpu_foo_genericNTestTargetHierarchySelection.test_intrinsic_selection.<locals>.cpu_foo_genericR      &q!,,r5      c                     > T " SS5      $ r  r/   r  s   r6   cpu_foo_dpuJTestTargetHierarchySelection.test_intrinsic_selection.<locals>.cpu_foo_dpuY  r  r5   N)r  zintrinsic intrin_math_dpufor the current targetTr  c                     > T " SS5      $ r  r/   r  s   r6   dpu_foo_specificOTestTargetHierarchySelection.test_intrinsic_selection.<locals>.dpu_foo_specifich  r  r5   r   c                     > T " SS5      $ r  r/   r  s   r6   dpu_foo_genericNTestTargetHierarchySelection.test_intrinsic_selection.<locals>.dpu_foo_generico  r  r5   c                     > T " SS5      $ r  r/   r  s   r6   dpu_foo_cpuJTestTargetHierarchySelection.test_intrinsic_selection.<locals>.dpu_foo_cpuv  r  r5   )r  zintrinsic intrin_math_cpur  )r
   r   rM  r   r  r  r  r  r  r  r   )rP   r  r  r  r  r  r  r  r  r  r  r  r  r  s              @@@r6   test_intrinsic_selection5TestTargetHierarchySelection.test_intrinsic_selection(  s    
)	$	  
%	  
%	 	  
!	  
%	 	  
!	  
	) 
	) 	)+Q/ 
	- 
	- 	*B/ 
	) 
	) ))6+=+=>v&&M '+ CMM#s6#3#345  
t		) 
	) 	)+R0 
t		- 
	- 	*B/ 
t		) 
	) ))6+=+=>v&&M '+ CMM#s6#3#345 G '&: '&s   =GG)
G&)
G7c                   ^^	^
 S m	[        SS9U	4S j5       m
[        [        R                  SSSS0S9U
4S	 j5       n[	        [
        R                  SSS0S9S
 5       nS m[	        TSSS0S9S 5       nSSKJn  U" S5         [        SS9U4S j5       nU" 5       nS S S 5        SSK
Jn  U R                  WUR                  5        g ! , (       d  f       N1= f)Nc                 P   UR                   UR                   :X  a  U$ UR                   UR                   :  a   UR                  X R                  U5      5      $ UR                  (       a   UR	                  X R                  U5      5      $ UR                  X R                  U5      5      $ rX   )bitwidthtruncr(  signedsextzext)r  r   r#  r!  r"  s        r6   cast_integerKTestTargetHierarchySelection.test_overload_allocation.<locals>.cast_integer  s~    }}/
0}}S*@*@*FGG||C)?)?)EFF ||C)?)?)EFFr5   r7   r/  c                 |   > U4S jnSSK Jn  [        R                  " [        R                  5      nU" XQU5      nXc4$ )z;Intrinsic to call into the allocator for Array
            c                    > Uu  pET" XUUR                   S   [        R                  5      nU R                  R	                  XU5      nU$ r[  )r   r   uint32nrtmeminfo_alloc_aligned)	r  r   	signaturer   	allocsizealign	align_u32meminfor  s	           r6   r   \TestTargetHierarchySelection.test_overload_allocation.<locals>.intrin_alloc.<locals>.codegen  sM    %)" )5)2):ELLJ	!++;;G<EGr5   r   )r  )numba.core.typingr  r   MemInfoPointervoidptr)	typingctxr  r  r   r  miprT   r  s          r6   intrin_allocKTestTargetHierarchySelection.test_overload_allocation.<locals>.intrin_alloc  s5    	 4&&u}}5CCE2C<r5   	_allocater  T)r0  jit_optionsc                    > U4S jnU$ )Nc                    > T" X5      $ rX   r/   )clsr  r  r  s      r6   rA  aTestTargetHierarchySelection.test_overload_allocation.<locals>._ol_arr_allocate_dpu.<locals>.impl  s    #I55r5   r/   )r  r  r  rA  r  s       r6   _ol_arr_allocate_dpuSTestTargetHierarchySelection.test_overload_allocation.<locals>._ol_arr_allocate_dpu  s    6Kr5   c                     S nU$ )Nc                 B    [         R                  R                  U S5      $ rm  )r   Arrayr  )ns    r6   rA  ZTestTargetHierarchySelection.test_overload_allocation.<locals>.ol_empty_impl.<locals>.impl  s    {{,,Q22r5   r/   )r  rA  s     r6   ol_empty_implLTestTargetHierarchySelection.test_overload_allocation.<locals>.ol_empty_impl  s    3Kr5   c                      g rX   r/   r/   r5   r6   buffer_funcJTestTargetHierarchySelection.test_overload_allocation.<locals>.buffer_func  rW  r5   c                      S n U $ )Nc                  .    [         R                  " S5      $ re  )npemptyr/   r5   r6   rA  `TestTargetHierarchySelection.test_overload_allocation.<locals>.ol_buffer_func_impl.<locals>.impl  s    xx|#r5   r/   )rA  s    r6   ol_buffer_func_implRTestTargetHierarchySelection.test_overload_allocation.<locals>.ol_buffer_func_impl  s    $Kr5   r   )r   r  c                     > T " 5       $ rX   r/   )r   s   r6   r  BTestTargetHierarchySelection.test_overload_allocation.<locals>.foo  s    "}$r5   )r  )r
   r	   r   r  r   r$  r%  numba.core.target_extensionr   r   numba.core.runtimer  assertIsInstanceMemInfo)rP   r  r  r'  r   r  rr  r   r  r  s           @@@r6   test_overload_allocation5TestTargetHierarchySelection.test_overload_allocation  s    	G 
%	 	  
!	 & 
ekk;u+5d*;
=	
=	
 
"((5z$6G	H	 
I	
	 
+e*T9J	K	 
L	
 	@ U#4 % !%A	 $
 	+a- $#s   C
Cc                    U R                  5       u  p[        U" 5       5      n[        USSS9S 5       nU R                  [        R
                  S5         [        [        R                  " U5      5      S 5       nS S S 5        [        [        R                  " [        R                  S S S2   U5      5      S 5       ng ! , (       d  f       NI= f)	Ndpu_onlyr7   r/  c                     S nU$ )Nc                     g)N*   r/   )objs    r6   impcTestTargetHierarchySelection.test_overload_attribute_target.<locals>.ov_dummy_dpu_attr.<locals>.imp  s    r5   r/   )r7  r8  s     r6   ov_dummy_dpu_attrVTestTargetHierarchySelection.test_overload_attribute_target.<locals>.ov_dummy_dpu_attr  s     Jr5   zUnknown attribute 'dpu_only'c                     U R                   $ rX   r3  rT  s    r6   illegal_target_attr_use\TestTargetHierarchySelection.test_overload_attribute_target.<locals>.illegal_target_attr_use  s    zz!r5   rd   c                 "    UR                   U S'   g )Nr   r=  )resdummys     r6   cuda_target_attr_useYTestTargetHierarchySelection.test_overload_attribute_target.<locals>.cuda_target_attr_use  s    ^^CFr5   )make_dummy_typer"   r   assertRaisesRegexr   r  r   r   int64r   void)rP   MyDummyMyDummyTypemydummy_typer:  r>  rC  s          r6   test_overload_attribute_target;TestTargetHierarchySelection.test_overload_attribute_target  s    #335gi(	KE	B	 
C	 ##F$6$6$BD%++l+," -"D 
ejjSqS)<8	9	$ 
:	$D Ds   &B==
Cr/   N)r0   r1   r2   r3   __doc__rP  ry  r  r  r  r  r  r0  rL  r4   r/   r5   r6   rG  rG    s;    9B .D(.T68
22*Z6xC.J$r5   rG  c                       \ rS rSrSrS rSrg)TestTargetOffloadi  a  In this use case the CPU compilation pipeline is extended with a new
compilation pass that runs just prior to lowering. The pass looks for
function calls and when it finds one it sees if there's a DPU function
available that is a valid overload for the function call. If there is one
then it swaps the CPU implementation out for a DPU implementation. This
producing an "offload" effect.
c                   ^^ Sm[        [        R                  SS9S 5       n[        SS9S 5       nU R	                  U" S5      S	5        [
        S
 5       nU R	                  U" S5      [        R                  " S5      5        [        SSS9 " U4S jS[        5      5       m " U4S jS[        5      n[        US9S 5       nU R	                  U" S5      S	[        R                  " S5      45        g )NFr7   r/  c                     S nU$ )Nc                     g)N    ,Ar/   rT  s    r6   dpu_sin_implQTestTargetOffload.test_basic_offload.<locals>.ol_np_sin_DPU.<locals>.dpu_sin_impl  s    r5   r/   )r<  rU  s     r6   ol_np_sin_DPU;TestTargetOffload.test_basic_offload.<locals>.ol_np_sin_DPU  s       r5   Tr  c                 .    [         R                  " U 5      $ rX   r$  sinrT  s    r6   r  1TestTargetOffload.test_basic_offload.<locals>.foo      66!9r5      rT  c                 .    [         R                  " U 5      $ rX   rZ  rT  s    r6   r  r\    r]  r5   )mutates_CFGanalysis_onlyc                   ,   > \ rS rSrSrS rU 4S jrSrg)@TestTargetOffload.test_basic_offload.<locals>.DispatcherSwitcheri  DispatcherSwitcherc                 0    [         R                  " U 5        g rX   )r)   r  rZ   s    r6   r  ITestTargetOffload.test_basic_offload.<locals>.DispatcherSwitcher.__init__  s    %%d+r5   c           
        > UR                   nSnUR                  R                  5        H  nUR                  S5       H  nUR                  UR
                  R                     nSn[        U5          UR                  UR                  UR                  U   R                  0 5      n[        U5      n	U	R                  R                  n
U
R                  Xh5        UR$                  R&                  U
l        Xul        SnS S S 5        M     M     U$ ! [          a1  nT(       a  SU SU SU 3n[#        U5         S nAS S S 5        GM  S nAff = f! , (       d  f       GM  = f)NFcallr7   z+Failed to find and compile an overload for z for z due to T)func_irblocksvalues
find_exprstypemapr  rJ   r   get_call_typer  	calltypesr   r   r   r   get_function	Exceptionprint	targetctx_codelib_stackr0  )rP   stateri  mutatedblkrh  functiontnamerT   r  hw_ctxer  _DEBUGs                r6   run_passITestTargetOffload.test_basic_offload.<locals>.DispatcherSwitcher.run_pass  sJ   --">>002C !$v 6#(==#@ % -U3)&.&<&<$)OO$)OOD$9$>$>$&'"
 (C5'I)-)9)9)H)H & 3 3H B !& > > #1 +0K&*GE 43 !7 3Z 3 $- )#)+88@zug N223)6 %(
 %*#J (' 43) 43s7   1E3A&D#E
E		E	5EE			E
Er/   N)r0   r1   r2   r3   _namer  r}  r4   )r|  s   r6   rd  rc    s    (E,0 0r5   rd  c                   "   > \ rS rSrU 4S jrSrg)@TestTargetOffload.test_basic_offload.<locals>.DPUOffloadCompileriB  c                    > [         R                  " U R                  5      nUR                  T[        5        UR                  5         U/$ rX   )r(   define_nopython_pipelineru  add_pass_afterr+   finalize)rP   pmrd  s     r6   define_pipelinesQTestTargetOffload.test_basic_offload.<locals>.DPUOffloadCompiler.define_pipelinesC  s9    '@@L!!"46GHtr5   r/   N)r0   r1   r2   r3   r  r4   )rd  s   r6   DPUOffloadCompilerr  B  s     r5   r  )r  c                 Z    [         R                  " U 5      [         R                  " U 5      4$ rX   )r$  r[  cosrT  s    r6   r  r\  K  s    66!9bffQi''r5   )
r   r$  r[  r   rx  r   r*   r)   r'   r  )rP   rW  r  r  rd  r|  s       @@r6   test_basic_offload$TestTargetOffload.test_basic_offload  s     
"&&	'	  
(	  
t		 
	 	A1 
	 
	 	Aq	2	5	>6	 6	 
?6	r	 	 
/	0	( 
1	( 	A266!9(=>r5   r/   N)r0   r1   r2   r3   rN  r  r4   r/   r5   r6   rP  rP    s    b?r5   rP  __main__)_rN  unittestnumba.tests.supportr   rf   operator	functoolsr   numpyr$  numbar   r   numba.extendingr   r   r	   r
   r+  r   r   r   r   r   r   r   
numba.corer   r   r   numba.core.dispatcherr   numba.core.descriptorsr   r   r   r   numba.core.baser   numba.core.compiler_lockr   r   numba.core.codegenr   r   numba.core.callwrapperr   numba.core.imputilsr    r!   numba.core.typing.typeofr"   r#   llvmlite.bindingbindingrG   llvmliter$   r   r,  r%   r&   numba.core.compilerr'   r(   numba.core.compiler_machineryr)   r*   numba.core.typed_passesr+   r-   r9   r   rr   r   r   r   r   lower_constantDummyr  
lower_castIntegerLiteralrD  r%  r+  Floatr.  r>  r  rE  rG  rP  r0   mainr/   r5   r6   <module>r     s    (   %  > >   3 2 , 3 + + ' 9  9 0 8 +    $  @ E 5#   
*-J *-\ !
 
E ER-  -: u
(J ( /< OE* +
%
< %
T (,_U# $ %%ekk2% 3%
 !!%"6"6FA GA %%emm4 5 %%ekk2 3 %  
(,,u% &T$8 T$n
k? k?\ zMMO r5   