Skip to content

Commit

Permalink
[caffe2][be] migrate gloabl static initializer (pytorch#128784)
Browse files Browse the repository at this point in the history
Summary:
Caffe2 lib has 200+ global static initializer usage, which are papar-cut reference to startup perf. Detail in this post https://fb.workplace.com/groups/arglassesperf/permalink/623909116287154.

This Diff migrate StorageImpl.cpp

Addtional Context: https://fb.workplace.com/groups/arglassesperf/permalink/623909116287154

Test Plan: CI

Differential Revision: D58639283

Pull Request resolved: pytorch#128784
Approved by: https://github.com/aaronenyeshi
  • Loading branch information
klein-meta authored and pytorchmergebot committed Jun 25, 2024
1 parent fd4af87 commit 0e6bb7f
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions c10/core/StorageImpl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,12 @@ namespace c10 {
C10_API std::array<StorageImplCreateHelper, at::COMPILE_TIME_MAX_DEVICE_TYPES>
StorageImplCreate;

// A allowlist of device type, currently available is PrivateUse1.
static ska::flat_hash_set<c10::DeviceType> DeviceTypeAllowList{
DeviceType::PrivateUse1};
// A allowlist of device type, currently available is PrivateUse1
inline ska::flat_hash_set<c10::DeviceType>& GetBackendMetaAllowlist() {
static ska::flat_hash_set<c10::DeviceType> DeviceTypeAllowList{
DeviceType::PrivateUse1};
return DeviceTypeAllowList;
}

void throwNullDataPtrError() {
TORCH_CHECK(
Expand Down Expand Up @@ -41,8 +44,9 @@ void SetStorageImplCreate(DeviceType t, StorageImplCreateHelper fptr) {
// Allowlist verification.
// Only if the devicetype is in the allowlist,
// we allow the extension to be registered for storageImpl create.
const auto& DeviceTypeAllowlist = GetBackendMetaAllowlist();
TORCH_CHECK(
DeviceTypeAllowList.find(t) != DeviceTypeAllowList.end(),
DeviceTypeAllowlist.find(t) != DeviceTypeAllowlist.end(),
"It is only allowed to register the storageImpl create method ",
"for PrivateUse1. ",
"If you have related storageImpl requirements, ",
Expand Down

0 comments on commit 0e6bb7f

Please sign in to comment.