Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
dgl
Commits
cf8fea73
Unverified
Commit
cf8fea73
authored
Jan 18, 2024
by
Muhammed Fatih BALIN
Committed by
GitHub
Jan 18, 2024
Browse files
[GraphBolt][CUDA] Hot fix after inplace pinning PR (#6974)
parent
aad12df6
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
8 additions
and
5 deletions
+8
-5
python/dgl/graphbolt/impl/fused_csc_sampling_graph.py
python/dgl/graphbolt/impl/fused_csc_sampling_graph.py
+8
-5
No files found.
python/dgl/graphbolt/impl/fused_csc_sampling_graph.py
View file @
cf8fea73
...
@@ -34,16 +34,17 @@ class FusedCSCSamplingGraph(SamplingGraph):
...
@@ -34,16 +34,17 @@ class FusedCSCSamplingGraph(SamplingGraph):
):
):
super
().
__init__
()
super
().
__init__
()
self
.
_c_csc_graph
=
c_csc_graph
self
.
_c_csc_graph
=
c_csc_graph
self
.
_is_inplace_pinned
=
set
()
def
__del__
(
self
):
def
__del__
(
self
):
# torch.Tensor.pin_memory() is not an inplace operation. To make it
# torch.Tensor.pin_memory() is not an inplace operation. To make it
# truly in-place, we need to use cudaHostRegister. Then, we need to use
# truly in-place, we need to use cudaHostRegister. Then, we need to use
# cudaHostUnregister to unpin the tensor in the destructor.
# cudaHostUnregister to unpin the tensor in the destructor.
# https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
# https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
if
hasattr
(
self
,
"_is_inplace_pinned"
):
for
tensor
in
self
.
_is_inplace_pinned
:
for
tensor
in
self
.
_is_inplace_pinned
:
assert
(
assert
(
torch
.
cuda
.
cudart
().
cudaHostUnregister
(
tensor
.
data_ptr
())
==
0
torch
.
cuda
.
cudart
().
cudaHostUnregister
(
tensor
.
data_ptr
())
==
0
)
)
@
property
@
property
...
@@ -990,6 +991,8 @@ class FusedCSCSamplingGraph(SamplingGraph):
...
@@ -990,6 +991,8 @@ class FusedCSCSamplingGraph(SamplingGraph):
# cudaHostUnregister to unpin the tensor in the destructor.
# cudaHostUnregister to unpin the tensor in the destructor.
# https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
# https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
cudart
=
torch
.
cuda
.
cudart
()
cudart
=
torch
.
cuda
.
cudart
()
if
not
hasattr
(
self
,
"_is_inplace_pinned"
):
self
.
_is_inplace_pinned
=
set
()
def
_pin
(
x
):
def
_pin
(
x
):
if
hasattr
(
x
,
"pin_memory_"
):
if
hasattr
(
x
,
"pin_memory_"
):
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment