[jit] _script_compile and _script_class_compile add to the python CU #22221
Conversation
Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python.
pytorchbot
added
oncall: jit
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
zdevito
left a comment
zdevito
left a comment
There was a problem hiding this comment.
This one looks fine. I think the previous one should not be necessary.
torch/csrc/jit/script/compiler.cpp
Outdated
| name, is_optimized(), std::make_shared<Graph>(), creator); | ||
| if (self) { | ||
| // Register this as a method on `self`'s type | ||
| self->getClassType()->registerMethod(fn.get()); |
There was a problem hiding this comment.
If this is the only place where getClassType is used, then I think it should be accomplished in a different way. compilation unit define is not the right place for adding it as a method. It might be better to have the mega-define return a list of things it defined instead and let the actual class method compiler add it as a function.
There was a problem hiding this comment.
yeah that's fair—I was a little concerned that you could get them desynchronized if some random caller did define without knowing there was an associated registerMethod that was required, but this is prbably too intrusive
There was a problem hiding this comment.
Actually, I think we have to do it this way. Otherwise we won't have registered the methods when we called ensure_defined(), so methods won't be able to find each other.
| classes_.push_back(std::move(classType)); | ||
| void register_class(c10::NamedTypePtr namedType) { | ||
| if (auto classType = namedType->cast<c10::ClassType>()) { | ||
| // TODO: class types cannot be redefined because we have no way right now |
There was a problem hiding this comment.
yuck -- we should have consistent behavior. I don't think we should let anything get redefined.
There was a problem hiding this comment.
Definitely agree this is gross, but we're in a difficult spot. Basically:
- Our tests frequently redefine Functions. That worked before because they all went in a different CU, but now there is a global one.
- But aren't these functions all different because they occur nested in different functions? No, because in Python there is no way of knowing whether a function is nested within another.
- Okay, fine. But then shouldn't we clear the Python CU between every test for isolation anyway? But we can't do that, because we define script funcs in the standard library and the resolvers for our methods will capture a dangling StrongFunctionPtr.
I think the last one can be untangled, but I didn't want to couple it with this PR. Also, in the long run, do we really want the behavior to be inconsistent with Python (where you can redefine pretty much anything)?
There was a problem hiding this comment.
There is a way to figure out whether something is nested or not: get its qualified name, try to access it starting from the global environment, did you get the function again? If so, it has that global name, if not, it was nested in something. Global names are important for export/import but should never be used as the way we link functions together in the Python environment. Given that we do not consistently keep track of when we inline, it will result in all sorts of bugs. If something is not a globally exposed function (i.e. accessible in python following the module hierarchy), we should be comfortable giving it a mangled name.
There was a problem hiding this comment.
Ah I see what you mean now. As an aside, one issue I'm running into with what you proposed is that when we do:
@torch.jit.script
def foo():
...
The script() decorator cannot figure out whether it can access foo from the global environment, because while script() is running foo has not been defined yet. So we can't, at definition time, let the CU know whether the function name needs to be mangled or not.
Let me know if you have any ideas. Possibly there is some Python trick that I need to Google for.
| std::shared_ptr<std::vector<bool>> parameterSlots_; | ||
|
|
||
| // List of methods associated with this class. | ||
| std::vector<Function*> methods_; |
There was a problem hiding this comment.
Ok, this appears in this PR.
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. Additional changes that need to be made as part of this: 1. All scripted functions will get their name mangled so that they are uniquely named in the global CU 2. Classes keep track of their methods 3. Clearing the Python CU involves deleting all classes *and* their methods. 4. The global python CU was moved into Python, because of static destruction order issues. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
pytorchbot
added
module: pybind
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. Additional changes that need to be made as part of this: 1. All scripted functions will get their name mangled so that they are uniquely named in the global CU 2. Classes keep track of their methods 3. Clearing the Python CU involves deleting all classes *and* their methods. 4. The global python CU was moved into Python, because of static destruction order issues. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. Additional changes that need to be made as part of this: 1. All scripted functions will get their name mangled so that they are uniquely named in the global CU 2. Classes keep track of their methods 3. Clearing the Python CU involves deleting all classes *and* their methods. 4. The global python CU was moved into Python, because of static destruction order issues. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. Additional changes that need to be made as part of this: 1. All scripted functions will get their name mangled so that they are uniquely named in the global CU 2. Classes keep track of their methods 3. Clearing the Python CU involves deleting all classes *and* their methods. 4. The global python CU was moved into Python, because of static destruction order issues. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
… python CU" [jit] _script_compile and _script_class_compile add to the python CU Actually wire things up so when you do `@torch.jit.script` from Python things will be put in the Python CU. This removes the hack where we have to keep the StrongFunctionPtrs alive in Python. Additional changes that need to be made as part of this: 1. All scripted functions will get their name mangled so that they are uniquely named in the global CU 2. Classes keep track of their methods 3. Clearing the Python CU involves deleting all classes *and* their methods. 4. The global python CU was moved into Python, because of static destruction order issues. gh-metadata: pytorch pytorch 22221 gh/suo/72/head
Stack from ghstack:
Actually wire things up so when you do
@torch.jit.scriptfrom Pythonthings will be put in the Python CU. This removes the hack where we have
to keep the StrongFunctionPtrs alive in Python.
Additional changes that need to be made as part of this:
uniquely named in the global CU
methods.
destruction order issues.
gh-metadata: pytorch pytorch 22221 gh/suo/72/head
Differential Revision: D15998757