Address comments

Author: larryliu0820

backends/cuda/cuda_backend.py

@@ -183,9 +183,9 @@ def preprocess(
             elif path.endswith(".wrapper_weights.blob"):
                 blob_path = path
 
-        if so_path is None:
+        if so_path is None or blob_path is None:
             raise RuntimeError(
-                f"Could not find .wrapper.so file in compiled paths, got {paths}"
+                f"Could not find required files in compiled paths, got {paths}"
             )
 
         # pyre-ignorep[6]: Incompatible parameter type
@@ -198,15 +198,14 @@ def preprocess(
         # Keep the so file in the NamedDataStore, so that it can be packaged into the .pte file.
         named_data_store.add_named_data(method_name + "_so_blob", so_data, 1, None)
 
-        # Add weights blob to named data store if it exists
-        if blob_path is not None:
-            with open(blob_path, "rb") as f:
-                blob_data = f.read()
-            named_data_store.add_named_data(
-                method_name + "_weights_blob", blob_data, 1, "aoti_cuda_blob"
-            )
-            # Clean up the weights blob file
-            os.remove(blob_path)
+        # Add weights blob to named data store
+        with open(blob_path, "rb") as f:
+            blob_data = f.read()
+        named_data_store.add_named_data(
+            method_name + "_weights_blob", blob_data, 1, "aoti_cuda_blob"
+        )
+        # Clean up the weights blob file
+        os.remove(blob_path)
 
         # Clean up the generated so file; it has been packaged into the NamedDataStore
         # pyre-ignorep[6]: Incompatible parameter type

backends/cuda/runtime/cuda_backend.cpp

@@ -28,13 +28,6 @@
 
 namespace executorch::backends::cuda {
 
-#define LOAD_SYMBOL(handle, member, name, so_handle)                        \
-  do {                                                                      \
-    handle->member = reinterpret_cast<name##Func>(dlsym(so_handle, #name)); \
-    ET_CHECK_OR_RETURN_ERROR(                                               \
-        handle->member != nullptr, AccessFailed, "Failed to load " #name);  \
-  } while (0)
-
 using namespace std;
 using namespace aoti;
 
@@ -60,35 +53,32 @@ class ET_EXPERIMENTAL CudaBackend final
   Error load_function_pointers_into_handle(
       void* so_handle,
       AOTIDelegateHandle* handle) const {
-    LOAD_SYMBOL(
-        handle,
-        create_with_device,
-        AOTInductorModelContainerCreateWithDevice,
-        so_handle);
-
-    LOAD_SYMBOL(
-        handle, delete_container, AOTInductorModelContainerDelete, so_handle);
-
-    LOAD_SYMBOL(
-        handle,
-        get_num_inputs,
-        AOTInductorModelContainerGetNumInputs,
-        so_handle);
-
-    LOAD_SYMBOL(
-        handle,
-        get_num_outputs,
-        AOTInductorModelContainerGetNumOutputs,
-        so_handle);
-
-    LOAD_SYMBOL(handle, run, AOTInductorModelContainerRun, so_handle);
-
-    LOAD_SYMBOL(
-        handle,
-        update_constants_from_blob,
-        AOTInductorModelUpdateConstantsFromBlob,
-        so_handle);
+#define LOAD_SYMBOL(member, name)                                           \
+  do {                                                                      \
+    handle->member = reinterpret_cast<name##Func>(dlsym(so_handle, #name)); \
+    ET_CHECK_OR_RETURN_ERROR(                                               \
+        handle->member != nullptr, AccessFailed, "Failed to load " #name);  \
+  } while (0)
+    LOAD_SYMBOL(create_with_device, AOTInductorModelContainerCreateWithDevice);
+
+    LOAD_SYMBOL(delete_container, AOTInductorModelContainerDelete);
 
+    LOAD_SYMBOL(get_num_inputs, AOTInductorModelContainerGetNumInputs);
+
+    LOAD_SYMBOL(get_num_outputs, AOTInductorModelContainerGetNumOutputs);
+
+    LOAD_SYMBOL(run, AOTInductorModelContainerRun);
+#undef LOAD_SYMBOL
+
+    handle->update_constants_from_blob =
+        reinterpret_cast<AOTInductorModelUpdateConstantsFromBlobFunc>(
+            dlsym(so_handle, "AOTInductorModelUpdateConstantsFromBlob"));
+    if (handle->update_constants_from_blob == nullptr) {
+      ET_LOG(
+          Info,
+          "Failed to load AOTInductorModelUpdateConstantsFromBlob"
+          " this .so is probably compiled on an old version of torch (<2.9.0)");
+    }
     return Error::Ok;
   }
 
@@ -184,7 +174,7 @@ class ET_EXPERIMENTAL CudaBackend final
     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());
-    if (buffer_res.ok()) {
+    if (buffer_res.ok() && handle->update_constants_from_blob != nullptr) {
       ET_LOG(Info, "Found %s in named data map", weights_blob_key.c_str());
       const void* weights_blob = buffer_res->data();
       // Feed the weights blob into the container. Under the hood it's copying