[jit] Add weak script mode for script functions #11963

driazati · 2018-09-22T01:52:12Z

This PR is the start of weak script mode for functions

Weak scripts allow you to compile a graph from Python code at runtime by annotating with torch.jit.weak_script for use in the JIT without affecting eager execution. Scripts are compiled lazily on the first call in a graph to avoid long Python startup times.

@apaszke @zdevito @ezyang

Summary: Weak scripts allow you to compile a graph from Python code at runtime by annotating with torch.jit.weak_script for use in the JIT without affecting Python execution. Scripts are compiled lazily on the first call in a graph to avoid long Python startup times. ```python import torch @torch.jit.weak_script def fn(x): return x + 5 @torch.jit.script def run_in_graph(x): return fn(x) + 1 x = torch.randn(4) # this should cause fn to compile and run graph script_out = run_in_graph(x) # this should run fn from Python (without calling into C++) expected_out = fn(x) ``` Test Plan: Reviewers: Subscribers: Tasks: Tags:

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

torch/nn/functional.py


 _VF = torch._C._VariableFunctions

+def weak_script(fn):


Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

zdevito

See comments for suggestions about how to structure the code in the compiler. In particular, it should mirror how (strong) script functions flow through the compiler closely.

test/test_jit.py

torch/csrc/jit/interned_strings.h

  static Symbol aten(const std::string & s);
  static Symbol onnx(const std::string & s);
  static Symbol prim(const std::string & s);
+  static Symbol weak(const std::string & s);


torch/csrc/jit/script/init.cpp

  return std::make_shared<SimpleValue>(v);
 }

+static py::function compile_func_;


torch/csrc/jit/script/init.cpp

+  const Def& def = py::cast<const Def&>(tuple[0]);
+  ResolutionCallback rcb = py::cast<ResolutionCallback>(tuple[1]);
+  auto m = std::make_shared<Module>();
+  defineMethodsInModule(*m, {def}, {pythonResolver(rcb)}, nullptr);


torch/csrc/jit/script/init.cpp

-    std::string cconv(inputs.size(), 'd');
-    Node* new_node = m.graph()->insertNode(m.graph()->createPythonOp(
-      THPObjectPtr(func.release().ptr()), cconv, {}));
+


torch/jit/__init__.py

+def weak_script(fn):
+    # register op for jit
+    # no op if called from python
+    fn.is_weak = True


Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

…_script

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

torch/csrc/jit/script/init.cpp

+
+  if (is_weak) {
+    auto compiled_fn =
+        py::module::import("torch.jit").attr("script")(obj, true, 0, true);


torch/jit/__init__.py

+    return fn
+
+
+def script(fn, optimize=True, _frames_up=0, is_weak=False):


Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

facebook-github-bot

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

zdevito

Approach looks good. I have minor comments on how the code is factored.

torch/jit/__init__.py


-def script(fn, optimize=True, _frames_up=0):
+def weak_script(fn):
+    fn._jit_is_weak_script = True


torch/jit/__init__.py

+    return fn
+
+
+def script(fn, optimize=True, _frames_up=0, _is_weak=False):


torch/jit/__init__.py



+def weak_script(fn):
+    compiled_weak_fns[fn] = "pending";


apaszke · 2018-10-04T00:04:39Z

I'm also pretty sure that we need to create resolution callbacks at the time when weak_script is called and not when it's compiled. That's what we do for script_methods at least.

torch/jit/__init__.py

+    if entry is None:
+        return None
+    if entry["status"] == COMPILATION_PENDING:
+        entry["status"] = COMPILED


torch/jit/__init__.py

+        entry["status"] = COMPILED
+        compiled_fn = torch.jit.script(fn, True, 0, entry["rcb"])
+        compiled_weak_fns[fn]["compiled_fn"] = compiled_fn
+    return compiled_fn


torch/jit/__init__.py

+    return compiled_fn
+
+
+def script(fn, optimize=True, _frames_up=0, rcb=None):


driazati · 2018-10-04T23:31:06Z

@pytorchbot retest this please

zdevito

This looks good. Minor nits to keep the amount of data we hold around lower.

torch/csrc/jit/script/init.cpp

  } else if (py::isinstance<py::module>(obj)) {
    return std::make_shared<PythonModuleValue>(obj);
  }
