1212 AT_INT_ARRAY_REF ,
1313 AT_SCALAR ,
1414 AT_TENSOR ,
15+ AT_TENSOR_LIST ,
1516 BOOL ,
1617 CppTestFileGen ,
1718 DOUBLE ,
2829 THREE_TENSOR_TUPLE ,
2930 TWO_TENSOR_TUPLE ,
3031)
32+
3133from torchgen .api import cpp
3234from torchgen .api .types import CppSignatureGroup
3335
@@ -75,6 +77,8 @@ class ValueRef:
7577
7678ValueRefList = Union [ValueRef , List [ValueRef ]]
7779
80+ InableCppType = frozenset ([AT_TENSOR , AT_TENSOR_LIST ])
81+
7882
7983class ComputeGraphGen :
8084 def __init__ (self , op_reg_name : str , f : NativeFunction , suite_def : TestSuite ):
@@ -114,7 +118,7 @@ def __init__(self, op_reg_name: str, f: NativeFunction, suite_def: TestSuite):
114118 name = f"{ arg .name } ,
115119 src_cpp_name = arg .name ,
116120 src_cpp_type = cpp_type ,
117- is_in = (cpp_type == AT_TENSOR ),
121+ is_in = (cpp_type in InableCppType ),
118122 requires_prepack = requires_prepack ,
119123 supports_prepack = supports_prepack ,
120124 )
@@ -244,6 +248,25 @@ def create_value_for(self, ref: ValueRefList) -> str: # noqa: C901
244248 ret_str += f"{ self .graph } { self .dot }
245249 ret_str += f"({ ref .src_cpp_name } \n "
246250 return ret_str
251+ elif ref .src_cpp_type == AT_TENSOR_LIST :
252+ assert ref .is_in , "AT_TENSOR_LIST must be an input"
253+ # This logic is a bit convoluted. We need to create a IOValueRef for
254+ # each tensor, to facilate staging. On the other hand, we will
255+ # use the .value tensor to create a ValueList, which will be passed
256+ # to the corresponding ops.
257+ ret_str = f"std::vector<IOValueRef> { ref .name } \n "
258+ ret_str += f"std::vector<ValueRef> { ref .name } \n "
259+ ret_str += f"for (int i=0; i < { ref .src_cpp_name } \n "
260+ ret_str += f" { cpp_type } { self .graph } { self .dot } \n "
261+ ret_str += f" { ref .src_cpp_name } \n "
262+ ret_str += (
263+ f" from_at_scalartype({ ref .src_cpp_name } \n "
264+ )
265+ ret_str += f" { ref .name } \n "
266+ ret_str += f" { ref .name } \n "
267+ ret_str += "}\n "
268+ ret_str += f"ValueRef { ref .name } { self .graph } { self .dot } { ref .name } \n "
269+ return ret_str
247270
248271 ret_str = f"{ cpp_type } { ref .name } { self .graph } { self .dot }
249272 if ref .src_cpp_type == AT_TENSOR and not prepack :
@@ -288,11 +311,16 @@ def create_op_call(self) -> str:
288311
289312 for aten_arg in self .args :
290313 ref = self .refs [aten_arg .name ]
291- op_create_code += (
292- f"{ ref .name }
293- if (ref .is_in and not self .prepack_ref (ref )) or ref .is_out
294- else f"{ ref .name }
295- )
314+ if ref .src_cpp_type == AT_TENSOR_LIST :
315+ # Special case. Underlying tensors are input tensors, but the
316+ # container itself is just a normal value.
317+ op_create_code += f"{ ref .name }
318+ else :
319+ op_create_code += (
320+ f"{ ref .name }
321+ if (ref .is_in and not self .prepack_ref (ref )) or ref .is_out
322+ else f"{ ref .name }
323+ )
296324
297325 op_create_code += "out_ref});\n "
298326 return op_create_code
@@ -311,22 +339,46 @@ def set_output(self, ref: ValueRefList) -> str:
311339
312340 def virtual_resize (self , ref : ValueRefList ) -> str :
313341 assert isinstance (ref , ValueRef )
314- assert ref .src_cpp_type == AT_TENSOR and ref .is_in
342+ assert ref .src_cpp_type in InableCppType and ref .is_in
315343 if self .prepack_ref (ref ):
316344 return ""
317- ret_str = f"{ self .graph } { self .dot } { ref .name }
318- ret_str += f"->virtual_resize({ ref .src_cpp_name } \n "
345+
346+ if ref .src_cpp_type == AT_TENSOR :
347+ ret_str = f"{ self .graph } { self .dot } { ref .name }
348+ ret_str += f"->virtual_resize({ ref .src_cpp_name } \n "
349+ elif ref .src_cpp_type == AT_TENSOR_LIST :
350+ ret_str = ""
351+ ret_str += f"for (int i=0; i < { ref .name } \n "
352+ ret_str += f" { self .graph } { self .dot } { ref .name }
353+ ret_str += f"->virtual_resize({ ref .src_cpp_name } \n "
354+ ret_str += "}\n "
355+ else :
356+ raise AssertionError (f"{ ref .src_cpp_type } )
357+
319358 return ret_str
320359
321360 def copy_into_staging (self , ref : ValueRefList ) -> str :
322361 assert isinstance (ref , ValueRef )
323- assert ref .src_cpp_type == AT_TENSOR and ref .is_in
362+ assert ref .src_cpp_type in InableCppType and ref .is_in
363+
324364 if self .prepack_ref (ref ):
325365 return ""
326- ret_str = f"{ self .graph } { self .dot }
327- ret_str += f"{ ref .name }
328- ret_str += f"{ ref .src_cpp_name }
329- ret_str += f"{ ref .src_cpp_name } \n "
366+
367+ if ref .src_cpp_type == AT_TENSOR :
368+ ret_str = f"{ self .graph } { self .dot }
369+ ret_str += f"{ ref .name }
370+ ret_str += f"{ ref .src_cpp_name }
371+ ret_str += f"{ ref .src_cpp_name } \n "
372+ elif ref .src_cpp_type == AT_TENSOR_LIST :
373+ ret_str = ""
374+ ret_str += f"for (int i=0; i < { ref .name } \n "
375+ ret_str += f" { self .graph } { self .dot }
376+ ret_str += f"{ ref .name }
377+ ret_str += f"{ ref .src_cpp_name }
378+ ret_str += f"{ ref .src_cpp_name } \n "
379+ ret_str += "}\n "
380+ else :
381+ raise AssertionError (f"{ ref .src_cpp_type } )
330382 return ret_str
331383
332384 def declare_vk_out_for (self , ref : Union [ValueRef , List [ValueRef ]]) -> str :
@@ -547,8 +599,10 @@ def gen_parameterization(self) -> str:
547599 if (!is_close && t1.numel() < 500) {
548600 std::cout << "reference: " << std::endl;
549601 print(t1, 150);
602+ std::cout << std::endl;
550603 std::cout << "vulkan: " << std::endl;
551604 print(t2, 150);
605+ std::cout << std::endl;
552606 }
553607 return is_close;
554608}
0 commit comments