ghstack-source-id: 35ecbd2
Pull-Request-resolved: #159456

ghstack-source-id: e38c1fd
Pull-Request-resolved: #159456

ghstack-source-id: 8d0bb39
Pull-Request-resolved: #159456

ghstack-source-id: a575c17
Pull-Request-resolved: #159456

ghstack-source-id: 3dc75a0
Pull-Request-resolved: #159456

This reverts commit 19f1f99.

Reverted #159456 on behalf of https://github.com/davidberard98 due to Sorry - this causes a merge conflict with #159798, which I'm trying to land with co-dev to resolve a sev ([comment](#159456 (comment)))

Pull Request resolved: #159695
Approved by: https://github.com/malfet
ghstack dependencies: #159456

In some cases we have mps kernels which are reused across higher-order-op subgraphs and the toplevel code. However, currently we initialize the variable for the mps kernel the first time we use it, which runs into an issue if we run into the mps kernel within a subgraph since the kernel will only be initialized within the subgraph scope. For instance:
```
if ...
    auto mps_lib_0_func = ...
    mps_lib_0_func->run()

// since we already used mps_lib_0 once, we don't re-initialize it
mps_lib_0_func->run()  // error, mps_lib_0_func not initialized
```

So the solution we took here is to initialize all the kernels at the beginning:
```
const std::shared_ptr<at::native::mps::MetalKernelFunction> get_mps_lib_0() {
    static const auto func = mps_lib_0.getKernelFunction("generated_kernel");
    return func;
}
AOTIMetalKernelFunctionHandle get_mps_lib_0_handle() {
    static const auto handle = AOTIMetalKernelFunctionHandle(get_mps_lib_0().get());
    return handle;
}
...
if ...
    get_mps_lib_0()->run()

get_mps_lib_0()->run()  // success
```

Pull Request resolved: #159753
Approved by: https://github.com/malfet
ghstack dependencies: #159456, #159695