U
    •Ê¹jT  ã                   @   sŒ   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
 d dlmZmZmZmZmZmZmZ dZeedœdd„Zed	d
�G dd„ dƒƒZdS )é    )Ú	dataclass)ÚCallableÚListÚSequenceÚTuple)ÚBindingÚCTypeÚ
NamedCType)ÚArgumentÚBaseTyÚBaseTypeÚListTypeÚNativeFunctionÚOptionalTypeÚTypez
	)ÚfÚreturnc                 C   s   | j j ¡ S )N)ÚfuncÚnameZunambiguous_name)r   © r   ú_/var/www/html/bwave/python/venv/lib/python3.8/site-packages/torchgen/executorch/api/unboxing.pyr      s    r   T)Úfrozenc                
   @   sä   e Zd ZU dZedef ed< ee e	e
e e
e f dœdd„Zddœeeee	eee
e e
e f d	œd
d„Zeeee	e
e e
e f dœdd„Zeeeee	e
e e
e f dœdd„Zeeeee	e
e e
e f dœdd„ZdS )ÚUnboxingaý  
    Takes a sequence of Bindings and unbox EValues to these Bindings. Return generated code that performs correct unboxing.
    A sample generated code:
    // aten::mul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
    void mul_out(EValue** stack) {
        EValue& self = *stack[0];
        EValue& other = *stack[1];
        EValue& out = *stack[2];
        const torch::executor::Tensor & self_base = self.to<torch::executor::Tensor>();
        const torch::executor::Tensor & other_base = other.to<torch::executor::Tensor>();
        torch::executor::Tensor & out_base = out.to<torch::executor::Tensor>();

        EXECUTORCH_SCOPE_PROF("native_call_mul.out");
        torch::executor::mul_outf(self_base, other_base, out_base);


    }
    .Úargument_type_gen)Úargsr   c           
         s’   ‡ fdd„t tˆ ƒƒD ƒ}g }ˆ D ]f}t|jtƒs@td|› �ƒ‚|j}| j|j|j|j	d�\}}}}	| 
|	¡ | 
|¡ | | |¡¡ q"||fS )Nc                    s$   g | ]}d ˆ | j › d|› d�‘qS )zEValue& z
 = *stack[z];)r   )Ú.0Úi©r   r   r   Ú
<listcomp>7   s     z.Unboxing.convert_arguments.<locals>.<listcomp>z7Unexpected argument type, expecting `Argument` but got ©Úmutable)ÚrangeÚlenÚ
isinstanceÚargumentr
   Ú	ExceptionÚargumenttype_evalue_convertÚtyper   Zis_writeÚextendÚappendÚ	with_name)
Úselfr   Z	code_listZbinding_listÚargr$   Zunboxed_nameÚ_ÚcodeÚdeclr   r   r   Úconvert_arguments4   s"    ÿ  ÿ

zUnboxing.convert_argumentsFr   )ÚtÚarg_namer    r   c                C   s´   | j |||d�j}t|tƒr<|› d�}| j|||d�\}}nlt|tƒrh|› d�}| j||||d�\}}n@t|tƒr”|› d�}| j||||d�\}}nt	d|› d|› �ƒ‚||||fS )	aÕ  
        Takes in the type, name and mutability corresponding to an argument, and generates a tuple of:
        (1) the C++ code necessary to unbox the argument
        (2) A Binding corresponding to the newly created unboxed variable, including variable name and its CType
        :param t: a `Type` of an argument
        :param arg_name: argument name
        :param mutable: boolean for whether this argument type is mutable
        :return: unboxed result
        )r    ZbindsZ_base)r2   Úout_nameÚctypeZ_opt_out)r2   r3   r1   r4   Z	_list_outzCannot handle type z. arg_name: )
r   r'   r#   r   Ú_gen_code_base_typer   Ú_gen_code_optional_typer   Ú_gen_code_list_typer%   )r+   r1   r2   r    r4   r3   r.   r/   r   r   r   r&   H   s4    

  ÿ

   ÿ

   ÿz$Unboxing.argumenttype_evalue_convert)r2   r3   r4   r   c              	   C   s.   |  ¡ › d|› d|› d|j dd�› d�gg fS )Nú ú = z.to<T©Z	strip_refz>();)Úcpp_type)r+   r2   r3   r4   r   r   r   r5   i   s    &ÿþzUnboxing._gen_code_base_type)r2   r3   r1   r4   r   c           
   
   C   sX   |› d�}|   |j|¡\}}}}	d|jdd�› d|› d|› d|jdd�› d�	 d	¡|	fS )
NZ_opt_inú
    Tr:   r8   r9   z.toOptional<z>();
            Ú
)r&   Úelemr;   Úsplit)
r+   r2   r3   r1   r4   Úin_nameÚres_nameZ	base_typeÚres_coder/   r   r   r   r6   p   s$    
 ÿ
ÿÿÿ
ÿýúz Unboxing._gen_code_optional_typec                 C   sd  |› d�}|› d�}g }|   |j|¡\}}	}
}t|jtƒrx|jjtjkrx| d|jdd�› d|› d|› d� 	d	¡¡ �nät|jtƒrÐ|jjtj
ks |jjtjkrÐ| d|jdd�› d|› d|› d
� 	d	¡¡ �nŒt|jtƒ�r|jjtjk�r| d|jdd�› d|› d|› d� 	d	¡¡ �n>t|jtƒ�rj|jjtjk�rj| d|jdd�› d|› d|› d� 	d	¡¡ nòt|jtƒ�rät|jjtƒ�rä|jjjtjk�rä| d|› d|› d|› d|› d|› d|› d|› d|› d|› d� 	d	¡¡ nx|d }| d|	jdd�› d|› d�¡ | d|› d|› dt |
¡› d|› d|› d|jdd�› d|› d|› d� 	d	¡¡ ||fS )NZ_list_inÚ_elemr<   Tr:   r8   r9   z!.toTensorList();
                r=   z.toIntList();
                z!.toDoubleList();
                z.toBoolList();
                z=
#ifdef USE_ATEN_LIB
at::ArrayRef<c10::optional<at::Tensor>> z>.toListOptionalTensor();
c10::List<c10::optional<at::Tensor>> z;
for (auto z: z) {
    z.push_back(zY);
}
#else
torch::executor::ArrayRef<torch::executor::optional<torch::executor::Tensor>> z0.toListOptionalTensor();
#endif
                Z_veczstd::vector<z> ú;z
    for (EValue z) {
        z	
        z);
    }
    ú(z);
                )r&   r>   r#   r   r   r   ZTensorr(   r;   r?   ÚintZSymIntÚfloatÚboolr   r)   Ú	connectorÚjoin)r+   r2   r3   r1   r4   r@   Z	elem_namer.   rA   Z	res_ctyperB   r/   Zvec_namer   r   r   r7   €   sä    

 ÿ
ÿÿÿýÿÿÿ
ÿÿÿýÿ
ÿÿÿýÿ
ÿÿÿýÿ

ÿþýþþýüþýüýþõÿÿÿÿþýý
ûûýùÿzUnboxing._gen_code_list_typeN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   r	   Ú__annotations__r   r   r   r   Ústrr0   r   rH   r   r&   r5   r   r6   r   r7   r   r   r   r   r      sB   
ÿÿþÿ  þ"  þ   þ   þr   N)Údataclassesr   Útypingr   r   r   r   Ztorchgen.api.typesr   r   r	   Ztorchgen.modelr
   r   r   r   r   r   r   rI   rP   r   r   r   r   r   r   Ú<module>   s   $
