Fixes for profiling JIT code (#38453)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/38453
Two fixes:
- RecordFunction in JIT interpreter should exist during the execution
of the frame, and not just when we enter the frame
- When creating a JIT continuation in wait instruction, we'd want to
preserve the original thread local context, right now when we resume
execution in continuation we preserve the thread local state of the
thread that set future value (i.e. executed a forked task)
Test Plan: unittest, CI
Reviewed By: ngimel
Differential Revision: D21565959
Pulled By: ilia-cher
fbshipit-source-id: 206b98e3bfb0052fc8e4031da778e372cc71afc1
