commit 96c9bfb66cf24883f359d051281cc99d10a7a69e
author Ben Vanik <ben.vanik@gmail.com> Mon Jul 08 15:44:51 2024 -0700
committer GitHub <noreply@github.com> Mon Jul 08 22:44:51 2024 +0000
tree 8e74da8617ec68a749d1796c50aef4872a3175e2
parent 129878fcf1c0356d736443d1196a1f11b783a12b

Optimizing queries for optional VM functions. (#17823)

Queries with IMPORT_OPTIONAL bypass reporting and full status creation
as a way to avoid generating a bunch of stack traces for cases where the
caller can gracefully handle the absence of an import. This change adds
EXPORT_OPTIONAL to provide the same functionality to callers that know
the export may not exist. This speeds up startup (when __init is omitted
on a module), shutdown (when __deinit is omitted on a module), and event
notification (__notify).

runtime/bindings/python/py_module.cc

diff --git a/runtime/bindings/python/py_module.cc b/runtime/bindings/python/py_module.cc
index 86eb1e1..85dbb2e 100644
--- a/runtime/bindings/python/py_module.cc
+++ b/runtime/bindings/python/py_module.cc
@@ -86,7 +86,8 @@
       iree_vm_function_t* out_function, iree_string_view_t* out_name,
       iree_vm_function_signature_t* out_signature) {
     auto self = AsSelf(vself);
-    if (IREE_LIKELY(linkage == IREE_VM_FUNCTION_LINKAGE_EXPORT)) {
+    if (IREE_LIKELY(linkage == IREE_VM_FUNCTION_LINKAGE_EXPORT ||
+                    linkage == IREE_VM_FUNCTION_LINKAGE_EXPORT_OPTIONAL)) {
       if (IREE_LIKELY(ordinal < self->export_functions_.size())) {
         std::unique_ptr<PyFunction>& f = self->export_functions_[ordinal];
         if (IREE_LIKELY(out_function)) {
@@ -114,7 +115,8 @@
       iree_vm_function_t* out_function) {
     auto self = AsSelf(vself);
     std::string_view name_cpp(name.data, name.size);
-    if (linkage == IREE_VM_FUNCTION_LINKAGE_EXPORT) {
+    if (linkage == IREE_VM_FUNCTION_LINKAGE_EXPORT ||
+        linkage == IREE_VM_FUNCTION_LINKAGE_EXPORT_OPTIONAL) {
       auto found_it = self->export_name_to_ordinals_.find(name_cpp);
       if (found_it != self->export_name_to_ordinals_.end()) {
         out_function->linkage = linkage;

runtime/bindings/python/vm.cc

diff --git a/runtime/bindings/python/vm.cc b/runtime/bindings/python/vm.cc
index f8b758c..0be8912 100644
--- a/runtime/bindings/python/vm.cc
+++ b/runtime/bindings/python/vm.cc
@@ -819,6 +819,7 @@
       .value("IMPORT", IREE_VM_FUNCTION_LINKAGE_IMPORT)
       .value("IMPORT_OPTIONAL", IREE_VM_FUNCTION_LINKAGE_IMPORT_OPTIONAL)
       .value("EXPORT", IREE_VM_FUNCTION_LINKAGE_EXPORT)
+      .value("EXPORT_OPTIONAL", IREE_VM_FUNCTION_LINKAGE_EXPORT_OPTIONAL)
       .export_values();
 
   auto vm_buffer = py::class_<VmBuffer>(m, "VmBuffer");