vfdev-5/CMakeLists.txt Secret

## readme.md

      
    Raw
  

              readme.md
            
          
    cmake -S. -Bbuild \
    -DCMAKE_BUILD_TYPE=Debug \
    -DPython_EXECUTABLE=$(which python3.14t) \
    -DUSE_TSAN=ON \
    -DCMAKE_C_COMPILER=clang-18 \
    -DCMAKE_CXX_COMPILER=clang++-18


cmake --build build


PYTHONPATH=build python3.14t check.py

  
## binding.cpp
#include <Python.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <exception>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "nanobind/nanobind.h"
#include "nanobind/stl/shared_ptr.h"  // IWYU pragma: keep
#include "nanobind/stl/string.h"  // IWYU pragma: keep
#include "nanobind/stl/vector.h"  // IWYU pragma: keep


namespace nb = nanobind;

namespace {

class PmapFunction {
public:
    PmapFunction(nb::callable fun) : fun_(std::move(fun)) {
        function_name_ = nb::cast<std::string>(nb::str(nb::getattr(fun_, "__name__", fun_)));
    }
    PmapFunction(const PmapFunction&) = delete;
    PmapFunction& operator=(const PmapFunction& other) = delete;

    const nb::callable& fun() const { return fun_; }
    const std::string& function_name() const { return function_name_; }

 private:
    nb::callable fun_;  // The Python function to pmap.
    std::string function_name_;
};


struct JaxPmapFunctionObject {
  PyObject_HEAD;
  PyObject* dict;      // Dictionary for __dict__
  PyObject* weakrefs;  // Weak references; for use by the Python interpreter.
  vectorcallfunc vectorcall;
  PmapFunction fun;
};

PyObject* JaxPmapFunction_Type = nullptr;

extern "C" {


PyObject* JaxPmapFunction_tp_new(PyTypeObject* subtype, PyObject* args,
                                 PyObject* kwds) {
  JaxPmapFunctionObject* self =
      reinterpret_cast<JaxPmapFunctionObject*>(subtype->tp_alloc(subtype, 0));
  if (!self) return nullptr;
  self->dict = nullptr;
  self->weakrefs = nullptr;
  return reinterpret_cast<PyObject*>(self);
}

void JaxPmapFunction_tp_dealloc(PyObject* self) {
  PyObject_GC_UnTrack(self);
  PyTypeObject* tp = Py_TYPE(self);
  JaxPmapFunctionObject* o = reinterpret_cast<JaxPmapFunctionObject*>(self);
  if (o->weakrefs) {
    PyObject_ClearWeakRefs(self);
  }
  Py_CLEAR(o->dict);
  o->fun.~PmapFunction();
  tp->tp_free(self);
  Py_DECREF(tp);
}

}  // extern "C"


nb::object MakePmapFunction(nb::callable fun) {
  nb::object obj = nb::steal<nb::object>(JaxPmapFunction_tp_new(
      reinterpret_cast<PyTypeObject*>(JaxPmapFunction_Type), nullptr, nullptr));
  JaxPmapFunctionObject* buf =
      reinterpret_cast<JaxPmapFunctionObject*>(obj.ptr());
  new (&buf->fun) PmapFunction(std::move(fun));
  return obj;
}

void BuildPmapSubmodule(nb::module_& m) {

    nb::module_ pmap_lib = m.def_submodule("pmap_lib", "Jax C++ pmap library");

    // We need to use heap-allocated type objects because we want to add
    // additional methods dynamically.
    nb::object cfun;
    {
        nb::str name = nb::str("PmapFunction");
        nb::str qualname = nb::str("PmapFunction");
        PyHeapTypeObject* heap_type = reinterpret_cast<PyHeapTypeObject*>(
            PyType_Type.tp_alloc(&PyType_Type, 0));
        // Caution: we must not call any functions that might invoke the GC until
        // PyType_Ready() is called. Otherwise the GC might see a half-constructed
        // type object.
        assert(heap_type);
        heap_type->ht_name = name.release().ptr();
        heap_type->ht_qualname = qualname.release().ptr();
        PyTypeObject* type = &heap_type->ht_type;
        type->tp_name = "PmapFunction";
        type->tp_basicsize = sizeof(JaxPmapFunctionObject);
        type->tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HEAPTYPE;
        type->tp_new = JaxPmapFunction_tp_new;
        type->tp_dealloc = JaxPmapFunction_tp_dealloc;
        type->tp_dictoffset = offsetof(JaxPmapFunctionObject, dict);
        assert(PyType_Ready(type) == 0);
        JaxPmapFunction_Type = reinterpret_cast<PyObject*>(type);
        cfun = nb::borrow<nb::object>(JaxPmapFunction_Type);
    }
    nb::object cfun_type = nb::borrow<nb::object>(JaxPmapFunction_Type);

    // Add PmapFunction to the xla_extension module so it can be pickled.
    m.attr("PmapFunction") = cfun_type;

    pmap_lib.def(
        "pmap",
        [](nb::callable fun) -> nb::object {
            return MakePmapFunction(std::move(fun));
        },
        nb::arg("fun"));
}

NB_MODULE(fake_pmap_ext, m) {
    BuildPmapSubmodule(m);
}

}

