codegen: Allow string arguments to have defaults (#45665)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/45665

Fixes #43944

Note that the codegen doesn't use a proper parser so, in the same way as with lists, the string `, ` cannot appear in defaults or it will be interpreted as a splitting point between arguments.

Test Plan: Imported from OSS

Reviewed By: albanD

Differential Revision: D24141835

Pulled By: ezyang

fbshipit-source-id: 578127861fd2504917f4486c44100491a2c40343

aten/src/ATen/native/TestOps.cpp

@@ -42,5 +42,13 @@ Tensor _test_optional_floatlist(
   return output;
 }
 
+// Test default strings can handle escape sequences properly (although commas are broken)
+Tensor _test_string_default(const Tensor& dummy, std::string a, std::string b) {
+  const c10::string_view expect = "\"'\\";
+  TORCH_CHECK(a == expect, "Default A failed");
+  TORCH_CHECK(b == expect, "Default B failed");
+  return dummy;
+}
+
 } // namespace native
 } // namespace at

aten/src/ATen/native/native_functions.yaml

@@ -8308,3 +8308,8 @@
   python_module: nn
   dispatch:
     CPU: _test_optional_floatlist
+
+# Note: this function is only for testing.
+- func: _test_string_default(Tensor dummy, str a="\"'\\", str b='"\'\\') -> Tensor
+  use_c10_dispatcher: full
+  python_module: nn

test/test_native_functions.py

@@ -176,6 +176,22 @@ class TestNativeFunctions(TestCase):
 
         self.do_test_optional_filled_intlist_with_module(fake_module)
 
+    def test_string_defaults(self):
+        dummy = torch.rand(1)
+        fn = torch._C._nn._test_string_default
+        fn(dummy)
+
+        with self.assertRaisesRegex(RuntimeError, "A"):
+            fn(dummy, a="")
+
+        with self.assertRaisesRegex(RuntimeError, "B"):
+            fn(dummy, b="")
+
+        def f(x):
+            torch._C._nn._test_string_default(x)
+        scripted_fn = torch.jit.script(f)
+        scripted_fn(dummy)
+
 
 if __name__ == '__main__':
     run_tests()

tools/autograd/gen_python_functions.py

@@ -36,6 +36,7 @@ from .gen_variable_type import should_trace
 from .utils import write, is_tensor_method
 
 from tools.codegen.code_template import CodeTemplate
+from tools.codegen.gen import cpp_string
 
 #
 # declarations blocklist
@@ -964,7 +965,7 @@ def method_impl(name, declarations, is_python_method, module):
     dispatch = []
     for i, dictionary in enumerate(grouped):
         signature = dictionary['signature']
-        signatures.append('"{}",'.format(signature))
+        signatures.append(f'{cpp_string(str(signature))},')
         overload_index = i if not is_singleton else None
         dispatch.append(emit_dispatch_case(overload_index, dictionary, is_python_method))
 

tools/codegen/api/cpp.py

