Fix pack_padded_sequence because the PyTorch restricts the batch size tensor shoudl be cpu

Author: EikanWang

scripts/cpu/gen-dense-cpu-ops.py

@@ -74,11 +74,13 @@
     'aten::index_select(Tensor self, int dim, Tensor index) -> Tensor',
     'aten::_unsafe_view(Tensor self, int[] size) -> Tensor',
     'aten::native_layer_norm(Tensor input, Tensor? weight, Tensor? bias, int M, int N, float eps) -> (Tensor, Tensor, Tensor)',
-    'aten::native_layer_norm_backward(Tensor grad_out, Tensor input, Tensor mean, Tensor rstd, Tensor? weight, int M, int N, bool[3] output_mask) -> (Tensor, Tensor, Tensor)'
+    'aten::native_layer_norm_backward(Tensor grad_out, Tensor input, Tensor mean, Tensor rstd, Tensor? weight, int M, int N, bool[3] output_mask) -> (Tensor, Tensor, Tensor)',
+    'aten::_pack_padded_sequence(Tensor input, Tensor lengths, bool batch_first) -> (Tensor, Tensor)'
 ]
 
 _FN_IPEX_FUNCS_WITH_SIMPLE_ATEN_SIG = [
     'aten::index_select(Tensor self, int dim, Tensor index) -> Tensor',
+    'aten::_pack_padded_sequence(Tensor input, Tensor lengths, bool batch_first) -> (Tensor, Tensor)'
 ]
 
 _SHALLOW_FALLBACK_TO_CPU_TENSOR_LIST = 'shallowFallbackToCPUTensorList'
@@ -312,10 +314,17 @@ def is_out_func(fname):
         code += '  try {\n'
 
         code += '    if (check_auto_dnnl()) {\n'
-        code += '      std::vector<at::Tensor> dnnl_input_tensors;\n'
-        if len(dnnl_tensor_param_vars) > 0:
-            for dnnl_tensor_param_var in dnnl_tensor_param_vars:
-                code += '      dnnl_input_tensors.push_back({});\n'.format(dnnl_tensor_param_var)
+
+        if not self.is_ipex_func(aten_func_sig_str):
+            # There are two different kind of DevOPs in IPEX
+            #    1. DNNL Operator
+            #    2. CPU BF16/INT8 Operator in Vanilla PyTorch. IPEX itegrates this kind of operators in IPEX for
+            #       mixture precision.
+            # For the type 2, IPEX does not need to check if DNNL supports these tensors.
+            code += '      std::vector<at::Tensor> dnnl_input_tensors;\n'
+            if len(dnnl_tensor_param_vars) > 0:
+                for dnnl_tensor_param_var in dnnl_tensor_param_vars:
+                    code += '      dnnl_input_tensors.push_back({});\n'.format(dnnl_tensor_param_var)
 
         fname = cpp_sig.def_name
         if fname.endswith('_'):
@@ -328,19 +337,18 @@ def is_out_func(fname):
             for param_var in param_vars:
                 param_seq_str = param_var
                 param_seq_str_vec.append(param_seq_str)
-            code += '      if (dbl::chk::dnnl_support_the_tensors(dnnl_input_tensors)) {\n'
 
             if self.is_ipex_func(aten_func_sig_str):
-                code += '        auto _result = AtenIpexCPUDev::dil_{}({});\n'.format(fname, ', '.join(param_seq_str_vec))
-                code += '        if (is_ipex_func_success()) {\n'
-                code += '          return _result;\n'
-                code += '        } else {\n'
-                code += '          reset_ipex_func_status();\n'
-                code += '        }\n'
+                code += '      auto _result = AtenIpexCPUDev::dil_{}({});\n'.format(fname, ', '.join(param_seq_str_vec))
+                code += '      if (is_ipex_func_success()) {\n'
+                code += '        return _result;\n'
+                code += '      } else {\n'
+                code += '        reset_ipex_func_status();\n'
+                code += '      }\n'
             else:
+                code += '      if (dbl::chk::dnnl_support_the_tensors(dnnl_input_tensors)) {\n'
                 code += '        return AtenIpexCPUDev::dil_{}({});\n'.format(fname, ', '.join(param_seq_str_vec))
-
-            code += '      }\n' # Check support tensors
+                code += '      }\n' # Check support tensors
         code += '    }\n' # Check auto dnnl
         code += '  } catch (std::exception& e) {\n'
         code += '#if defined(_DEBUG)\n'