## check.log
==================
WARNING: ThreadSanitizer: data race (pid=276288)
  Atomic read of size 8 at 0x7fffb5778000 by thread T9:
    #0 _Py_atomic_load_uintptr_relaxed(unsigned long const*) /tmp/output-python/include/python3.14t/cpython/pyatomic_gcc.h:375:10 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7) (BuildId: 59a98a527d0b7f6e107e3fd69a5bde83cd34f3cd)
    #1 _Py_IsOwnedByCurrentThread(_object*) /tmp/output-python/include/python3.14t/object.h:252:12 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7)
    #2 Py_DECREF(_object*) /tmp/output-python/include/python3.14t/refcount.h:361:9 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7)
    #3 JaxPmapFunction_tp_dealloc /project/playground/check_jax/pmap/binding.cpp:75:3 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7)
    #4 _Py_Dealloc /project/cpython/Objects/object.c:2991:5 (python3.14t+0x2944a2) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #5 _Py_MergeZeroLocalRefcount /project/cpython/Objects/object.c (python3.14t+0x2944a2)
    #6 Py_DECREF /project/cpython/./Include/refcount.h:365:13 (python3.14t+0x45953e) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #7 _PyFrame_ClearLocals /project/cpython/Python/frame.c:99:9 (python3.14t+0x45953e)
    #8 _PyFrame_ClearExceptCode /project/cpython/Python/frame.c:124:5 (python3.14t+0x459762) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #9 clear_thread_frame /project/cpython/Python/ceval.c:1733:5 (python3.14t+0x411b77) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #10 _PyEval_FrameClearAndPop /project/cpython/Python/ceval.c:1760:9 (python3.14t+0x411b77)
    #11 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:7103:13 (python3.14t+0x40a479) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #12 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #13 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
    #14 _PyFunction_Vectorcall /project/cpython/Objects/call.c (python3.14t+0x1e99df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #15 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:167:11 (python3.14t+0x1ee17f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #16 method_vectorcall /project/cpython/Objects/classobject.c:71:20 (python3.14t+0x1ee17f)
    #17 _PyVectorcall_Call /project/cpython/Objects/call.c:273:16 (python3.14t+0x1e9653) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #18 _PyObject_Call /project/cpython/Objects/call.c:348:16 (python3.14t+0x1e9653)
    #19 PyObject_Call /project/cpython/Objects/call.c:373:12 (python3.14t+0x1e96d5) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #20 thread_run /project/cpython/./Modules/_threadmodule.c:346:21 (python3.14t+0x581662) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #21 pythread_wrapper /project/cpython/Python/thread_pthread.h:242:5 (python3.14t+0x4dc607) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

  Previous write of size 8 at 0x7fffb5778000 by thread T8:
    #0 mi_block_set_nextx /project/cpython/./Include/internal/mimalloc/mimalloc/internal.h:652:15 (python3.14t+0x2c5f4e) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #1 _mi_free_block_mt /project/cpython/Objects/mimalloc/alloc.c:467:9 (python3.14t+0x2c5f4e)
    #2 _mi_free_block /project/cpython/Objects/mimalloc/alloc.c:506:5 (python3.14t+0x2a039a) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #3 _mi_free_generic /project/cpython/Objects/mimalloc/alloc.c:524:3 (python3.14t+0x2a039a)
    #4 mi_free /project/cpython/Objects/mimalloc/alloc.c (python3.14t+0x2bee7b) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #5 _PyObject_MiFree /project/cpython/Objects/obmalloc.c:284:5 (python3.14t+0x2bee7b)
    #6 PyObject_Free /project/cpython/Objects/obmalloc.c:1368:5 (python3.14t+0x2c1bb4) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #7 PyObject_GC_Del /project/cpython/Python/gc_free_threading.c:1925:5 (python3.14t+0x45e16d) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #8 type_dealloc /project/cpython/Objects/typeobject.c:6156:5 (python3.14t+0x2f0f68) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #9 _Py_Dealloc /project/cpython/Objects/object.c:2991:5 (python3.14t+0x2943ab) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #10 _Py_DecRefSharedDebug /project/cpython/Objects/object.c:418:9 (python3.14t+0x2943ab)
    #11 _Py_DecRefShared /project/cpython/Objects/object.c:425:5 (python3.14t+0x2943ab)
    #12 Py_DECREF(_object*) /tmp/output-python/include/python3.14t/refcount.h:369:9 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x700d) (BuildId: 59a98a527d0b7f6e107e3fd69a5bde83cd34f3cd)
    #13 JaxPmapFunction_tp_dealloc /project/playground/check_jax/pmap/binding.cpp:75:3 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x700d)
    #14 _Py_Dealloc /project/cpython/Objects/object.c:2991:5 (python3.14t+0x2944a2) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #15 _Py_MergeZeroLocalRefcount /project/cpython/Objects/object.c (python3.14t+0x2944a2)
    #16 Py_DECREF /project/cpython/./Include/refcount.h:365:13 (python3.14t+0x45953e) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #17 _PyFrame_ClearLocals /project/cpython/Python/frame.c:99:9 (python3.14t+0x45953e)
    #18 _PyFrame_ClearExceptCode /project/cpython/Python/frame.c:124:5 (python3.14t+0x459762) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #19 clear_thread_frame /project/cpython/Python/ceval.c:1733:5 (python3.14t+0x411b77) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #20 _PyEval_FrameClearAndPop /project/cpython/Python/ceval.c:1760:9 (python3.14t+0x411b77)
    #21 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:7103:13 (python3.14t+0x40a479) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #22 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #23 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
    #24 _PyFunction_Vectorcall /project/cpython/Objects/call.c (python3.14t+0x1e99df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #25 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:167:11 (python3.14t+0x1ee17f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #26 method_vectorcall /project/cpython/Objects/classobject.c:71:20 (python3.14t+0x1ee17f)
    #27 _PyVectorcall_Call /project/cpython/Objects/call.c:273:16 (python3.14t+0x1e9653) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #28 _PyObject_Call /project/cpython/Objects/call.c:348:16 (python3.14t+0x1e9653)
    #29 PyObject_Call /project/cpython/Objects/call.c:373:12 (python3.14t+0x1e96d5) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #30 thread_run /project/cpython/./Modules/_threadmodule.c:346:21 (python3.14t+0x581662) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #31 pythread_wrapper /project/cpython/Python/thread_pthread.h:242:5 (python3.14t+0x4dc607) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

  Thread T9 'ThreadPoolExecu' (tid=276298, running) created by main thread at:
    #0 pthread_create <null> (python3.14t+0xe01df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #1 do_start_joinable_thread /project/cpython/Python/thread_pthread.h:289:14 (python3.14t+0x4db498) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #2 PyThread_start_joinable_thread /project/cpython/Python/thread_pthread.h:313:9 (python3.14t+0x4db2ba) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #3 ThreadHandle_start /project/cpython/./Modules/_threadmodule.c:431:9 (python3.14t+0x5811f7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #4 do_start_new_thread /project/cpython/./Modules/_threadmodule.c:1798:9 (python3.14t+0x5811f7)
    #5 thread_PyThread_start_joinable_thread /project/cpython/./Modules/_threadmodule.c:1921:14 (python3.14t+0x57ff51) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #6 cfunction_call /project/cpython/Objects/methodobject.c:551:18 (python3.14t+0x28dee7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #7 _PyObject_MakeTpCall /project/cpython/Objects/call.c:242:18 (python3.14t+0x1e87cc) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #8 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:165:16 (python3.14t+0x1e9428) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #9 PyObject_Vectorcall /project/cpython/Objects/call.c:327:12 (python3.14t+0x1e9428)
    #10 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:2215:35 (python3.14t+0x3f6bc8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #11 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #12 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
    #13 PyEval_EvalCode /project/cpython/Python/ceval.c:658:21 (python3.14t+0x3eb1d6) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #14 run_eval_code_obj /project/cpython/Python/pythonrun.c:1338:9 (python3.14t+0x4bcaee) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #15 run_mod /project/cpython/Python/pythonrun.c:1423:19 (python3.14t+0x4bc215) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #16 pyrun_file /project/cpython/Python/pythonrun.c:1256:15 (python3.14t+0x4b8310) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #17 _PyRun_SimpleFileObject /project/cpython/Python/pythonrun.c:491:13 (python3.14t+0x4b8310)
    #18 _PyRun_AnyFileObject /project/cpython/Python/pythonrun.c:78:15 (python3.14t+0x4b79d8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #19 pymain_run_file_obj /project/cpython/Modules/main.c:410:15 (python3.14t+0x4f771f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #20 pymain_run_file /project/cpython/Modules/main.c:429:15 (python3.14t+0x4f771f)
    #21 pymain_run_python /project/cpython/Modules/main.c:697:21 (python3.14t+0x4f696c) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #22 Py_RunMain /project/cpython/Modules/main.c:776:5 (python3.14t+0x4f696c)
    #23 pymain_main /project/cpython/Modules/main.c:806:12 (python3.14t+0x4f6da8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #24 Py_BytesMain /project/cpython/Modules/main.c:830:12 (python3.14t+0x4f6e2b) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #25 main /project/cpython/./Programs/python.c:15:12 (python3.14t+0x15e7eb) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

  Thread T8 'ThreadPoolExecu' (tid=276297, running) created by main thread at:
    #0 pthread_create <null> (python3.14t+0xe01df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #1 do_start_joinable_thread /project/cpython/Python/thread_pthread.h:289:14 (python3.14t+0x4db498) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #2 PyThread_start_joinable_thread /project/cpython/Python/thread_pthread.h:313:9 (python3.14t+0x4db2ba) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #3 ThreadHandle_start /project/cpython/./Modules/_threadmodule.c:431:9 (python3.14t+0x5811f7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #4 do_start_new_thread /project/cpython/./Modules/_threadmodule.c:1798:9 (python3.14t+0x5811f7)
    #5 thread_PyThread_start_joinable_thread /project/cpython/./Modules/_threadmodule.c:1921:14 (python3.14t+0x57ff51) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #6 cfunction_call /project/cpython/Objects/methodobject.c:551:18 (python3.14t+0x28dee7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #7 _PyObject_MakeTpCall /project/cpython/Objects/call.c:242:18 (python3.14t+0x1e87cc) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #8 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:165:16 (python3.14t+0x1e9428) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #9 PyObject_Vectorcall /project/cpython/Objects/call.c:327:12 (python3.14t+0x1e9428)
    #10 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:2215:35 (python3.14t+0x3f6bc8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #11 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #12 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
    #13 PyEval_EvalCode /project/cpython/Python/ceval.c:658:21 (python3.14t+0x3eb1d6) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #14 run_eval_code_obj /project/cpython/Python/pythonrun.c:1338:9 (python3.14t+0x4bcaee) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #15 run_mod /project/cpython/Python/pythonrun.c:1423:19 (python3.14t+0x4bc215) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #16 pyrun_file /project/cpython/Python/pythonrun.c:1256:15 (python3.14t+0x4b8310) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #17 _PyRun_SimpleFileObject /project/cpython/Python/pythonrun.c:491:13 (python3.14t+0x4b8310)
    #18 _PyRun_AnyFileObject /project/cpython/Python/pythonrun.c:78:15 (python3.14t+0x4b79d8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #19 pymain_run_file_obj /project/cpython/Modules/main.c:410:15 (python3.14t+0x4f771f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #20 pymain_run_file /project/cpython/Modules/main.c:429:15 (python3.14t+0x4f771f)
    #21 pymain_run_python /project/cpython/Modules/main.c:697:21 (python3.14t+0x4f696c) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #22 Py_RunMain /project/cpython/Modules/main.c:776:5 (python3.14t+0x4f696c)
    #23 pymain_main /project/cpython/Modules/main.c:806:12 (python3.14t+0x4f6da8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #24 Py_BytesMain /project/cpython/Modules/main.c:830:12 (python3.14t+0x4f6e2b) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
    #25 main /project/cpython/./Programs/python.c:15:12 (python3.14t+0x15e7eb) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

