gemini-code-assist[bot] commented on code in PR #655:
URL: https://github.com/apache/tvm-ffi/pull/655#discussion_r3524167868
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -113,6 +105,7 @@ void ORCJITDynamicLibraryObj::AddObjectFile(const String&
path) {
// Add object file to this JITDylib
TVM_FFI_ORCJIT_LLVM_CALL(jit_->addObjectFile(*dylib_,
std::move(*buffer_or_err)));
+ context_symbol_refreshed_ = false;
Review Comment:

Protect the modification of `context_symbol_refreshed_` with the mutex to
prevent data races with concurrent calls to `GetFunction`.
```c
{
std::lock_guard<std::mutex> lock(mutex_);
context_symbol_refreshed_ = false;
}
```
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h:
##########
@@ -107,6 +107,9 @@ class ORCJITDynamicLibraryObj : public ModuleObj {
/*! \brief Link order tracking (to support incremental linking) */
llvm::orc::JITDylibSearchOrder link_order_;
+ /*! \brief Whether context slots have been refreshed since the last object
add. */
+ bool context_symbol_refreshed_{false};
+
Review Comment:

Since `GetFunction` can be called concurrently from multiple threads on a
shared `Module` instance, accessing and modifying `context_symbol_refreshed_`
without synchronization introduces a data race (undefined behavior). We should
introduce a `std::mutex` to protect the lazy initialization and refresh of
context symbols. Note that you will also need to include `<mutex>` at the top
of this header file.
```c
/*! \brief Whether context slots have been refreshed since the last object
add. */
bool context_symbol_refreshed_{false};
/*! \brief Mutex to protect context symbol refresh and avoid data races. */
std::mutex mutex_;
```
##########
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc:
##########
@@ -181,6 +180,18 @@ Optional<Function>
ORCJITDynamicLibraryObj::GetFunction(const String& name) {
});
}
+ if (!context_symbol_refreshed_) {
+ if (void** ctx_addr =
reinterpret_cast<void**>(GetSymbol(symbol::tvm_ffi_library_ctx))) {
+ *ctx_addr = this;
+ }
+ Module::VisitContextSymbols([this](const String& name, void* symbol) {
+ if (void** ctx_addr = reinterpret_cast<void**>(GetSymbol(name))) {
+ *ctx_addr = symbol;
+ }
+ });
+ context_symbol_refreshed_ = true;
+ }
Review Comment:

Acquire the mutex lock before checking and updating
`context_symbol_refreshed_` to ensure thread-safe lazy initialization of the
context symbols.
```c
{
std::lock_guard<std::mutex> lock(mutex_);
if (!context_symbol_refreshed_) {
if (void** ctx_addr =
reinterpret_cast<void**>(GetSymbol(symbol::tvm_ffi_library_ctx))) {
*ctx_addr = this;
}
Module::VisitContextSymbols([this](const String& name, void* symbol) {
if (void** ctx_addr = reinterpret_cast<void**>(GetSymbol(name))) {
*ctx_addr = symbol;
}
});
context_symbol_refreshed_ = true;
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]