3j++SSKJr SSKJrJr SSKJrJr SSKrSSK rSSK J r SSK J r SSK JrJrJr \(a SS KJr SSKr"S S \5r"S S \5rg)) annotations)ABCabstractmethod)Any TYPE_CHECKINGN)ir)V) FixedLayoutFlexibleLayoutLayout)Sequencec\rSrSrSrSSSjjrSSSjjr\SSj5r\SSj5r SSjr SS jr SS jr SS jr SS jrSS jrSSjrSS Sjjr\S!Sj5rS"Sjr\S#S$Sjj5rSrg)% KernelInputsz Class to store and provide access to input nodes for kernels. This class takes in a tuple of input nodes and provides methods to access information about these nodes, such as their device type and device. NcnXlSUlUbUO0UlX0l[ U5S:dS5eg)z Initialize with a tuple of input nodes. Args: input_nodes: A tuple of input nodes to store out_dtype: Optional output dtype to store Nrz Expected at least one input node) _input_nodes _device_name_scalars _out_dtypelen)self input_nodesscalars out_dtypes W/home/wildlama/miniconda3/lib/python3.13/site-packages/torch/_inductor/kernel_inputs.py__init__KernelInputs.__init__s>((,#*#6B #;!#G%GG#cUc UR$[UR5[U5:Xd)S[UR5S[U535eUVs/sHo RUPM sn$s snf)z Return the stored input nodes, optionally reordered. Args: reorder: Optional sequence of indices to reorder the nodes. For example, (2, 0, 1) would return nodes in that order. Returns: The tuple of input nodes, optionally reordered zReorder length mismatch:  vs )rr)rreorderis rnodesKernelInputs.nodes.s| ?$$ $4$$%W5 'D,=,=(>'?tCL> R 5/66g!!!$g666sA:c,[UR5$)zH Get the number of input nodes. Returns: The number of input nodes )rrrs rcountKernelInputs.count@s4$$%%rcH[R"URS5$)z\ Get the device type of the first node. Returns: The device type (e.g., 'cuda', 'cpu') r)rget_device_typerr's r device_typeKernelInputs.device_typeJs !!$"3"3A"677rc<URSR5$)zN Get the device of the first node. Returns: The device of the first node r)r get_devicer's rdeviceKernelInputs.deviceUs  #..00rcURcPUR5nURS:Xa0[RR U5nUR UlUR$)zU Get the device name information. Returns: A tuple of (gpu_name, vendor, model) cuda)rr0r,torchr3get_device_properties gcnArchName)rr0device_propertiess r device_nameKernelInputs.device_name^sX    $[[]F6)$)JJ$D$DV$L!$5$A$A!   rc:[SUR55$)zg Get the symbolic shapes of all input nodes. Returns: A tuple of shape tuples for each input node c3@# UHoR5v M g7fN)get_size.0nodes r /KernelInputs.shapes_symbolic..ssC1B]]__1Btuplerr's rshapes_symbolicKernelInputs.shapes_symboliclsC1B1BCCCrc:[SUR55$)z Get the size hints for shapes of all input nodes. Returns: A tuple of shape tuples with integer hints for each input node c3# UH<n[RRRUR 55v M> g7fr<)r graphsizevarsoptimization_hintsr=r>s rrA-KernelInputs.shapes_hinted..|s5 ) GG   / /  @ @)AArDr's r shapes_hintedKernelInputs.shapes_hintedu$ ))   rc:[SUR55$)zi Get the symbolic strides of all input nodes. Returns: A tuple of stride tuples for each input node c3@# UHoR5v M g7fr<) get_strider>s rrA0KernelInputs.strides_symbolic..sE3D4__&&3DrCrDr's rstrides_symbolicKernelInputs.strides_symbolicsE43D3DEEErc:[SUR55$)z Get the size hints for strides of all input nodes. Returns: A tuple of stride tuples with integer hints for each input node c3# UH<n[RRRUR 55v M> g7fr<)r rJrKrLrTr>s rrA.KernelInputs.strides_hinted..s6 ) GG   / /0A B B)rNrDr's rstrides_hintedKernelInputs.strides_hintedrQrc:[SUR55$)zX Get the dtypes of all input nodes. Returns: A tuple of dtypes for each input node c3@# UHoR5v M g7fr<) get_dtyper>s rrA&KernelInputs.dtypes..sD2C$^^%%2CrCrDr's rdtypesKernelInputs.dtypessD$2C2CDDDrc<URUR5$)z Get the dtype of a specific input node. Args: idx: Index of the node to get the dtype from (default: 0) Returns: The dtype of the specified input node )rr_)ridxs rdtypeKernelInputs.dtypes  %//11rcg)c Get the output dtype, whether passed in or inferred from the nodes Returns: The output dtype Nr's rrKernelInputs.out_dtypercTXR;d SUS35eURU$)zr Get the scalar value for a given name. Args: name: Name of the scalar to get Returns: The scalar value zScalar z not found, but required)r)rnames r get_scalarKernelInputs.get_scalars2}}$Nv5M&NN$}}T""rcg)z Abstract method to handle output layout generation. Args: out_dtype: Optional output dtype. If not provided, infer from inputs flexible: If True, return FlexibleLayout, otherwise FixedLayout Nri)rflexibles r output_layoutKernelInputs.output_layoutrkr)rrrr)NN)r list[Any]rdict[str, float | int] | Nonertorch.dtype | Noner<)r"zSequence[int] | Nonereturnrtrwint)rwz str | None)rwz torch.device)rwztuple[tuple[Any, ...], ...])rwztuple[tuple[int, ...], ...])rwz%tuple[tuple[sympy.Integer, ...], ...])rwztuple[torch.dtype, ...])r)rdryrw torch.dtyperwrz)rmstrrwz float | intTrqboolrwr )__name__ __module__ __qualname____firstlineno____doc__rr$propertyr(r,r0r8rFrOrVr[rarerrrnrr__static_attributes__rirrrrs26(, HH/H& H&7$&&881 !D  F  E 2   #  rrc^\rSrSrSrS S U4SjjjrS SjrSSjrSSSjjrSSjr SSjr SS jr S r U=r $)MMKernelInputszz Specialized KernelInputs for matrix multiplication operations. Provides additional methods to access M, N, K dimensions. cZ>[TU]XU5 [UR5S:dS5eXEpvUS:aU[U5- nUS:aU[U5- nSUs=::a[U5:d O SU35eSUs=::a[U5:d O SU35eX@lXPlg)z Initialize with a tuple of input nodes. By default, we assume the last 2 input nodes are mat1 and mat2, but the caller can adjust when necessary zExpected at least 2 input nodesrzInvalid mat1_idx: zInvalid mat2_idx: N)superrrr _mat1_idx _mat2_idx) rrrrmat1_idxmat2_idxm1_idxm2_idx __class__s rrMMKernelInputs.__init__s y94$$%*M,MM*" a< c+& &F a< c+& &FF-S--N1CH:/NN-F-S--N1CH:/NN-!!rchUR5URnUR5URnUR5SnUR5SnUR5SnUR5Sn[R R RXF5 X5U4$)aq Get the symbolic M, N, K dimensions for matrix multiplication. Handles both 2D (MM) and 3D (BMM) tensors. M is extracted from the second-to-last dimension of the first operand (mat1). N is extracted from the last dimension of the second operand (mat2). K is extracted from the last dimension of the first operand (mat1). Returns: A tuple of (M, N, K) dimensions )r$rrr=r rJrK check_equals)rmat1mat2mknk0s r mnk_symbolicMMKernelInputs.mnk_symbolicszz|DNN+zz|DNN+ MMOB  MMOB  MMOB ]]_R  %%a,ayrcvURb UR$UR5SR5$)rhr)rmat1mat2r_r's rrMMKernelInputs.out_dtype s2 ?? &?? "}}q!++--rcUR5up#UR5nUR5GtpVnUR5Gtpn [XX5V V s/sH.up[R R RX5PM0 sn n n /U QUPU Pn U(a[UR5XM5$[UR5XM5$s sn n f)z Handle output layout generation for matrix multiplication. Args: out_dtype: Optional output dtype. If not provided, infer from inputs flexible: If True, return FlexibleLayout, otherwise FixedLayout ) rrr=zipr rJrKcheck_equals_and_simplifyr r0r )rrqrrrb1rk1b2k2rabsizes rrrMMKernelInputs.output_layouts]]_ NN$ ]]_ ]]_ JMb+ V+$!QWW   7 7 =+ VzAzqz !$++-A At{{}i> > Ws5CcZUR5nXRXR4$)zJ Get the mat1 and mat2 nodes. Returns: A tuple of (mat1, mat2) nodes )r$rr)rr$s rrMMKernelInputs.mat1mat2,s( ^^$eNN&;;;rcUR5nXRnXRnUSnUSnUSnUSnXW:Xd SUSU35eXFU4$)z Get the hinted M, N, K dimensions for matrix multiplication. Handles both 2D (MM) and 3D (BMM) tensors. Uses shapes_hinted from the base class to get integer hints for dimensions. Returns: A tuple of (M, N, K) dimensions as integers rrzK dimensions don't match: r!)rOrr)r hinted_shapes mat1_shape mat2_shaperrrk_checks r mnk_hintedMMKernelInputs.mnk_hinted6sw**, ">>2 ">>2 rN rN rNR.|J9!D JJ|ayrchUR5nXRn[U5S:aUS$g)z Get the hinted batch size for batched matrix multiplication. Returns 1 for non-batched (2D) operations. Returns: The batch size as an integer r )rOrr)rrrs r batch_hintedMMKernelInputs.batch_hintedNs7**, ">>2 z?a b> !r)rr)NNrr) rrtrrurrvrryrry)rwz2tuple[sympy.Integer, sympy.Integer, sympy.Integer]r{r}r~)rwztuple[Any, Any])rwztuple[int, int, int]rx)rrrrrrrrrrrrrr __classcell__)rs@rrrs|26(, ""/"& "  "  ""B ;4 .?(<0  rr) __future__rabcrrtypingrrr4torch._inductor.configtorch._inductorrtorch._inductor.virtualizedr r r r collections.abcrsympyrrrirrrsH"#% )33(u 3u pO\Or