Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make check_feature_gate_key PT2 compatible #3426

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Make check_feature_gate_key PT2 compatible #3426

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b80208b
Make check_feature_gate_key PT2 compatible (#3425)
sryap Nov 30, 2024
8f561dd
Make check_feature_gate_key PT2 compatible
sryap Nov 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion fbgemm_gpu/fbgemm_gpu/config/feature_list.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,10 @@ class FeatureGate:
FeatureGate.is_enabled(FeatureGateName.TBE_V2)
"""

dummy_tensor = torch.empty(1)

@classmethod
def is_enabled(cls, feature: FeatureGateName) -> bool:
return torch.ops.fbgemm.check_feature_gate_key(feature.name)
return torch.ops.fbgemm.check_feature_gate_key_pt2(
cls.dummy_tensor, feature.name
)[0].item()
25 changes: 25 additions & 0 deletions fbgemm_gpu/src/config/feature_gates.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,21 @@ DLL_PUBLIC bool check_feature_gate_key(const std::string& key) {
}
}

DLL_PUBLIC at::Tensor check_feature_gate_key_pt2(
at::Tensor& tensor,
const std::string& key) {
auto output = at::empty({1}, tensor.options().dtype(at::kBool));
output.data_ptr<bool>()[0] = check_feature_gate_key(key);
return output;
}

DLL_PUBLIC at::Tensor check_feature_gate_key_pt2_meta(
at::Tensor& tensor,
const std::string& key) {
auto output = at::empty({1}, tensor.options().dtype(at::kBool));
return output;
}

DLL_PUBLIC bool is_feature_enabled(const FeatureGateName& feature) {
return check_feature_gate_key(to_string(feature));
}
Expand All @@ -81,4 +96,14 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
m.def(
"check_feature_gate_key(str key) -> bool",
fbgemm_gpu::config::check_feature_gate_key);
m.def("check_feature_gate_key_pt2(Tensor tensor, str key) -> Tensor");
DISPATCH_TO_CPU(
"check_feature_gate_key_pt2",
fbgemm_gpu::config::check_feature_gate_key_pt2);
DISPATCH_TO_CUDA(
"check_feature_gate_key_pt2",
fbgemm_gpu::config::check_feature_gate_key_pt2);
DISPATCH_TO_META(
"check_feature_gate_key_pt2",
fbgemm_gpu::config::check_feature_gate_key_pt2_meta);
}
Loading