Skip to content

Commit ce7b919

Browse files
xw285cornellfacebook-github-bot
authored andcommitted
Use gpu_library_selector for permute_pooled_embedding_ops_gpu (#1799)
Summary: X-link: pytorch/FBGEMM#2435 Pull Request resolved: #1799 Using gpu_library_selector can avoid divergence in the usercode A few notable changes: * the original permute_pooled_embedding_ops_gpu is missing link_whole * Make permute_pooled_embedding_function.cpp as its own library. This is because permute_pooled_embedding_function.cpp is included in both permute_pooled_embedding_ops_gpu and permute_pooled_embedding_ops_cpu. So when we link the _cpu and _gpu together, it'll complain about duplicate symbol. Reviewed By: houseroad Differential Revision: D54986292 fbshipit-source-id: b029c689eaf5f143b6d0b710aa30a06d1dd141f5
1 parent 7fb1d62 commit ce7b919

File tree

1 file changed

+3
-14
lines changed

1 file changed

+3
-14
lines changed

torchrec/distributed/utils.py

+3-14
Original file line numberDiff line numberDiff line change
@@ -35,20 +35,9 @@
3535
torch.ops.load_library(
3636
"//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_cpu"
3737
)
38-
try:
39-
if torch.version.hip:
40-
torch.ops.load_library(
41-
"//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_gpu_hip"
42-
)
43-
else:
44-
torch.ops.load_library(
45-
"//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_gpu_cuda"
46-
)
47-
except OSError:
48-
# For backward compatibility
49-
torch.ops.load_library(
50-
"//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_gpu"
51-
)
38+
torch.ops.load_library(
39+
"//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_gpu"
40+
)
5241
except OSError:
5342
pass
5443

0 commit comments

Comments
 (0)