+


torch/jit/__init__.py

+        compiled_fn = torch.jit.script(fn, True, 0, entry["rcb"])
+        compiled_weak_fns[fn]["compiled_fn"] = compiled_fn
+        entry["status"] = COMPILED
+        return compiled_fn


torch/jit/__init__.py

+    compiled_weak_fns[fn] = {
+        "status": COMPILATION_PENDING,
+        "compiled_fn": None,
+        "rcb": createResolutionCallback(_frames_up + 1)


facebook-github-bot

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

David Riazati added 6 commits September 21, 2018 13:24

Make some nn.functional functions weak script

6d04f40

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Refactor

8c4ab22

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Refactor

0e43d56

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Add weak:: namespace for Symbol

3c93f87

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Fix test_jit.py test

fbf9edc

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

driazati requested review from apaszke, colesbury, ezyang, gchanan, soumith and zdevito as code owners September 22, 2018 01:52

driazati commented Sep 22, 2018

View reviewed changes

torch/nn/functional.py Outdated

_VF = torch._C._VariableFunctions

def weak_script(fn):

This comment was marked as off-topic.

Sign in to view

Lint fix

e0f50ab

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

zdevito reviewed Sep 24, 2018

View reviewed changes

ailzhang added the oncall: jit Add this issue/PR to JIT oncall triage queue label Sep 25, 2018

David Riazati added 5 commits September 26, 2018 16:57

Address comments

2a287e5

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Merge branch 'master' of https://github.com/pytorch/pytorch into weak…

f177e77

…_script

Fix tests

a7363be

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Update expected output

ccaaace

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Fix test again

9bf1d15

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

driazati changed the title ~~[jit][WIP] Add weak script mode for script functions~~ [jit] Add weak script mode for script functions Sep 27, 2018

apaszke reviewed Sep 27, 2018

View reviewed changes

Change param name

6a43336

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

facebook-github-bot reviewed Oct 3, 2018

View reviewed changes

zdevito reviewed Oct 3, 2018

View reviewed changes

David Riazati added 2 commits October 3, 2018 16:18

Clear changes from script()

4d4297f

Style

15c31c7

apaszke reviewed Oct 3, 2018

View reviewed changes

torch/jit/__init__.py Outdated

def weak_script(fn):

compiled_weak_fns[fn] = "pending";

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

David Riazati added 3 commits October 3, 2018 17:10

Comments

33bf8b8

Get rcb eagerly

d4faee1

Merge

3f5d168

apaszke reviewed Oct 4, 2018

View reviewed changes

David Riazati added 2 commits October 4, 2018 09:30

Fix ref before assignment, ordering

7369cda

param name

daab3ac

zdevito approved these changes Oct 5, 2018

View reviewed changes

David Riazati added 2 commits October 5, 2018 14:51

Capture only locals/globals for resolution cb

1d16caf

Rearrage init.cpp

d356b03

facebook-github-bot reviewed Oct 5, 2018

View reviewed changes

facebook-github-bot closed this in 92b0e70 Oct 6, 2018

ezyang added the merged label Jun 26, 2019

		return fn


		def script(fn, optimize=True, _frames_up=0, is_weak=False):

		return compiled_fn


		def script(fn, optimize=True, _frames_up=0, rcb=None):

[jit] Add weak script mode for script functions #11963

[jit] Add weak script mode for script functions #11963

Uh oh!

Conversation

driazati commented Sep 22, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

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!

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!

apaszke commented Oct 4, 2018

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!

driazati commented Oct 4, 2018

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!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

driazati commented Sep 22, 2018 •

edited

Loading