[MPS] Move MetalShaderLibrary to its own header (pytorch#141475)

In preparation to be used from libtorch_python
Pull Request resolved: pytorch#141475
Approved by: https://github.com/Skylion007
ghstack dependencies: pytorch#141474

Author: malfet

aten/src/ATen/native/mps/MetalShaderLibrary.h

@@ -0,0 +1,61 @@
+#pragma once
+#include <Metal/Metal.h>
+#include <unordered_map>
+#include <vector>
+
+namespace at::native::mps {
+class MetalShaderLibrary {
+ public:
+  MetalShaderLibrary(const std::string& src)
+      : shaderSource(src), nparams(0), compile_options(nullptr) {}
+  MetalShaderLibrary(const std::string& src, unsigned nparams_)
+      : shaderSource(src), nparams(nparams_), compile_options(nullptr) {}
+  MetalShaderLibrary(
+      const std::string& src,
+      unsigned nparams_,
+      MTLCompileOptions* compile_options_)
+      : shaderSource(src),
+        nparams(nparams_),
+        compile_options(compile_options_) {}
+  MetalShaderLibrary(const MetalShaderLibrary&) = delete;
+  inline id<MTLComputePipelineState> getPipelineStateForFunc(
+      const std::string& fname) {
+    return getLibraryPipelineState(getLibrary(), fname).first;
+  }
+  id<MTLComputePipelineState> getPipelineStateForFunc(
+      const std::string& fname,
+      const std::initializer_list<std::string>& params) {
+    return getLibraryPipelineState(getLibrary(params), fname).first;
+  }
+  inline id<MTLFunction> getMTLFunction(const std::string& fname) {
+    return getLibraryPipelineState(getLibrary(), fname).second;
+  }
+  id<MTLFunction> getMTLFunction(
+      const std::string& fname,
+      const std::initializer_list<std::string>& params) {
+    return getLibraryPipelineState(getLibrary(params), fname).second;
+  }
+  static MetalShaderLibrary& getBundledLibrary();
+
+ protected:
+  virtual id<MTLLibrary> getLibrary();
+  virtual id<MTLLibrary> getLibrary(
+      const std::initializer_list<std::string>& params);
+  id<MTLLibrary> library = nil;
+
+ private:
+  std::pair<id<MTLComputePipelineState>, id<MTLFunction>>
+  getLibraryPipelineState(id<MTLLibrary> lib, const std::string& fname);
+
+  id<MTLLibrary> compileLibrary(const std::string& src);
+  std::string shaderSource;
+  unsigned nparams;
+  MTLCompileOptions* compile_options;
+  std::unordered_map<std::string, id<MTLLibrary>> libMap;
+  std::unordered_map<
+      std::string,
+      std::pair<id<MTLComputePipelineState>, id<MTLFunction>>>
+      cplMap;
+};
+
+} // namespace at::native::mps

aten/src/ATen/native/mps/OperationUtils.h

@@ -7,6 +7,7 @@
 #include <ATen/Tensor.h>
 #include <ATen/Utils.h>
 #include <ATen/mps/MPSStream.h>
+#include <ATen/native/mps/MetalShaderLibrary.h>
 #include <ATen/native/mps/TensorFactory.h>
 #include <c10/core/ScalarType.h>
 #include <torch/library.h>
@@ -342,46 +343,6 @@ inline bool is_dense_in_storage(const TensorBase& t) {
   return compute_storage_numel_distance(t) == static_cast<size_t>(t.numel());
 }
 
-class MetalShaderLibrary {
- public:
-  MetalShaderLibrary(const std::string& src) : shaderSource(src), nparams(0), compile_options(nullptr) {}
-  MetalShaderLibrary(const std::string& src, unsigned nparams_)
-      : shaderSource(src), nparams(nparams_), compile_options(nullptr) {}
-  MetalShaderLibrary(const std::string& src, unsigned nparams_, MTLCompileOptions* compile_options_)
-      : shaderSource(src), nparams(nparams_), compile_options(compile_options_) {}
-  MetalShaderLibrary(const MetalShaderLibrary&) = delete;
-  inline id<MTLComputePipelineState> getPipelineStateForFunc(const std::string& fname) {
-    return getLibraryPipelineState(getLibrary(), fname).first;
-  }
-  id<MTLComputePipelineState> getPipelineStateForFunc(const std::string& fname,
-                                                      const std::initializer_list<std::string>& params) {
-    return getLibraryPipelineState(getLibrary(params), fname).first;
-  }
-  inline id<MTLFunction> getMTLFunction(const std::string& fname) {
-    return getLibraryPipelineState(getLibrary(), fname).second;
-  }
-  id<MTLFunction> getMTLFunction(const std::string& fname, const std::initializer_list<std::string>& params) {
-    return getLibraryPipelineState(getLibrary(params), fname).second;
-  }
-  static MetalShaderLibrary& getBundledLibrary();
-
- protected:
-  virtual id<MTLLibrary> getLibrary();
-  virtual id<MTLLibrary> getLibrary(const std::initializer_list<std::string>& params);
-  id<MTLLibrary> library = nil;
-
- private:
-  std::pair<id<MTLComputePipelineState>, id<MTLFunction>> getLibraryPipelineState(id<MTLLibrary> lib,
-                                                                                  const std::string& fname);
-
-  id<MTLLibrary> compileLibrary(const std::string& src);
-  std::string shaderSource;
-  unsigned nparams;
-  MTLCompileOptions* compile_options;
-  std::unordered_map<std::string, id<MTLLibrary>> libMap;
-  std::unordered_map<std::string, std::pair<id<MTLComputePipelineState>, id<MTLFunction>>> cplMap;
-};
-
 namespace detail {
 template <typename T>
 class has_size_type {