frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Operator Versioning][Edge] Initialize upgrader thread safe (#70161)

Browse Source 
Summary:
Upgrader should only be initialized once when runtime loads the first module. It no longer needs to initialized afterwards.

Previously, instead of using an atomic variable, the upgrader will be initialized depends on whether byteCodeFunctionWithOperator.function.get_code().operators_ is empty. If it's empty, it means the operator from the upgrader is not initialized yet. However, it's not thread safe. When multiple thread loads module together, it's possible that they all consider it's the first module. Use an atomic variable here to make sure it's thread safe.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/70161

ghstack-source-id: 146012642

Test Plan:
```
buck test mode/opt //papaya/integration/service/test/analytics/histogram:generic_histogram_system_test -- --exact 'papaya/integration/service/test/analytics/histogram:generic_histogram_system_test - SumHistogramSystemTest.test' --run-disabled
buck test mode/opt //caffe2/test/cpp/jit:jit
```

Reviewed By: iseeyuan

Differential Revision: D33220320

fbshipit-source-id: 10f2397c3b358d5a1d39a2ce25457e3fdb640d2c
This commit is contained in:
Chen Lai
2021-12-19 20:14:40 -08:00
committed by Facebook GitHub Bot
parent 7ea86dfdb1
commit 8a912014b1
3 changed files with 23 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
torch/csrc/jit/mobile/import.cpp
View File

@ -226,6 +226,7 @@ class BytecodeDeserializer final {
// operators from the given model. If it's less than the current runtime,
// upgrader will be applied at loading stage.
uint64_t operator_version_;
std::atomic<bool> _upgrader_initialized{false};
};
BytecodeDeserializer::BytecodeDeserializer(
@ -301,13 +302,21 @@ void BytecodeDeserializer::parseFunctionSchema(
}
void BytecodeDeserializer::init_upgrader(mobile::Function* function) {
for (auto& byteCodeFunctionWithOperator : getUpgraderBytecodeList()) {
// When kUpgraderByteCode is initialized in upgrader_mobile.h, the mobile
// function is initialized with everything (instruction, constants, types,
// registerer size and etc), except operator. The operator function is also
// static initialized and is available later. The oprator for the upgrader
// function will be initialized when the first module is loaded.
if (byteCodeFunctionWithOperator.function.get_code().operators_.empty()) {
// Upgrader should only be initialized once when runtime loads the first
// module. It no longer needs to initialized afterwards. Previously, instead
// of using an atomic variable, the upgrader will be initailized depends on
// whether byteCodeFunctionWithOperator.function.get_code().operators_ is
// empty. If it's empty, it means the operator from the upgrader is not
// initialized yet. However, it's not thread safe. When multiple thread loads
// module together, it's possible that they all consider it's the first
// module. Use an atomic variable here to make sure it's thread safe.
if (!_upgrader_initialized.load(std::memory_order_seq_cst)) {
for (auto& byteCodeFunctionWithOperator : getUpgraderBytecodeList()) {
// When kUpgraderByteCode is initialized in upgrader_mobile.h, the mobile
// function is initialized with everything (instruction, constants, types,
// registerer size and etc), except operator. The operator function is
// also static initialized and is available later. The oprator for the
// upgrader function will be initialized when the first module is loaded.
for (const auto& op : byteCodeFunctionWithOperator.operators) {
byteCodeFunctionWithOperator.function.append_operator(
op.name,
@ -315,8 +324,12 @@ void BytecodeDeserializer::init_upgrader(mobile::Function* function) {
op.num_specified_args,
caffe2::serialize::kMaxSupportedFileFormatVersion);
}
// Add the upgrader function in code.
function->append_function(byteCodeFunctionWithOperator.function);
}
function->append_function(byteCodeFunctionWithOperator.function);
// Set the flag _upgrader_initialized to ture, and for the 2, 3, ...
// modules, no need to initalized again
_upgrader_initialized.store(true, std::memory_order_seq_cst);
}
}