mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
6a866c3ed1
Summary: this mostly consisted of adding __all__ to files without them. A few functions in X.utils were made private too Test Plan: python test/test_public_bindings.py Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D40814548](https://our.internmc.facebook.com/intern/diff/D40814548) Pull Request resolved: https://github.com/pytorch/pytorch/pull/87883 Approved by: https://github.com/jcaip, https://github.com/anjali411
43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
import threading
|
|
|
|
__all__ = [
|
|
"LinearBlockSparsePattern"
|
|
]
|
|
|
|
def _is_valid_linear_block_sparse_pattern(row_block_size, col_block_size):
|
|
return (row_block_size == 1 and col_block_size == 4) or \
|
|
(row_block_size == 8 and col_block_size == 1)
|
|
|
|
# This is a stop-gap measure as current flow does not allow module
|
|
# specific block sparse pattern.
|
|
# Infact there is no way to convey sparse pattern via module config
|
|
# of quantization flow. Thus using the global context to convey
|
|
# sparsity pattern.
|
|
# Once the flow supports it, this should be removed.
|
|
class LinearBlockSparsePattern:
|
|
rlock = threading.RLock()
|
|
row_block_size = 1
|
|
col_block_size = 4
|
|
prev_row_block_size = 1
|
|
prev_col_block_size = 4
|
|
|
|
def __init__(self, row_block_size=1, col_block_size=4):
|
|
assert(_is_valid_linear_block_sparse_pattern(row_block_size, col_block_size))
|
|
LinearBlockSparsePattern.rlock.acquire()
|
|
LinearBlockSparsePattern.prev_row_block_size = LinearBlockSparsePattern.row_block_size
|
|
LinearBlockSparsePattern.prev_col_block_size = LinearBlockSparsePattern.col_block_size
|
|
LinearBlockSparsePattern.row_block_size = row_block_size
|
|
LinearBlockSparsePattern.col_block_size = col_block_size
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, backtrace):
|
|
LinearBlockSparsePattern.row_block_size = LinearBlockSparsePattern.prev_row_block_size
|
|
LinearBlockSparsePattern.col_block_size = LinearBlockSparsePattern.prev_col_block_size
|
|
LinearBlockSparsePattern.rlock.release()
|
|
|
|
@staticmethod
|
|
def block_size():
|
|
return LinearBlockSparsePattern.row_block_size, LinearBlockSparsePattern.col_block_size
|