perf: do not mark task inputs/outputs as multi-threaded unless necessary

src/runtime/io.cpp

@@ -707,19 +707,6 @@ static inline atomic<object*> * mt_ref_val_addr(object * o) {
     return reinterpret_cast<atomic<object*> *>(&(lean_to_ref(o)->m_value));
 }
 
-/*
-  Important: we have added support for initializing global constants
-  at program startup. This feature is particularly useful for
-  initializing `ST.Ref` values. Any `ST.Ref` value created during
-  initialization will be marked as persistent. Thus, to make `ST.Ref`
-  API thread-safe, we must treat persistent `ST.Ref` objects created
-  during initialization as a multi-threaded object. Then, whenever we store
-  a value `val` into a global `ST.Ref`, we have to mark `va`l as a multi-threaded
-  object as we do for multi-threaded `ST.Ref`s. It makes sense since
-  the global `ST.Ref` may be used to communicate data between threads.
-*/
-static inline bool ref_maybe_mt(b_obj_arg ref) { return lean_is_mt(ref) || lean_is_persistent(ref); }
-
 extern "C" LEAN_EXPORT obj_res lean_st_ref_get(b_obj_arg ref, obj_arg) {
     if (ref_maybe_mt(ref)) {
         atomic<object *> * val_addr = mt_ref_val_addr(ref);

src/runtime/object.cpp

@@ -772,7 +772,9 @@ class task_manager {
         } else if (v != nullptr) {
             lean_assert(t->m_imp->m_closure == nullptr);
             handle_finished(t);
-            mark_mt(v);
+            if (ref_maybe_mt((lean_object *)t)) {
+                mark_mt(v);
+            }
             t->m_value = v;
             /* After the task has been finished and we propagated
                dependecies, we can release `m_imp` and keep just the value */
@@ -941,7 +943,9 @@ static inline void lean_set_task_header(lean_object * o) {
 }
 
 static lean_task_object * alloc_task(obj_arg c, unsigned prio, bool keep_alive) {
-    lean_mark_mt(c);
+    if (!lean_is_exclusive(c)) {
+        lean_mark_mt(c);
+    }
     lean_task_object * o = (lean_task_object*)lean_alloc_small_object(sizeof(lean_task_object));
     lean_set_task_header((lean_object*)o);
     o->m_value = nullptr;
@@ -1018,9 +1022,8 @@ static obj_res task_bind_fn1(obj_arg x, obj_arg f, obj_arg) {
     lean_assert(lean_is_task(new_task));
     lean_assert(g_current_task_object->m_imp);
     lean_assert(g_current_task_object->m_imp->m_closure == nullptr);
-    obj_res c = mk_closure_2_1(task_bind_fn2, new_task);
-    mark_mt(c);
-    g_current_task_object->m_imp->m_closure = c;
+    // NOTE: closure is unique, so no need to mark multi-threaded
+    g_current_task_object->m_imp->m_closure = mk_closure_2_1(task_bind_fn2, new_task);
     return nullptr; /* notify queue that task did not finish yet. */
 }
 

src/runtime/object.h

@@ -42,6 +42,21 @@ inline bool is_scalar(object * o) { return lean_is_scalar(o); }
 inline object * box(size_t n) { return lean_box(n); }
 inline size_t unbox(object * o) { return lean_unbox(o); }
 
+/*
+  Important: we have added support for initializing global constants
+  at program startup. This feature is particularly useful for
+  initializing `ST.Ref` values. Any `ST.Ref` value created during
+  initialization will be marked as persistent. Thus, to make `ST.Ref`
+  API thread-safe, we must treat persistent `ST.Ref` objects created
+  during initialization as a multi-threaded object. Then, whenever we store
+  a value `val` into a global `ST.Ref`, we have to mark `va`l as a multi-threaded
+  object as we do for multi-threaded `ST.Ref`s. It makes sense since
+  the global `ST.Ref` may be used to communicate data between threads.
+  Similar issues can happen with other types with interior mutability
+  such as `Task`s.
+*/
+inline bool ref_maybe_mt(b_obj_arg ref) { return lean_is_mt(ref) || lean_is_persistent(ref); }
+
 inline bool is_mt_heap_obj(object * o) { return lean_is_mt(o); }
 inline bool is_st_heap_obj(object * o) { return lean_is_st(o); }
 inline bool is_heap_obj(object * o) { return is_st_heap_obj(o) || is_mt_heap_obj(o); }