[torch.compile] remove compilation_context and simplify code

[torch.compile] remove compilation_context and simplify code #2579