mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
375687839e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56617 This migrates the sparsity to the open source Test Plan: `buck test mode/opt //caffe2/test:ao` Reviewed By: raghuramank100 Differential Revision: D27812207 fbshipit-source-id: cc87d9d2b486269901a4ad9b483615741a1cd712
39 lines
1.7 KiB
Python
39 lines
1.7 KiB
Python
import threading
|
|
|
|
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 QNNPACKLinearBlockSparsePattern:
|
|
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))
|
|
QNNPACKLinearBlockSparsePattern.rlock.acquire()
|
|
QNNPACKLinearBlockSparsePattern.prev_row_block_size = QNNPACKLinearBlockSparsePattern.row_block_size
|
|
QNNPACKLinearBlockSparsePattern.prev_col_block_size = QNNPACKLinearBlockSparsePattern.col_block_size
|
|
QNNPACKLinearBlockSparsePattern.row_block_size = row_block_size
|
|
QNNPACKLinearBlockSparsePattern.col_block_size = col_block_size
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, backtrace):
|
|
QNNPACKLinearBlockSparsePattern.row_block_size = QNNPACKLinearBlockSparsePattern.prev_row_block_size
|
|
QNNPACKLinearBlockSparsePattern.col_block_size = QNNPACKLinearBlockSparsePattern.prev_col_block_size
|
|
QNNPACKLinearBlockSparsePattern.rlock.release()
|
|
|
|
@staticmethod
|
|
def block_size():
|
|
return QNNPACKLinearBlockSparsePattern.row_block_size, QNNPACKLinearBlockSparsePattern.col_block_size
|