SUMMARY: ThreadSanitizer: data race /tmp/output-python/include/python3.14t/cpython/pyatomic_gcc.h:375:10 in _Py_atomic_load_uintptr_relaxed(unsigned long const*)
==================
ThreadSanitizer: reported 1 warnings

## check.py
import math
import numpy as np

import concurrent.futures
import threading

from fake_pmap_ext import pmap_lib


if __name__ == "__main__":
    num_workers = 40

    barrier = threading.Barrier(num_workers)

    def _pmap(fun):
        cpp_mapped_f = pmap_lib.pmap(fun)

        return cpp_mapped_f

    def closure():
        barrier.wait()

        in_shape = (1, 2, 4)
        x = np.arange(math.prod(in_shape)).reshape(in_shape)
        fn = _pmap(lambda x: x)

    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = []
        for i in range(num_workers):
            futures.append(executor.submit(closure))
        assert len(list(f.result() for f in futures)) == num_workers

## CMakeLists.txt
cmake_minimum_required(VERSION 3.16)

project(fake_pmap_ext_project)

# https://nanobind.readthedocs.io/en/latest/building.html


if (CMAKE_VERSION VERSION_LESS 3.18)
  set(DEV_MODULE Development)
else()
  set(DEV_MODULE Development.Module)
endif()

