Update

malfet · malfet · commit 9f7dde6af652 · 2024-12-02T15:01:14.000-08:00
[ghstack-poisoned]
diff --git a/aten/src/ATen/native/mps/MetalShaderLibrary.h b/aten/src/ATen/native/mps/MetalShaderLibrary.h
@@ -45,9 +45,9 @@ constexpr bool has_size_type_v = has_size_type<T>::value;
 
 class MetalKernelFunction {
  public:
-  MetalKernelFunction(MTLComputePipelineState_t cps_) : cps(cps_) {}
-  MetalKernelFunction(MetalKernelFunction&) = delete;
+  MetalKernelFunction(MTLComputePipelineState_t cps_);
   ~MetalKernelFunction();
+  MetalKernelFunction(MetalKernelFunction&) = delete;
   // Shader properties
   uint64_t getMaxThreadsPerThreadgroup() const;
   uint64_t getThreadExecutionWidth() const;
@@ -152,6 +152,7 @@ class DynamicMetalShaderLibrary : public MetalShaderLibrary {
     // Compile right away
     getLibrary();
   }
+  ~DynamicMetalShaderLibrary();
 };
 
 } // namespace at::native::mps
diff --git a/aten/src/ATen/native/mps/OperationUtils.mm b/aten/src/ATen/native/mps/OperationUtils.mm
@@ -923,8 +923,17 @@ static dispatch_data_t getSectionData(const std::string& name) {
   return l;
 }
 
+// DynamicMetalShaderLibrary implementation
+DynamicMetalShaderLibrary::~DynamicMetalShaderLibrary() {
+  [library release];
+}
+
 // MetalKernelFunction implementation
-MetalKernelFunction::~MetalKernelFunction() {}
+MetalKernelFunction::MetalKernelFunction(MTLComputePipelineState_t cps_) : cps([cps_ retain]) {}
+
+MetalKernelFunction::~MetalKernelFunction() {
+  [cps release];
+}
 
 void MetalKernelFunction::runCommandBlock(std::function<void(void)> run) {
   dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() {