@@ -161,6 +161,26 @@ JIT_TO_CPP_DEFAULT = {
 def default_expr(d: str, t: Type) -> str:
     if d == 'None' and str(t) == 'Tensor?':
         return '{}'
+    if isinstance(t, BaseType) and t.name is BaseTy.str:
+        # Schema allows single quotes but C++ needs double
+        if len(d) >= 2 and d[0] == "'" and d[-1] == "'":
+            s = ''
+            i = 1
+            while i + 1 < len(d):
+                if d[i] != '\\':
+                    if d[i] == '"':
+                        s += '\\"'
+                    else:
+                        s += d[i]
+                    i += 1
+                else:
+                    if d[i + 1] == "'":
+                        s += "'"
+                    else:
+                        s += d[i:i + 2]
+                    i += 2
+
+            return f'"{s}"'
     return JIT_TO_CPP_DEFAULT.get(d, d)
 
 # Convert an argument into its C++ API form

tools/codegen/gen.py

@@ -9,6 +9,7 @@ from collections import OrderedDict
 import argparse
 import pathlib
 import functools
+import json
 
 from tools.codegen.code_template import CodeTemplate
 from tools.codegen.model import *
@@ -124,6 +125,18 @@ def concatMap(func: Callable[[T], Sequence[S]], xs: Sequence[T]) -> Iterator[S]:
         for r in func(x):
             yield r
 
+def cpp_string(s: str) -> str:
+    """Convert a python string into a c++ string literal """
+    s = s.replace('\\', '\\\\')
+    s = s.replace('"', '\\"')
+    s = s.replace('\a', '\\a')
+    s = s.replace('\b', '\\b')
+    s = s.replace('\f', '\\f')
+    s = s.replace('\n', '\\n')
+    s = s.replace('\v', '\\v')
+    s = s.replace('\t', '\\t')
+    return f'"{s}"'
+
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
 #
 #                        C++ CODE GENERATION
@@ -268,7 +281,7 @@ def compute_type_method(
             # def registration only happens in TypeDefault
             def_registration = ""
             if dispatch is None:
-                def_registration = f'm.def("{f.func}");\n'
+                def_registration = f'm.def({cpp_string(str(f.func))});\n'
 
             impl_registration = ""
             if not def_only and not f.manual_kernel_registration and (dispatch is not None or f.dispatch is None):
@@ -881,9 +894,12 @@ def compute_registration_declarations(f: NativeFunction) -> str:
     returns_type = dispatcher.returns_type(f.func.returns)
     args = dispatcher.arguments(f.func)
     args_str = ', '.join(map(str, args))
-    dispatch = f.dispatch is not None
+    comment_data : Dict[str, str] = {
-    math = dispatch and 'Math' in f.dispatch  # type: ignore
+        'schema': f'aten::{f.func}',
-    return f"""{returns_type} {name}({args_str}); // {{"schema": "aten::{f.func}", "dispatch": "{dispatch}", "math": "{math}"}}
+        'dispatch': str(f.dispatch is not None),
+        'math': str(f.dispatch is not None and 'Math' in f.dispatch)
+    }
+    return f"""{returns_type} {name}({args_str}); // {json.dumps(comment_data)}
 """
 
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

torch/csrc/utils/python_arg_parser.cpp

@@ -503,6 +503,62 @@ static inline std::vector<int64_t> parse_intlist_args(const std::string& s, int6
   return args;
 }
 
+// Parse a string literal to remove quotes and escape sequences
+static std::string parse_string_literal(c10::string_view str) {
+  TORCH_CHECK(str.length() >= 2, "String defaults must be quoted");
+
+  if (str.front() == '"') {
+    TORCH_CHECK(str.back() == '"',
+                "Mismatched quotes in string default: ", str);
+  } else {
+    TORCH_CHECK(str.front() == '\'' && str.back() == '\'',
+                "Invalid quotes in string default: ", str)
+  }
+
+  std::string parsed;
+  parsed.reserve(str.size());
+  for (size_t i = 1; i < str.size() - 1;) {
+    if (str[i] != '\\') {
+      parsed.push_back(str[i]);
+      ++i;
+      continue;
+    }
+
+    // Handle escape sequences
+    TORCH_CHECK(i < str.size() - 2, "String ends with escaped final quote: ", str)
+    char c = str[i + 1];
+    switch (c) {
+      case '\\':
+      case '\'':
+      case '\"':
+        break;
+      case 'a':
+        c = '\a';
+        break;
+      case 'b':
+        c = '\b';
+        break;
+      case 'f':
+        c = '\f';
+        break;
+      case 'n':
+        c = '\n';
+        break;
+      case 'v':
+        c = '\v';
+        break;
+      case 't':
+        c = '\t';
+        break;
+      default:
+        TORCH_CHECK(false, "Unsupported escape sequence in string default: \\", str[i + 1]);
+    }
+    parsed.push_back(c);
+    i += 2;
+  }
+  return parsed;
+}
+
 void FunctionParameter::set_default_str(const std::string& str) {
   if (str == "None") {
     allow_none = true;
@@ -558,8 +614,8 @@ void FunctionParameter::set_default_str(const std::string& str) {
       throw std::runtime_error("invalid device: " + str);
     }
   } else if (type_ == ParameterType::STRING) {
-    if (str != "None" && str != "") {
+    if (str != "None") {
-      throw std::runtime_error("invalid default string: " + str);
+      default_string = parse_string_literal(str);
     }
   }
 }

torch/csrc/utils/python_arg_parser.h

@@ -188,6 +188,7 @@ struct PythonArgs {
   inline c10::optional<at::MemoryFormat> memoryformatOptional(int i);
   inline at::QScheme toQScheme(int i);
   inline std::string string(int i);
+  inline std::string stringWithDefault(int i, const std::string& default_str);
   inline c10::optional<std::string> stringOptional(int i);
   inline PyObject* pyobject(int i);
   inline int64_t toInt64(int i);
@@ -226,6 +227,7 @@ struct FunctionParameter {
   at::SmallVector<PyObject *, 5> numpy_python_names;
   at::Scalar default_scalar;
   std::vector<int64_t> default_intlist;
+  std::string default_string;
   union {
     bool default_bool;
     int64_t default_int;
@@ -530,7 +532,11 @@ inline at::QScheme PythonArgs::toQScheme(int i) {
 }
 
 inline std::string PythonArgs::string(int i) {
-  if (!args[i]) return "";
+  return stringWithDefault(i, signature.params[i].default_string);
+}
+
+inline std::string PythonArgs::stringWithDefault(int i, const std::string& default_str) {
+  if (!args[i]) return default_str;
   return THPUtils_unpackString(args[i]);
 }