find_package(Python 3.8 COMPONENTS Interpreter ${DEV_MODULE} REQUIRED)

if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
  set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
endif()

option(USE_TSAN "Compile with TSAN" ON)

function(append value)
  foreach(variable ${ARGN})
    set(${variable} "${${variable}} ${value}" PARENT_SCOPE)
  endforeach(variable)
endfunction()

if (USE_TSAN)
  append("-fsanitize=thread -g -O2" CMAKE_C_FLAGS CMAKE_CXX_FLAGS)
  append("-fsanitize=thread" CMAKE_EXE_LINKER_FLAGS CMAKE_MODULE_LINKER_FLAGS CMAKE_SHARED_LINKER_FLAGS)
endif()

# Detect the installed nanobind package and import it into CMake
execute_process(
  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)

find_package(nanobind CONFIG REQUIRED)

# Abseil requires C++14
set(CMAKE_CXX_STANDARD 14)

nanobind_add_module(
    fake_pmap_ext
    FREE_THREADED
    binding.cpp
)

set_target_properties(fake_pmap_ext PROPERTIES POSITION_INDEPENDENT_CODE ON)
# target_link_libraries(fake_pmap_ext PRIVATE absl::strings absl::node_hash_map absl::base absl::synchronization absl::cleanup)
	#include <Python.h>

	#include <algorithm>
	#include <cstddef>
	#include <cstdint>
	#include <exception>
	#include <memory>
	#include <optional>
	#include <stdexcept>
	#include <string>
	#include <tuple>
	#include <utility>
	#include <vector>

	#include "nanobind/nanobind.h"
	#include "nanobind/stl/shared_ptr.h" // IWYU pragma: keep
	#include "nanobind/stl/string.h" // IWYU pragma: keep
	#include "nanobind/stl/vector.h" // IWYU pragma: keep


	namespace nb = nanobind;

	namespace {

	class PmapFunction {
	public:
	PmapFunction(nb::callable fun) : fun_(std::move(fun)) {
	function_name_ = nb::cast<std::string>(nb::str(nb::getattr(fun_, "__name__", fun_)));
	}
	PmapFunction(const PmapFunction&) = delete;
	PmapFunction& operator=(const PmapFunction& other) = delete;

	const nb::callable& fun() const { return fun_; }
	const std::string& function_name() const { return function_name_; }

	private:
	nb::callable fun_; // The Python function to pmap.
	std::string function_name_;
	};


	struct JaxPmapFunctionObject {
	PyObject_HEAD;
	PyObject* dict; // Dictionary for __dict__
	PyObject* weakrefs; // Weak references; for use by the Python interpreter.
	vectorcallfunc vectorcall;
	PmapFunction fun;
	};

	PyObject* JaxPmapFunction_Type = nullptr;

	extern "C" {


	PyObject* JaxPmapFunction_tp_new(PyTypeObject* subtype, PyObject* args,
	PyObject* kwds) {
	JaxPmapFunctionObject* self =
	reinterpret_cast<JaxPmapFunctionObject*>(subtype->tp_alloc(subtype, 0));
	if (!self) return nullptr;
	self->dict = nullptr;
	self->weakrefs = nullptr;
	return reinterpret_cast<PyObject*>(self);
	}

	void JaxPmapFunction_tp_dealloc(PyObject* self) {
	PyObject_GC_UnTrack(self);
	PyTypeObject* tp = Py_TYPE(self);
	JaxPmapFunctionObject* o = reinterpret_cast<JaxPmapFunctionObject*>(self);
	if (o->weakrefs) {
	PyObject_ClearWeakRefs(self);
	}
	Py_CLEAR(o->dict);
	o->fun.~PmapFunction();
	tp->tp_free(self);
	Py_DECREF(tp);
	}

	} // extern "C"


	nb::object MakePmapFunction(nb::callable fun) {
	nb::object obj = nb::steal<nb::object>(JaxPmapFunction_tp_new(
	reinterpret_cast<PyTypeObject*>(JaxPmapFunction_Type), nullptr, nullptr));
	JaxPmapFunctionObject* buf =
	reinterpret_cast<JaxPmapFunctionObject*>(obj.ptr());
	new (&buf->fun) PmapFunction(std::move(fun));
	return obj;
	}

	void BuildPmapSubmodule(nb::module_& m) {

	nb::module_ pmap_lib = m.def_submodule("pmap_lib", "Jax C++ pmap library");

	// We need to use heap-allocated type objects because we want to add
	// additional methods dynamically.
	nb::object cfun;
	{
	nb::str name = nb::str("PmapFunction");
	nb::str qualname = nb::str("PmapFunction");
	PyHeapTypeObject* heap_type = reinterpret_cast<PyHeapTypeObject*>(
	PyType_Type.tp_alloc(&PyType_Type, 0));
	// Caution: we must not call any functions that might invoke the GC until
	// PyType_Ready() is called. Otherwise the GC might see a half-constructed
	// type object.
	assert(heap_type);
	heap_type->ht_name = name.release().ptr();
	heap_type->ht_qualname = qualname.release().ptr();
	PyTypeObject* type = &heap_type->ht_type;
	type->tp_name = "PmapFunction";
	type->tp_basicsize = sizeof(JaxPmapFunctionObject);
	type->tp_flags = Py_TPFLAGS_DEFAULT \| Py_TPFLAGS_HEAPTYPE;
	type->tp_new = JaxPmapFunction_tp_new;
	type->tp_dealloc = JaxPmapFunction_tp_dealloc;
	type->tp_dictoffset = offsetof(JaxPmapFunctionObject, dict);
	assert(PyType_Ready(type) == 0);
	JaxPmapFunction_Type = reinterpret_cast<PyObject*>(type);
	cfun = nb::borrow<nb::object>(JaxPmapFunction_Type);
	}
	nb::object cfun_type = nb::borrow<nb::object>(JaxPmapFunction_Type);

	// Add PmapFunction to the xla_extension module so it can be pickled.
	m.attr("PmapFunction") = cfun_type;

	pmap_lib.def(
	"pmap",
	[](nb::callable fun) -> nb::object {
	return MakePmapFunction(std::move(fun));
	},
	nb::arg("fun"));
	}

	NB_MODULE(fake_pmap_ext, m) {
	BuildPmapSubmodule(m);
	}

	}
	==================
	WARNING: ThreadSanitizer: data race (pid=276288)
	Atomic read of size 8 at 0x7fffb5778000 by thread T9:
	#0 _Py_atomic_load_uintptr_relaxed(unsigned long const*) /tmp/output-python/include/python3.14t/cpython/pyatomic_gcc.h:375:10 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7) (BuildId: 59a98a527d0b7f6e107e3fd69a5bde83cd34f3cd)
	#1 _Py_IsOwnedByCurrentThread(_object*) /tmp/output-python/include/python3.14t/object.h:252:12 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7)
	#2 Py_DECREF(_object*) /tmp/output-python/include/python3.14t/refcount.h:361:9 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7)
	#3 JaxPmapFunction_tp_dealloc /project/playground/check_jax/pmap/binding.cpp:75:3 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x6fd7)
	#4 _Py_Dealloc /project/cpython/Objects/object.c:2991:5 (python3.14t+0x2944a2) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#5 _Py_MergeZeroLocalRefcount /project/cpython/Objects/object.c (python3.14t+0x2944a2)
	#6 Py_DECREF /project/cpython/./Include/refcount.h:365:13 (python3.14t+0x45953e) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#7 _PyFrame_ClearLocals /project/cpython/Python/frame.c:99:9 (python3.14t+0x45953e)
	#8 _PyFrame_ClearExceptCode /project/cpython/Python/frame.c:124:5 (python3.14t+0x459762) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#9 clear_thread_frame /project/cpython/Python/ceval.c:1733:5 (python3.14t+0x411b77) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#10 _PyEval_FrameClearAndPop /project/cpython/Python/ceval.c:1760:9 (python3.14t+0x411b77)
	#11 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:7103:13 (python3.14t+0x40a479) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#12 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#13 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
	#14 _PyFunction_Vectorcall /project/cpython/Objects/call.c (python3.14t+0x1e99df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#15 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:167:11 (python3.14t+0x1ee17f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#16 method_vectorcall /project/cpython/Objects/classobject.c:71:20 (python3.14t+0x1ee17f)
	#17 _PyVectorcall_Call /project/cpython/Objects/call.c:273:16 (python3.14t+0x1e9653) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#18 _PyObject_Call /project/cpython/Objects/call.c:348:16 (python3.14t+0x1e9653)
	#19 PyObject_Call /project/cpython/Objects/call.c:373:12 (python3.14t+0x1e96d5) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#20 thread_run /project/cpython/./Modules/_threadmodule.c:346:21 (python3.14t+0x581662) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#21 pythread_wrapper /project/cpython/Python/thread_pthread.h:242:5 (python3.14t+0x4dc607) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

	Previous write of size 8 at 0x7fffb5778000 by thread T8:
	#0 mi_block_set_nextx /project/cpython/./Include/internal/mimalloc/mimalloc/internal.h:652:15 (python3.14t+0x2c5f4e) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#1 _mi_free_block_mt /project/cpython/Objects/mimalloc/alloc.c:467:9 (python3.14t+0x2c5f4e)
	#2 _mi_free_block /project/cpython/Objects/mimalloc/alloc.c:506:5 (python3.14t+0x2a039a) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#3 _mi_free_generic /project/cpython/Objects/mimalloc/alloc.c:524:3 (python3.14t+0x2a039a)
	#4 mi_free /project/cpython/Objects/mimalloc/alloc.c (python3.14t+0x2bee7b) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#5 _PyObject_MiFree /project/cpython/Objects/obmalloc.c:284:5 (python3.14t+0x2bee7b)
	#6 PyObject_Free /project/cpython/Objects/obmalloc.c:1368:5 (python3.14t+0x2c1bb4) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#7 PyObject_GC_Del /project/cpython/Python/gc_free_threading.c:1925:5 (python3.14t+0x45e16d) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#8 type_dealloc /project/cpython/Objects/typeobject.c:6156:5 (python3.14t+0x2f0f68) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#9 _Py_Dealloc /project/cpython/Objects/object.c:2991:5 (python3.14t+0x2943ab) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#10 _Py_DecRefSharedDebug /project/cpython/Objects/object.c:418:9 (python3.14t+0x2943ab)
	#11 _Py_DecRefShared /project/cpython/Objects/object.c:425:5 (python3.14t+0x2943ab)
	#12 Py_DECREF(_object*) /tmp/output-python/include/python3.14t/refcount.h:369:9 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x700d) (BuildId: 59a98a527d0b7f6e107e3fd69a5bde83cd34f3cd)
	#13 JaxPmapFunction_tp_dealloc /project/playground/check_jax/pmap/binding.cpp:75:3 (fake_pmap_ext.cpython-314t-x86_64-linux-gnu.so+0x700d)
	#14 _Py_Dealloc /project/cpython/Objects/object.c:2991:5 (python3.14t+0x2944a2) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#15 _Py_MergeZeroLocalRefcount /project/cpython/Objects/object.c (python3.14t+0x2944a2)
	#16 Py_DECREF /project/cpython/./Include/refcount.h:365:13 (python3.14t+0x45953e) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#17 _PyFrame_ClearLocals /project/cpython/Python/frame.c:99:9 (python3.14t+0x45953e)
	#18 _PyFrame_ClearExceptCode /project/cpython/Python/frame.c:124:5 (python3.14t+0x459762) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#19 clear_thread_frame /project/cpython/Python/ceval.c:1733:5 (python3.14t+0x411b77) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#20 _PyEval_FrameClearAndPop /project/cpython/Python/ceval.c:1760:9 (python3.14t+0x411b77)
	#21 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:7103:13 (python3.14t+0x40a479) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#22 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#23 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
	#24 _PyFunction_Vectorcall /project/cpython/Objects/call.c (python3.14t+0x1e99df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#25 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:167:11 (python3.14t+0x1ee17f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#26 method_vectorcall /project/cpython/Objects/classobject.c:71:20 (python3.14t+0x1ee17f)
	#27 _PyVectorcall_Call /project/cpython/Objects/call.c:273:16 (python3.14t+0x1e9653) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#28 _PyObject_Call /project/cpython/Objects/call.c:348:16 (python3.14t+0x1e9653)
	#29 PyObject_Call /project/cpython/Objects/call.c:373:12 (python3.14t+0x1e96d5) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#30 thread_run /project/cpython/./Modules/_threadmodule.c:346:21 (python3.14t+0x581662) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#31 pythread_wrapper /project/cpython/Python/thread_pthread.h:242:5 (python3.14t+0x4dc607) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

	Thread T9 'ThreadPoolExecu' (tid=276298, running) created by main thread at:
	#0 pthread_create <null> (python3.14t+0xe01df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#1 do_start_joinable_thread /project/cpython/Python/thread_pthread.h:289:14 (python3.14t+0x4db498) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#2 PyThread_start_joinable_thread /project/cpython/Python/thread_pthread.h:313:9 (python3.14t+0x4db2ba) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#3 ThreadHandle_start /project/cpython/./Modules/_threadmodule.c:431:9 (python3.14t+0x5811f7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#4 do_start_new_thread /project/cpython/./Modules/_threadmodule.c:1798:9 (python3.14t+0x5811f7)
	#5 thread_PyThread_start_joinable_thread /project/cpython/./Modules/_threadmodule.c:1921:14 (python3.14t+0x57ff51) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#6 cfunction_call /project/cpython/Objects/methodobject.c:551:18 (python3.14t+0x28dee7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#7 _PyObject_MakeTpCall /project/cpython/Objects/call.c:242:18 (python3.14t+0x1e87cc) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#8 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:165:16 (python3.14t+0x1e9428) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#9 PyObject_Vectorcall /project/cpython/Objects/call.c:327:12 (python3.14t+0x1e9428)
	#10 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:2215:35 (python3.14t+0x3f6bc8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#11 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#12 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
	#13 PyEval_EvalCode /project/cpython/Python/ceval.c:658:21 (python3.14t+0x3eb1d6) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#14 run_eval_code_obj /project/cpython/Python/pythonrun.c:1338:9 (python3.14t+0x4bcaee) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#15 run_mod /project/cpython/Python/pythonrun.c:1423:19 (python3.14t+0x4bc215) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#16 pyrun_file /project/cpython/Python/pythonrun.c:1256:15 (python3.14t+0x4b8310) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#17 _PyRun_SimpleFileObject /project/cpython/Python/pythonrun.c:491:13 (python3.14t+0x4b8310)
	#18 _PyRun_AnyFileObject /project/cpython/Python/pythonrun.c:78:15 (python3.14t+0x4b79d8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#19 pymain_run_file_obj /project/cpython/Modules/main.c:410:15 (python3.14t+0x4f771f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#20 pymain_run_file /project/cpython/Modules/main.c:429:15 (python3.14t+0x4f771f)
	#21 pymain_run_python /project/cpython/Modules/main.c:697:21 (python3.14t+0x4f696c) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#22 Py_RunMain /project/cpython/Modules/main.c:776:5 (python3.14t+0x4f696c)
	#23 pymain_main /project/cpython/Modules/main.c:806:12 (python3.14t+0x4f6da8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#24 Py_BytesMain /project/cpython/Modules/main.c:830:12 (python3.14t+0x4f6e2b) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#25 main /project/cpython/./Programs/python.c:15:12 (python3.14t+0x15e7eb) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

	Thread T8 'ThreadPoolExecu' (tid=276297, running) created by main thread at:
	#0 pthread_create <null> (python3.14t+0xe01df) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#1 do_start_joinable_thread /project/cpython/Python/thread_pthread.h:289:14 (python3.14t+0x4db498) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#2 PyThread_start_joinable_thread /project/cpython/Python/thread_pthread.h:313:9 (python3.14t+0x4db2ba) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#3 ThreadHandle_start /project/cpython/./Modules/_threadmodule.c:431:9 (python3.14t+0x5811f7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#4 do_start_new_thread /project/cpython/./Modules/_threadmodule.c:1798:9 (python3.14t+0x5811f7)
	#5 thread_PyThread_start_joinable_thread /project/cpython/./Modules/_threadmodule.c:1921:14 (python3.14t+0x57ff51) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#6 cfunction_call /project/cpython/Objects/methodobject.c:551:18 (python3.14t+0x28dee7) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#7 _PyObject_MakeTpCall /project/cpython/Objects/call.c:242:18 (python3.14t+0x1e87cc) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#8 _PyObject_VectorcallTstate /project/cpython/./Include/internal/pycore_call.h:165:16 (python3.14t+0x1e9428) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#9 PyObject_Vectorcall /project/cpython/Objects/call.c:327:12 (python3.14t+0x1e9428)
	#10 _PyEval_EvalFrameDefault /project/cpython/Python/generated_cases.c.h:2215:35 (python3.14t+0x3f6bc8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#11 _PyEval_EvalFrame /project/cpython/./Include/internal/pycore_ceval.h:116:16 (python3.14t+0x3eb6ca) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#12 _PyEval_Vector /project/cpython/Python/ceval.c:1911:12 (python3.14t+0x3eb6ca)
	#13 PyEval_EvalCode /project/cpython/Python/ceval.c:658:21 (python3.14t+0x3eb1d6) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#14 run_eval_code_obj /project/cpython/Python/pythonrun.c:1338:9 (python3.14t+0x4bcaee) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#15 run_mod /project/cpython/Python/pythonrun.c:1423:19 (python3.14t+0x4bc215) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#16 pyrun_file /project/cpython/Python/pythonrun.c:1256:15 (python3.14t+0x4b8310) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#17 _PyRun_SimpleFileObject /project/cpython/Python/pythonrun.c:491:13 (python3.14t+0x4b8310)
	#18 _PyRun_AnyFileObject /project/cpython/Python/pythonrun.c:78:15 (python3.14t+0x4b79d8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#19 pymain_run_file_obj /project/cpython/Modules/main.c:410:15 (python3.14t+0x4f771f) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#20 pymain_run_file /project/cpython/Modules/main.c:429:15 (python3.14t+0x4f771f)
	#21 pymain_run_python /project/cpython/Modules/main.c:697:21 (python3.14t+0x4f696c) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#22 Py_RunMain /project/cpython/Modules/main.c:776:5 (python3.14t+0x4f696c)
	#23 pymain_main /project/cpython/Modules/main.c:806:12 (python3.14t+0x4f6da8) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#24 Py_BytesMain /project/cpython/Modules/main.c:830:12 (python3.14t+0x4f6e2b) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)
	#25 main /project/cpython/./Programs/python.c:15:12 (python3.14t+0x15e7eb) (BuildId: edd9563c30a0b5c4e5f9718de6c6baf395511e5b)

	SUMMARY: ThreadSanitizer: data race /tmp/output-python/include/python3.14t/cpython/pyatomic_gcc.h:375:10 in _Py_atomic_load_uintptr_relaxed(unsigned long const*)
	==================
	ThreadSanitizer: reported 1 warnings
	import math
	import numpy as np

	import concurrent.futures
	import threading

	from fake_pmap_ext import pmap_lib


	if __name__ == "__main__":
	num_workers = 40

	barrier = threading.Barrier(num_workers)

	def _pmap(fun):
	cpp_mapped_f = pmap_lib.pmap(fun)

	return cpp_mapped_f

	def closure():
	barrier.wait()

	in_shape = (1, 2, 4)
	x = np.arange(math.prod(in_shape)).reshape(in_shape)
	fn = _pmap(lambda x: x)

	with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
	futures = []
	for i in range(num_workers):
	futures.append(executor.submit(closure))
	assert len(list(f.result() for f in futures)) == num_workers
	cmake_minimum_required(VERSION 3.16)

	project(fake_pmap_ext_project)

	# https://nanobind.readthedocs.io/en/latest/building.html


	if (CMAKE_VERSION VERSION_LESS 3.18)
	set(DEV_MODULE Development)
	else()
	set(DEV_MODULE Development.Module)
	endif()

	find_package(Python 3.8 COMPONENTS Interpreter ${DEV_MODULE} REQUIRED)

	if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
	set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
	set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
	endif()

	option(USE_TSAN "Compile with TSAN" ON)

	function(append value)
	foreach(variable ${ARGN})
	set(${variable} "${${variable}} ${value}" PARENT_SCOPE)
	endforeach(variable)
	endfunction()

	if (USE_TSAN)
	append("-fsanitize=thread -g -O2" CMAKE_C_FLAGS CMAKE_CXX_FLAGS)
	append("-fsanitize=thread" CMAKE_EXE_LINKER_FLAGS CMAKE_MODULE_LINKER_FLAGS CMAKE_SHARED_LINKER_FLAGS)
	endif()

	# Detect the installed nanobind package and import it into CMake
	execute_process(
	COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
	OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)

	find_package(nanobind CONFIG REQUIRED)

	# Abseil requires C++14
	set(CMAKE_CXX_STANDARD 14)

	nanobind_add_module(
	fake_pmap_ext
	FREE_THREADED
	binding.cpp
	)

	set_target_properties(fake_pmap_ext PROPERTIES POSITION_INDEPENDENT_CODE ON)
	# target_link_libraries(fake_pmap_ext PRIVATE absl::strings absl::node_hash_map absl::base absl::synchronization absl::cleanup)