Metal backend: stop tracking blob memory once a model is destroyed (pytorch#22954)

### Summary

Fixes pytorch#22953

Tensors created from a model's constants blob are tracked in
`memory_to_n_tensor` as `NOT_OWN`, and `aoti_torch_delete_tensor_object`
leaves those entries in place when it deletes such a tensor.
`cleanup_memory()` cleared `tensors` but not that map, so the addresses
stayed tracked after the model was gone. The next model whose constants
landed on a reused address then failed to load with "Memory address ...
is already being tracked by another tensor". Loading six small Metal
models twice in one process failed 7 of 12 loads.

`cleanup_memory()` now clears `memory_to_n_tensor` once every tensor has
been deleted. The entry is not erased in the `NOT_OWN` branch of
`aoti_torch_delete_tensor_object` because several tensors can alias one
`NOT_OWN` address without a count, and erasing on the first delete would
make the later ones fail with "memory not found during deletion".

### Test plan

New `backends/apple/metal/runtime/test/test_memory.cpp`, wired through
`et_cxx_test` under `BUILD_TESTING` the same way `backends/cuda` does
it:

```
cmake -S . -B cmake-out-test -DEXECUTORCH_BUILD_METAL=ON -DAOTI_METAL=ON \
  -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
  -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
  -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON -DEXECUTORCH_BUILD_TESTS=ON
cmake --build cmake-out-test --target test_metal_memory
cmake-out-test/backends/apple/metal/test_metal_memory
```

- Before the change: `BlobAddressCanBeReusedAfterCleanup` and
`CleanupLeavesNoTrackedMemory` fail;
`BlobAddressIsTrackedWhileTensorIsAlive` (the control) passes.
- After: `[  PASSED  ] 3 tests.`
- End to end, the load/destroy/load loop from the issue went from `12
loads, 7 failures` to `12 loads, 0 failures` in three consecutive runs.
- `python -m unittest
backends.apple.metal.tests.test_modules.TestMetalBackendModules`: `Ran
134 tests ... OK`.
- `lintrunner` clean on the touched files.

Author: abdelaziz-mahdy

backends/apple/metal/CMakeLists.txt

@@ -139,3 +139,12 @@ install(
   EXPORT ExecuTorchTargets
   DESTINATION lib
 )
+
+if(BUILD_TESTING)
+  include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
+
+  et_cxx_test(
+    test_metal_memory SOURCES runtime/test/test_memory.cpp EXTRA_LIBS
+    metal_backend
+  )
+endif()

backends/apple/metal/runtime/shims/memory.cpp

@@ -706,6 +706,13 @@ void cleanup_memory() {
   // tensors map should now be empty, but ensure it's cleared
   tensors.clear();
 
+  // Tensors created from a blob are tracked as NOT_OWN and
+  // aoti_torch_delete_tensor_object leaves their address in the map, since
+  // several of them may alias it. With every tensor gone nothing is tracked
+  // anymore, and a stale entry would make the next model fail to load as soon
+  // as its constants land on an address used before.
+  memory_to_n_tensor.clear();
+
   // Clean up Metal resources
   metal_cleanup_resources();
 

backends/apple/metal/runtime/test/test_memory.cpp

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <vector>
+
+#include <executorch/backends/apple/metal/runtime/shims/memory.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/platform/platform.h>
+
+using namespace executorch::backends::metal;
+using executorch::runtime::Error;
+
+namespace {
+
+// ScalarType::Float, as AOTInductor passes it to the shims.
+constexpr int32_t kFloat32 = 6;
+// DeviceType::MPS.
+constexpr int32_t kDeviceMps = 13;
+
+} // namespace
+
+class MetalMemoryTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    et_pal_init();
+  }
+
+  void TearDown() override {
+    cleanup_memory();
+  }
+
+  // Wraps `data` the way a model's constants are wrapped: a tensor that
+  // borrows memory it does not own.
+  Error createFromBlob(void* data, AOTITensorHandle* tensor) {
+    const std::vector<int64_t> sizes = {2, 2};
+    const std::vector<int64_t> strides = {2, 1};
+    return aoti_torch_create_tensor_from_blob_v2(
+        data,
+        static_cast<int64_t>(sizes.size()),
+        sizes.data(),
+        strides.data(),
+        /*storage_offset=*/0,
+        kFloat32,
+        kDeviceMps,
+        /*device_index=*/0,
+        tensor,
+        /*layout=*/0,
+        /*opaque_metadata=*/nullptr,
+        /*opaque_metadata_size=*/0);
+  }
+
+  std::vector<float> blob_ = std::vector<float>(4, 1.0f);
+};
+
+TEST_F(MetalMemoryTest, BlobAddressIsTrackedWhileTensorIsAlive) {
+  AOTITensorHandle first = nullptr;
+  ASSERT_EQ(createFromBlob(blob_.data(), &first), Error::Ok);
+
+  AOTITensorHandle second = nullptr;
+  EXPECT_NE(createFromBlob(blob_.data(), &second), Error::Ok);
+}
+
+// A model's constants blob can be mapped at an address a previously destroyed
+// model used. cleanup_memory() runs when a model is destroyed, so it must leave
+// nothing tracked behind.
+TEST_F(MetalMemoryTest, BlobAddressCanBeReusedAfterCleanup) {
+  AOTITensorHandle first = nullptr;
+  ASSERT_EQ(createFromBlob(blob_.data(), &first), Error::Ok);
+
+  cleanup_memory();
+
+  AOTITensorHandle second = nullptr;
+  EXPECT_EQ(createFromBlob(blob_.data(), &second), Error::Ok);
+}
+
+TEST_F(MetalMemoryTest, CleanupLeavesNoTrackedMemory) {
+  AOTITensorHandle tensor = nullptr;
+  ASSERT_EQ(createFromBlob(blob_.data(), &tensor), Error::Ok);
+  ASSERT_FALSE(memory_to_n_tensor.empty());
+
+  cleanup_memory();
+
+  EXPECT_TRUE(tensors.empty());
+  EXPECT_TRUE(memory_to_n_tensor.empty());
+}