[JIT] Plumb type annotations through script compilation #9405

jamesr66a · 2018-07-13T00:28:00Z

No description provided.

facebook-github-bot

@jamesr66a has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

zdevito

Cool -- I have a bunch of organization comments to keep the complexity of this change down. Right now it adds too many different pathways to the compiler and I have some suggestions about how to reduce that.

torch/csrc/jit/script/compiler.h

+  DefAndTypes(Def def, std::vector<TypePtr> arg_types, TypePtr return_type)
+    : def(std::move(def)), arg_types(arg_types), return_type(return_type) {}
+  Def def;
+  std::vector<TypePtr> arg_types;


torch/csrc/jit/script/compiler.h

  bool allow_varargs;
 };

+struct DefAndTypes {


torch/csrc/jit/script/compiler.h

  const std::vector<Resolver>& resolvers, /* determines how we handle free variables in each definition*/
-  std::shared_ptr<SugaredValue> self /* if non-null, the first argument to each def, is bound to this value */
+  std::shared_ptr<SugaredValue> self, /* if non-null, the first argument to each def, is bound to this value */
+  bool pure_func=false /* If true, expect a single def which will become the 'forward' method on the module */


torch/jit/__init__.py

+    mod = ScriptModule()
+    rcb = createResolutionCallback(_frames_up + 1)
+    ast = get_jit_ast(fn)
+    arg_types, ret_type = annotations.get_signature(fn, ast.num_params(), 0)


torch/jit/__init__.py

-    graph = _script_graph(fn, _frames_up=_frames_up + 1)
    mod = ScriptModule()
-    mod._create_method_from_graph('forward', graph)
+    _script_pure_function(mod, fn, _frames_up=_frames_up + 1)


torch/csrc/jit/script/compiler.cpp

      auto& name = (*it).ident().name();
-      arguments.push_back({name, DynamicType::get()});
+      TypePtr arg_type = DynamicType::get();
+      if (def_and_types.arg_types.size()) {


torch/jit/__init__.py

+    ast = get_jit_ast(fn)
+    arg_types, return_type = annotations.get_signature(fn, ast.num_params() - 1, 0)
+    # Dumb handling for `self`
+    if len(arg_types) == ast.num_params():


torch/jit/__init__.py

                torch._C.ScriptModule.__init__(self)
            original_init(self, *args, **kwargs)
-            asts = [m.ast for m in methods]
+            defs = [torch._C.DefAndTypes(m.ast, m.arg_types, m.return_type) for m in methods]


torch/csrc/jit/script/compiler.cpp

-  for(Def def : definitions) {
-    const std::string& name = def.name().name();
+  for(DefAndTypes def : definitions) {
+    const std::string& name = pure_func ? "forward" : def.def.name().name();


torch/csrc/jit/script/compiler.cpp

+      std::vector<TypePtr> flattened_return_types;
+      if (def_and_types.return_type) {
+        if (def_and_types.return_type->kind() == TypeKind::TupleType) {
+          const auto &tuple_type_elmts = def_and_types.return_type->cast<TupleType>()->elements();


torch/csrc/jit/script/python_tree_views.cpp

                         wrap_list(r, std::move(body)));
-    }));
+    }))
+    .def("num_params", [](Def& self) {


torch/csrc/jit/script/compiler.cpp

  FunctionTable table;
  JIT_ASSERT(definitions.size() == resolvers.size());
+  if (pure_func) {
+    JIT_ASSERT(definitions.size() == 1);


torch/csrc/jit/script/compiler.h

  bool allow_varargs;
 };

+struct DefAndTypes {


torch/jit/__init__.py



-def _script_graph(fn, _frames_up=0):
+def _compile_fn(fn, _frames_up=0):


torch/jit/__init__.py

+    _jit_script_compile_pure_fn(mod, torch._C.DefAndTypes(ast, arg_types, ret_type), rcb)
+
+
+def _script_graph(fn, _frames_up=0):


torch/csrc/jit/script/init.cpp


-  m.def("_jit_script_compile", [](Def def, ResolutionCallback rcb) {
-    return compileFunction(def, pythonResolver(rcb));
+  m.def("_jit_script_compile_pure_fn", [](Module &m, DefAndTypes def, ResolutionCallback rcb) {


jamesr66a · 2018-07-18T18:27:31Z

Superseded by #9547

Summary: Supersedes #9405 Pull Request resolved: #9547 Reviewed By: zdevito Differential Revision: D8900327 Pulled By: jamesr66a fbshipit-source-id: a00a94615af4fbaec98ee3ede0cb54bcfd9108dd

Summary: Supersedes pytorch#9405 Pull Request resolved: pytorch#9547 Reviewed By: zdevito Differential Revision: D8900327 Pulled By: jamesr66a fbshipit-source-id: a00a94615af4fbaec98ee3ede0cb54bcfd9108dd

jamesr66a requested review from apaszke, colesbury, ezyang, gchanan, soumith and zdevito as code owners July 13, 2018 00:28

jamesr66a force-pushed the script_type_annotations2 branch from ce55be3 to f88188c Compare July 13, 2018 00:30

[JIT] Plumb type annotations through script compilation

8a01c6b

jamesr66a force-pushed the script_type_annotations2 branch from f88188c to 8a01c6b Compare July 13, 2018 00:32

facebook-github-bot reviewed Jul 13, 2018

View reviewed changes

zdevito reviewed Jul 13, 2018

View reviewed changes

apaszke reviewed Jul 13, 2018

View reviewed changes

jamesr66a mentioned this pull request Jul 18, 2018

[JIT] Plumb type annotations through script compilation (new) #9547

Closed

jamesr66a closed this Jul 18, 2018



		def _script_graph(fn, _frames_up=0):
		def _compile_fn(fn, _frames_up=0):

		_jit_script_compile_pure_fn(mod, torch._C.DefAndTypes(ast, arg_types, ret_type), rcb)


		def _script_graph(fn, _frames_up=0):

[JIT] Plumb type annotations through script compilation #9405

[JIT] Plumb type annotations through script compilation #9405

Uh oh!

Conversation

jamesr66a commented Jul 13, 2018

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

jamesr66a commented Jul 18, 2018

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants