Perf nitpicks on python_arg_parser's is_int_or_symint_list

swolchok · swolchok · commit 0ec9031b0f7a · 2025-09-02T12:46:25.000-07:00
This function has come up in DTensor perf work, and I had a nitpick on #160256 so here it is. I have neither compiled nor measured this, but am reasonably confident it's better nonetheless. [ghstack-poisoned]
diff --git a/torch/csrc/utils/python_arg_parser.cpp b/torch/csrc/utils/python_arg_parser.cpp
@@ -962,15 +962,15 @@ static bool is_int_or_symint_list(
     int broadcast_size,
     int64_t* failed_idx = nullptr,
     std::vector<PyObject*>* overloaded_args = nullptr) {
-  if (PyTuple_Check(obj) || PyList_Check(obj)) {
-    if (PySequence_Size(obj) == 0) {
+  const bool is_tuple = PyTuple_Check(obj);
+  if (is_tuple || PyList_Check(obj)) {
+    const auto size = is_tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj);
+    if (size == 0) {
       return true;
     }
 
     // Check all elements, not just the first one, when looking for torch
     // functions
-    const bool is_tuple = PyTuple_Check(obj);
-    const auto size = is_tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj);
     bool has_torch_func = false;
 
     for (Py_ssize_t idx = 0; idx < size; idx++) {
