Per-instance named-data keys for AOTI delegate blobs (pytorch#20424)

Author: shoumikhin

backends/aoti/aoti_backend.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import contextlib
+import hashlib
 import os
 import typing
 from abc import ABC, abstractmethod
@@ -276,18 +277,21 @@ def preprocess(
 
         # Create named data store
         named_data_store = NamedDataStore()
-        method_name = cls.method_name_from_compile_specs(compile_specs)
 
-        named_data_store.add_named_data(method_name + "_so_blob", so_data, 1, None)
+        # Key each blob by a content hash so partitions in one method get distinct
+        # keys (a method-name-only key collides). Runtime recovers them from
+        # processed_bytes below.
+        so_blob_key = hashlib.sha256(so_data).hexdigest() + "_so_blob"
+        weights_blob_key = hashlib.sha256(blob_data).hexdigest() + "_weights_blob"
+
+        named_data_store.add_named_data(so_blob_key, so_data, 1, None)
         # Determine whether to save named data externally based on backend setting
         # External: save to separate .ptd file, otherwise merge with .pte file
         external_tag = (
             f"aoti_{device_name}_blob" if cls.save_data_externally() else None
         )
 
-        named_data_store.add_named_data(
-            method_name + "_weights_blob", blob_data, 1, external_tag
-        )
+        named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag)
 
         # Clean up the generated files
         os.remove(so_path)
@@ -299,8 +303,11 @@ def preprocess(
         # the next preprocess call (e.g. for the next method).
         cls.release_moved_tensors(device_edge_program, compile_specs)
 
+        # The runtime cannot recompute these hash keys, so carry them (one per line).
+        processed_bytes = (so_blob_key + "\n" + weights_blob_key).encode("utf-8")
+
         return PreprocessResult(
-            processed_bytes=b"",
+            processed_bytes=processed_bytes,
             debug_handle_map={},
             data_store_output=named_data_store.get_named_data_store_output(),
         )

backends/aoti/aoti_delegate_handle.h

@@ -10,13 +10,15 @@
 
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
+#include <executorch/runtime/core/freeable_buffer.h>
 #include <string>
 
 namespace executorch {
 namespace backends {
 namespace aoti {
 
 using executorch::runtime::Error;
+using executorch::runtime::FreeableBuffer;
 using executorch::runtime::etensor::Tensor;
 
 extern "C" {
@@ -148,6 +150,30 @@ struct AOTIDelegateHandle {
       update_user_managed_constant_buffer_pairs;
 };
 
+// New-format payload is "<so_key>\n<weights_key>"; an empty payload is a
+// pre-this-change artifact, so fall back to the legacy method-name keys.
+inline Error resolve_blob_keys(
+    const FreeableBuffer* processed,
+    const std::string& method_name,
+    std::string& so_blob_key,
+    std::string& weights_blob_key) {
+  if (processed != nullptr && processed->size() > 0) {
+    const std::string keys(
+        static_cast<const char*>(processed->data()), processed->size());
+    const size_t newline = keys.find('\n');
+    if (newline == std::string::npos) {
+      return Error::Internal;
+    }
+    so_blob_key = keys.substr(0, newline);
+    weights_blob_key = keys.substr(newline + 1);
+  } else {
+    so_blob_key = method_name.empty() ? "so_blob" : method_name + "_so_blob";
+    weights_blob_key =
+        method_name.empty() ? "weights_blob" : method_name + "_weights_blob";
+  }
+  return Error::Ok;
+}
+
 } // namespace aoti
 } // namespace backends
 } // namespace executorch

backends/aoti/tests/TARGETS

@@ -3,6 +3,18 @@ load("@fbcode_macros//build_defs/lib:re_test_utils.bzl", "re_test_utils")
 
 oncall("executorch")
 
+cpp_unittest(
+    name = "test_resolve_blob_keys",
+    srcs = [
+        "test_resolve_blob_keys.cpp",
+    ],
+    deps = [
+        "//executorch/backends/aoti:delegate_handle",
+        "//executorch/runtime/core:core",
+        "//executorch/runtime/core:evalue",
+    ],
+)
+
 cpp_unittest(
     name = "test_common_shims",
     srcs = [

backends/aoti/tests/test_resolve_blob_keys.cpp

@@ -0,0 +1,53 @@
+/*
+ * 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 <executorch/backends/aoti/aoti_delegate_handle.h>
+
+#include <gtest/gtest.h>
+#include <string>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/freeable_buffer.h>
+
+using executorch::backends::aoti::resolve_blob_keys;
+using executorch::runtime::Error;
+using executorch::runtime::FreeableBuffer;
+
+TEST(ResolveBlobKeysTest, ParsesKeysFromPayload) {
+  const std::string payload = "aaa_so_blob\nbbb_weights_blob";
+  FreeableBuffer processed(payload.data(), payload.size(), nullptr);
+  std::string so_key;
+  std::string weights_key;
+
+  ASSERT_EQ(
+      resolve_blob_keys(&processed, "forward", so_key, weights_key), Error::Ok);
+  EXPECT_EQ(so_key, "aaa_so_blob");
+  EXPECT_EQ(weights_key, "bbb_weights_blob");
+}
+
+TEST(ResolveBlobKeysTest, FallsBackToMethodNameKeysWhenEmpty) {
+  FreeableBuffer processed; // size 0: a pre-this-change artifact
+  std::string so_key;
+  std::string weights_key;
+
+  ASSERT_EQ(
+      resolve_blob_keys(&processed, "forward", so_key, weights_key), Error::Ok);
+  EXPECT_EQ(so_key, "forward_so_blob");
+  EXPECT_EQ(weights_key, "forward_weights_blob");
+}
+
+TEST(ResolveBlobKeysTest, FailsOnMalformedPayload) {
+  const std::string payload = "missing_the_newline_separator";
+  FreeableBuffer processed(payload.data(), payload.size(), nullptr);
+  std::string so_key;
+  std::string weights_key;
+
+  EXPECT_EQ(
+      resolve_blob_keys(&processed, "forward", so_key, weights_key),
+      Error::Internal);
+}

backends/apple/metal/runtime/metal_backend.cpp

@@ -245,8 +245,12 @@ class ET_EXPERIMENTAL MetalBackend final
       }
     }
 
-    std::string so_blob_key =
-        method_name.empty() ? "so_blob" : method_name + "_so_blob";
+    std::string so_blob_key;
+    std::string weights_blob_key;
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        executorch::backends::aoti::resolve_blob_keys(
+            processed, method_name, so_blob_key, weights_blob_key),
+        "Malformed named-data key payload");
     ET_LOG(Info, "MetalBackend::init - so_blob_key: %s", so_blob_key.c_str());
 
     const NamedDataMap* named_data_map = context.get_named_data_map();
@@ -258,8 +262,6 @@ class ET_EXPERIMENTAL MetalBackend final
     // Prefetch the weights blob — trigger async readahead so pages are
     // resident by the time update_constants_from_blob memcpy's them.
     // This overlaps disk I/O with the .so write + dlopen (~200ms).
-    std::string weights_blob_key =
-        method_name.empty() ? "weights_blob" : method_name + "_weights_blob";
     {
       auto prefetch_buf = named_data_map->get_data(weights_blob_key.c_str());
       if (prefetch_buf.ok() && prefetch_buf->data() != nullptr) {

backends/cuda/runtime/cuda_backend.cpp

@@ -319,8 +319,12 @@ class ET_EXPERIMENTAL CudaBackend final
       }
     }
 
-    std::string so_blob_key =
-        method_name.empty() ? "so_blob" : method_name + "_so_blob";
+    std::string so_blob_key;
+    std::string weights_blob_key;
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        executorch::backends::aoti::resolve_blob_keys(
+            processed, method_name, so_blob_key, weights_blob_key),
+        "Malformed named-data key payload");
 
     const NamedDataMap* named_data_map = context.get_named_data_map();
     auto aoti_dso_buffer = named_data_map->get_data(so_blob_key.c_str());
@@ -394,11 +398,11 @@ class ET_EXPERIMENTAL CudaBackend final
     // methods are independent sub-graphs that may have FQN collisions
     // (e.g. parakeet).
     if (is_weight_sharing_across_methods_enabled()) {
-      ET_CHECK_OK_OR_RETURN_ERROR(
-          load_constants_with_cache(handle, named_data_map, method_name));
+      ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache(
+          handle, named_data_map, method_name, weights_blob_key));
     } else {
       ET_CHECK_OK_OR_RETURN_ERROR(
-          load_constants_legacy(handle, named_data_map, method_name));
+          load_constants_legacy(handle, named_data_map, weights_blob_key));
     }
 
     // Use shared CUDA stream if enabled via options, otherwise create one.
@@ -1011,13 +1015,14 @@ class ET_EXPERIMENTAL CudaBackend final
   Error load_constants_with_cache(
       cuda::CudaDelegateHandle* handle,
       const NamedDataMap* named_data_map,
-      const std::string& method_name) const {
+      const std::string& method_name,
+      const std::string& weights_blob_key) const {
     // Check if the required APIs are available
     if (!handle->get_num_constants || !handle->get_constant_name ||
         !handle->get_constant_original_fqn || !handle->extract_constants_map ||
         !handle->update_user_managed_constant_buffer_pairs) {
       // Fall back to the legacy path
-      return load_constants_legacy(handle, named_data_map, method_name);
+      return load_constants_legacy(handle, named_data_map, weights_blob_key);
     }
 
     // Step 1: Enumerate constants and partition into cached/uncached
@@ -1069,8 +1074,6 @@ class ET_EXPERIMENTAL CudaBackend final
     if (!uncached_fqns.empty()) {
       // Need to load from blob — use update_constants_from_blob for all,
       // then extract the new constants into the cache.
-      std::string weights_blob_key =
-          method_name.empty() ? "weights_blob" : method_name + "_weights_blob";
       auto buffer_res = named_data_map->get_data(weights_blob_key.c_str());
 
       ET_CHECK_OR_RETURN_ERROR(
@@ -1190,9 +1193,7 @@ class ET_EXPERIMENTAL CudaBackend final
   Error load_constants_legacy(
       cuda::CudaDelegateHandle* handle,
       const NamedDataMap* named_data_map,
-      const std::string& method_name) const {
-    std::string weights_blob_key =
-        method_name.empty() ? "weights_blob" : method_name + "_weights_blob";
+      const std::string& weights_blob_key) const {
     auto buffer_res = named_data_map->get_data(weights_blob_key.c_str());
     if (buffer_res.ok() && handle->update_constants_from_blob != nullptr) {
       ET_LOG(Info, "Found %s in named data map", weights_blob_key.c_str());