feat: Fix perf issue in Map.get (#341)

Should fix #326

Author: junrushao

python/tvm_ffi/_ffi_api.py

@@ -23,7 +23,7 @@
 from typing import Any, Callable, TYPE_CHECKING
 if TYPE_CHECKING:
     from collections.abc import Mapping, Sequence
-    from tvm_ffi import Module
+    from tvm_ffi import Module, Object
     from tvm_ffi.access_path import AccessPath
 # isort: on
 # fmt: on
@@ -50,6 +50,8 @@ def Map(*args: Any) -> Any: ...
     def MapCount(_0: Mapping[Any, Any], _1: Any, /) -> int: ...
     def MapForwardIterFunctor(_0: Mapping[Any, Any], /) -> Callable[..., Any]: ...
     def MapGetItem(_0: Mapping[Any, Any], _1: Any, /) -> Any: ...
+    def MapGetItemOrMissing(_0: Mapping[Any, Any], _1: Any, /) -> Any: ...
+    def MapGetMissingObject() -> Object: ...
     def MapSize(_0: Mapping[Any, Any], /) -> int: ...
     def ModuleClearImports(_0: Module, /) -> None: ...
     def ModuleGetFunction(_0: Module, _1: str, _2: bool, /) -> Callable[..., Any] | None: ...
@@ -95,6 +97,8 @@ def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
     "MapCount",
     "MapForwardIterFunctor",
     "MapGetItem",
+    "MapGetItemOrMissing",
+    "MapGetMissingObject",
     "MapSize",
     "ModuleClearImports",
     "ModuleGetFunction",

python/tvm_ffi/container.py

@@ -77,6 +77,8 @@
 V = TypeVar("V")
 _DefaultT = TypeVar("_DefaultT")
 
+MISSING = _ffi_api.MapGetMissingObject()
+
 
 def getitem_helper(
     obj: Any,
@@ -254,12 +256,11 @@ def __contains__(self, item: object) -> bool:
         if not isinstance(item, tuple) or len(item) != 2:
             return False
         key, value = item
-        try:
-            existing_value = self._backend_map[key]
-        except KeyError:
+        actual_value = self._backend_map.get(key, MISSING)
+        if actual_value is MISSING:
             return False
-        else:
-            return existing_value == value
+        # TODO(@junrus): Is `__eq__` the right method to use here?
+        return actual_value == value
 
 
 @register_object("ffi.Map")
@@ -349,10 +350,10 @@ def get(self, key: K, default: V | _DefaultT | None = None) -> V | _DefaultT | N
             The result value.
 
         """
-        try:
-            return self[key]
-        except KeyError:
+        ret = _ffi_api.MapGetItemOrMissing(self, key)
+        if MISSING.same_as(ret):
             return default
+        return ret
 
     def __repr__(self) -> str:
         """Return a string representation of the map."""

src/ffi/container.cc

@@ -55,6 +55,11 @@ class MapForwardIterFunctor {
   ffi::MapObj::iterator end_;
 };
 
+ObjectRef GetMissingObject() {
+  static ObjectRef missing_obj(make_object<Object>());
+  return missing_obj;
+}
+
 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef()
@@ -81,8 +86,17 @@ TVM_FFI_STATIC_INIT_BLOCK() {
            [](const ffi::MapObj* n, const Any& k) -> int64_t {
              return static_cast<int64_t>(n->count(k));
            })
-      .def("ffi.MapForwardIterFunctor", [](const ffi::MapObj* n) -> ffi::Function {
-        return ffi::Function::FromTyped(MapForwardIterFunctor(n->begin(), n->end()));
+      .def("ffi.MapForwardIterFunctor",
+           [](const ffi::MapObj* n) -> ffi::Function {
+             return ffi::Function::FromTyped(MapForwardIterFunctor(n->begin(), n->end()));
+           })
+      .def("ffi.MapGetMissingObject", GetMissingObject)
+      .def("ffi.MapGetItemOrMissing", [](const ffi::MapObj* n, const Any& k) -> Any {
+        try {
+          return n->at(k);
+        } catch (const tvm::ffi::Error& e) {
+          return GetMissingObject();
+        }
       });
 }
 }  // namespace ffi