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
21c4c29a
Unverified
Commit
21c4c29a
authored
Jan 12, 2023
by
czkkkkkk
Committed by
GitHub
Jan 12, 2023
Browse files
[Sparse] Reduce peak memory of SpSpAdd. (#5143)
* [Sparse] Reduce peak memory usage of SpSpAdd * Update * Update
parent
b743f767
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
9 additions
and
3 deletions
+9
-3
dgl_sparse/src/elemenwise_op.cc
dgl_sparse/src/elemenwise_op.cc
+9
-3
No files found.
dgl_sparse/src/elemenwise_op.cc
View file @
21c4c29a
...
...
@@ -22,9 +22,15 @@ c10::intrusive_ptr<SparseMatrix> SpSpAdd(
const
c10
::
intrusive_ptr
<
SparseMatrix
>&
A
,
const
c10
::
intrusive_ptr
<
SparseMatrix
>&
B
)
{
ElementwiseOpSanityCheck
(
A
,
B
);
torch
::
Tensor
sum
;
{
// TODO(#5145) This is a workaround to reduce peak memory usage. It is no
// longer needed after we address #5145.
auto
torch_A
=
COOToTorchCOO
(
A
->
COOPtr
(),
A
->
value
());
auto
torch_B
=
COOToTorchCOO
(
B
->
COOPtr
(),
B
->
value
());
auto
sum
=
(
torch_A
+
torch_B
).
coalesce
();
sum
=
torch_A
+
torch_B
;
}
sum
=
sum
.
coalesce
();
auto
indices
=
sum
.
indices
();
auto
row
=
indices
[
0
];
auto
col
=
indices
[
1
];
...
...
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