-
Notifications
You must be signed in to change notification settings - Fork 445
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add helper ops to support cache conflict misses (#2571)
Summary: This diff adds helper operators for the cache conflict miss support enablement in SSD TBE. Changes include: - Extend `get_unique_indices_cuda` to compute and return inverse linear indices (the tensor that contains the original positions of lienar indices before sorting) - Extend `lru_cache_find_uncached_cuda` to compute and return the inverse cache sets (the tensor that contains the original positions of cache sets of unique indices before sorting) - Update SSD backend to support cache conflict misses instead of failing. The rows that experience conflict misses will be stored in a scratch pad for TBE kernels to consume. They will be evicted to SSD once the backward+optimizer step of TBE is completed. - Add `ssd_generate_row_addrs` for generating row addresses of data that is fetched from SSD (data can be in either a scratch pad or LXU cache). Reviewed By: q10 Differential Revision: D55926421
- Loading branch information
1 parent
309e9e2
commit fb11951
Showing
16 changed files
with
618 additions
and
205 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,48 @@ | ||
# Copyright (c) Meta Platforms, Inc. and affiliates. | ||
# All rights reserved. | ||
# | ||
# This source code is licensed under the BSD-style license found in the | ||
# LICENSE file in the root directory of this source tree. | ||
|
||
from typing import Optional, Tuple, Union | ||
|
||
import torch | ||
|
||
lib = torch.library.Library("fbgemm", "FRAGMENT") | ||
lib.define( | ||
""" | ||
get_unique_indices( | ||
Tensor linear_indices, | ||
int max_indices, | ||
bool compute_count=False, | ||
bool compute_inverse_indices=False | ||
) -> (Tensor, Tensor, Tensor?, Tensor?) | ||
""" | ||
) | ||
|
||
|
||
@torch.library.impl(lib, "get_unique_indices", "CUDA") | ||
def get_unique_indices( | ||
linear_indices: torch.Tensor, | ||
max_indices: int, | ||
compute_count: bool = False, | ||
compute_inverse_indices: bool = False, | ||
) -> Union[ | ||
Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]], | ||
Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]], | ||
]: | ||
""" | ||
A wrapper for get_unique_indices for overloading the return type | ||
based on inputs | ||
""" | ||
ret = torch.ops.fbgemm.get_unique_indices_internal( | ||
linear_indices, | ||
max_indices, | ||
compute_count, | ||
compute_inverse_indices, | ||
) | ||
if not compute_inverse_indices: | ||
# Return only 3 tensors | ||
return ret[:-1] | ||
# Return all tensors | ||
return ret |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.