Verify types in custom op schemas (#124520)

zou3519 · pytorchmergebot · commit 5b98d43488be · 2024-04-23T14:18:35.000Z
Before this PR, we didn't check that types in a schema were valid. This is because TorchScript treats unknown types as type variables. This PR checks types in a schema for the TORCH_LIBRARY APIs. To do this, we add an `allow_typevars` flag to parseSchema so that TorchScript can use allow_typevars=True. We also add some error messages for common mistakes (e.g. using int64_t or double in schema). Test Plan: - new tests Differential Revision: [D56432690](https://our.internmc.facebook.com/intern/diff/D56432690) Pull Request resolved: #124520 Approved by: https://github.com/albanD
diff --git a/test/test_custom_ops.py b/test/test_custom_ops.py
@@ -1740,6 +1740,17 @@ def test_is_functional_schema(self):
             res = torch._library.utils.is_functional_schema(schema)
             self.assertEqual(res, expected)
 
+    def test_incorrect_schema_types(self):
+        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
+            with self.assertRaisesRegex(RuntimeError, "unknown type specifier"):
+                lib.define("foo12(Tensor a) -> asdfasdf")
+            with self.assertRaisesRegex(RuntimeError, "unknown type specifier"):
+                lib.define("foo12(asdf a) -> Tensor")
+            with self.assertRaisesRegex(RuntimeError, "Use `SymInt` or `int`"):
+                lib.define("foo12(int64_t a) -> Tensor")
+            with self.assertRaisesRegex(RuntimeError, "Use `float`"):
+                lib.define("foo12(double a) -> Tensor")
+
     def test_is_tensorlist_like_type(self):
         tensorlists = [
             # Tensor[]
diff --git a/torch/csrc/jit/frontend/function_schema_parser.cpp b/torch/csrc/jit/frontend/function_schema_parser.cpp
@@ -23,14 +23,14 @@ namespace torch::jit {
 
 namespace {
 struct SchemaParser {
-  explicit SchemaParser(const std::string& str)
+  explicit SchemaParser(const std::string& str, bool allow_typevars)
       : L(std::make_shared<Source>(
             c10::string_view(str),
             c10::nullopt,
             0,
             nullptr,
             Source::DONT_COPY)),
-        type_parser(L, /*parse_complete_tensor_types*/ false) {}
+        type_parser(L, /*parse_complete_tensor_types*/ false, allow_typevars) {}
 
   std::variant<OperatorName, FunctionSchema> parseDeclaration() {
     OperatorName name = parseName();
@@ -361,16 +361,19 @@ struct SchemaParser {
   }
   Lexer L;
   SchemaTypeParser type_parser;
+  bool allow_typevars_;
 };
 } // namespace
 
 std::variant<OperatorName, FunctionSchema> parseSchemaOrName(
-    const std::string& schemaOrName) {
-  return SchemaParser(schemaOrName).parseExactlyOneDeclaration();
+    const std::string& schemaOrName,
+    bool allow_typevars) {
+  return SchemaParser(schemaOrName, allow_typevars)
+      .parseExactlyOneDeclaration();
 }
 
-FunctionSchema parseSchema(const std::string& schema) {
-  auto parsed = parseSchemaOrName(schema);
+FunctionSchema parseSchema(const std::string& schema, bool allow_typevars) {
+  auto parsed = parseSchemaOrName(schema, allow_typevars);
   TORCH_CHECK(
       std::holds_alternative<FunctionSchema>(parsed),
       "Tried to parse a function schema but only the operator name was given");
diff --git a/torch/csrc/jit/frontend/function_schema_parser.h b/torch/csrc/jit/frontend/function_schema_parser.h
@@ -8,9 +8,15 @@
 namespace torch {
 namespace jit {
 
+// allow_typevars: If true, we assume that lowercase types that we don't
+// understand are type variables. This is only needed for TorchScript (and not
+// not needed for custom ops).
 TORCH_API std::variant<c10::OperatorName, c10::FunctionSchema> parseSchemaOrName(
-    const std::string& schemaOrName);
-TORCH_API c10::FunctionSchema parseSchema(const std::string& schema);
+    const std::string& schemaOrName,
+    bool allow_typevars = true);
+TORCH_API c10::FunctionSchema parseSchema(
+    const std::string& schema,
+    bool allow_typevars = true);
 TORCH_API c10::OperatorName parseName(const std::string& name);
 
 } // namespace jit
diff --git a/torch/csrc/jit/frontend/schema_type_parser.cpp b/torch/csrc/jit/frontend/schema_type_parser.cpp
@@ -82,12 +82,27 @@ TypePtr SchemaTypeParser::parseBaseType() {
 
   auto it = type_map.find(text);
   if (it == type_map.end()) {
-    if (!text.empty() && islower(text[0])) {
+    if (allow_typevars_ && !text.empty() && islower(text[0])) {
       // lower case identifiers that are not otherwise valid types
       // are treated as type variables
       return c10::TypeFactory::createNamed<VarType>(text);
     }
-    throw ErrorReport(tok.range) << "unknown type specifier";
+    if (text == "double") {
+      throw ErrorReport(tok.range)
+          << "Use `float` instead of `double` in an operator's schema string. "
+             "`float` in schema corresponds to the double type in C++";
+    }
+    if (text == "int64_t") {
+      throw ErrorReport(tok.range)
+          << "Use `SymInt` or `int` instead of `int64_t` in an operator's schema string. "
+             "`SymInt` corresponds to c10::SymInt in C++ while `int` in schema corresponds "
+             "to the int64_t type in C++.";
+    }
+    throw ErrorReport(tok.range)
+        << "unknown type specifier. Common valid schema types include "
+           "Tensor, SymInt, int, float, bool, Scalar; "
+           "for a full list, please see "
+           "https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#func ";
   }
   return it->second;
 }
diff --git a/torch/csrc/jit/frontend/schema_type_parser.h b/torch/csrc/jit/frontend/schema_type_parser.h
@@ -20,8 +20,13 @@ struct TORCH_API SchemaTypeParser {
   c10::optional<at::ScalarType> parseTensorDType(const std::string& dtype);
   TypePtr parseRefinedTensor();
 
-  SchemaTypeParser(Lexer& L, bool parse_complete_tensor_types)
-      : complete_tensor_types(parse_complete_tensor_types), L(L) {}
+  SchemaTypeParser(
+      Lexer& L,
+      bool parse_complete_tensor_types,
+      bool allow_typevars)
+      : complete_tensor_types(parse_complete_tensor_types),
+        L(L),
+        allow_typevars_(allow_typevars) {}
 
  private:
   c10::optional<bool> tryToParseRequiresGrad();
@@ -35,6 +40,7 @@ struct TORCH_API SchemaTypeParser {
   bool complete_tensor_types;
   Lexer& L;
   size_t next_id = 0;
+  bool allow_typevars_;
 };
 } // namespace jit
 } // namespace torch
diff --git a/torch/csrc/jit/ir/irparser.cpp b/torch/csrc/jit/ir/irparser.cpp
@@ -35,7 +35,10 @@ class IRParser {
       : L(std::make_shared<Source>(str)),
         g(graph),
         vmap(vmap),
-        type_parser(L, /*parse_complete_tensor_types*/ true),
+        type_parser(
+            L,
+            /*parse_complete_tensor_types*/ true,
+            /*allow_type_vars*/ true),
         parse_tensor_constants_(parse_tensor_constants) {}
 
   std::string parseVar();
diff --git a/torch/csrc/jit/python/init.cpp b/torch/csrc/jit/python/init.cpp
@@ -1765,7 +1765,11 @@ void initJITBindings(PyObject* module) {
       },
       py::arg("input"),
       py::arg("parse_tensor_constants") = false);
-  m.def("parse_schema", parseSchema);
+  m.def(
+      "parse_schema",
+      &parseSchema,
+      py::arg("schema"),
+      py::arg("allow_typevars") = true);
   m.def("unify_type_list", [](const std::vector<TypePtr>& types) {
     std::ostringstream s;
     auto type = unifyTypeList(types, s);
diff --git a/torch/csrc/jit/runtime/static/passes.cpp b/torch/csrc/jit/runtime/static/passes.cpp
@@ -1347,7 +1347,8 @@ bool isNoOpSlice(Node* node) {
 void EliminateNoOpSlice(std::shared_ptr<Graph>& graph) {
   DepthFirstGraphNodeIterator it(graph);
   auto schema = torch::schema(
-      "aten::slice.t(t[] l, int? start=None, int? end=None, int step=1) -> t[]");
+      "aten::slice.t(t[] l, int? start=None, int? end=None, int step=1) -> t[]",
+      /*allow_typevars*/ true);
   Node* node = nullptr;
   std::vector<Node*> to_delete;
   while ((node = it.next()) != nullptr) {
diff --git a/torch/library.h b/torch/library.h
@@ -406,17 +406,17 @@ inline CppFunction dispatch(c10::DeviceType type, Func&& raw_f) {
 /// ```
 ///
 /// \ingroup torch-schema-overloads
-inline c10::FunctionSchema schema(const char* str, c10::AliasAnalysisKind k) {
-  c10::FunctionSchema s = torch::jit::parseSchema(str);
+inline c10::FunctionSchema schema(const char* str, c10::AliasAnalysisKind k, bool allow_typevars=false) {
+  c10::FunctionSchema s = torch::jit::parseSchema(str, /*allow_typevars*/allow_typevars);
   s.setAliasAnalysis(k);
   return s;
 }
 
 /// Function schemas can be directly constructed from string literals.
 ///
 /// \ingroup torch-schema-overloads
-inline c10::FunctionSchema schema(const char* s) {
-  return schema(s, c10::AliasAnalysisKind::FROM_SCHEMA);
+inline c10::FunctionSchema schema(const char* s, bool allow_typevars=false) {
+  return schema(s, c10::AliasAnalysisKind::FROM_SCHEMA, allow_typevars);
 }
 
 /// \private
