From 2c10b09833c42798f10897ad828b2b9b8fec2820 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Wed, 13 Jan 2021 07:18:45 +0530 Subject: [PATCH 1/4] [Frontend][Tensorflow] Sparse dense matmul adjoint option added --- python/tvm/relay/frontend/tensorflow.py | 20 ++++--------------- .../frontend/tensorflow/test_forward.py | 12 ++++++----- 2 files changed, 11 insertions(+), 21 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index d5746a38582c..06b74f935b4a 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -942,7 +942,10 @@ def _impl(inputs, attr, params, mod): ) if sparse_lhs: - data = _op.transpose(data) + if attr.get("adjoint_a"): + weight_sp = csr_matrix(weight_sp.transpose()) + if not attr.get("adjoint_b"): + data = _op.transpose(data) else: weight_sp = csr_matrix(weight_sp.transpose()) @@ -955,21 +958,6 @@ def _impl(inputs, attr, params, mod): if not sparse_lhs: ret = _op.transpose(ret) - # Case 1. If both are true means first input was dense and second was sparse - # Case 2. If both are false means first input was sparse and second was dense - # TODO(ANSHUMAN87): Support other adjoint option too - if not ( - (attr.get("adjoint_a") and attr.get("adjoint_b")) - or ((not attr.get("adjoint_a")) and (not attr.get("adjoint_b"))) - ): - raise tvm.error.OpAttributeUnImplemented( - "Only tf.sparse.sparse_dense_matmul() with adjoint_a=True and adjoint_b=True" - "or with adjoint_a=False and adjoint_b=False" - " is supported, but adjoint_a={} and adjoint_b={} was supplied.".format( - attr.get("adjoint_a"), attr.get("adjoint_b") - ) - ) - return ret return _impl diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index d71405796ede..ccd17a4fab85 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1758,19 +1758,21 @@ def test_forward_batch_matmul(): # ---------------------------------- -def _test_sparse_dense_matmul(indices, values, A_shape, B_shape, dtype, flip=False): +def _test_sparse_dense_matmul(indices, values, A_inp_shape, B_inp_shape, dtype, flip=False): """ One iteration of sparse_dense_matmul """ - # TODO(ANSHUMAN87): Support adjoint options too - for adjoint_a in [False]: - for adjoint_b in [False]: + for adjoint_a in [False, True]: + for adjoint_b in [False, True]: + A_shape = A_inp_shape[::-1] if adjoint_a else A_inp_shape + B_shape = B_inp_shape[::-1] if adjoint_b else B_inp_shape + with tf.Graph().as_default(): A_sp = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=A_shape) B = tf.placeholder(shape=B_shape, dtype=dtype, name="B") if flip: result = tf.sparse.sparse_dense_matmul( - B, A_sp, adjoint_a=adjoint_a, adjoint_b=adjoint_b + B, A_sp, adjoint_a=adjoint_b, adjoint_b=adjoint_a ) else: result = tf.sparse.sparse_dense_matmul( From 8c3c5dc0b02ef34a519ca906709bf3d59b374945 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Fri, 22 Jan 2021 23:01:32 +0530 Subject: [PATCH 2/4] [1] Review comments handled --- python/tvm/relay/frontend/tensorflow.py | 44 +++++++++++++++++++++---- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 06b74f935b4a..b6d1fbc44c21 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -926,13 +926,6 @@ def _impl(inputs, attr, params, mod): data = inputs[3] - # By default, in tensorflow the first input ,i.e., data is sparse - sparse_lhs = True - - # If both are true means First input was dense and second was sparse - if attr.get("adjoint_a") and attr.get("adjoint_b"): - sparse_lhs = False - rows = [x[0] for x in indices_tensor] cols = [x[1] for x in indices_tensor] @@ -941,12 +934,48 @@ def _impl(inputs, attr, params, mod): (values_tensor, (rows, cols)), shape=tuple(dense_shape_tensor.tolist()) ) + # As per tensorflow implementation, we have 4 possible input combination + # and the first input(A) is always sparse and second input(B) is always dense. + # Case 1: A , B , adjoint_a=False, adjoint_b=False --> A * B + # Case 2: A , B , adjoint_a=True, adjoint_b=False --> A.T * B + # Case 3: A , B , adjoint_a=False, adjoint_b=True --> A * B.T + # Case 4: A , B , adjoint_a=True, adjoint_b=True --> (A.T * B.T).T + # + # Topi implementation for sparse_dense(matmul) has 2 possible input + # combination where first input(A) is always dense + # and second input(B) is always sparse. + # Case 1: A , B, sparse_lhs = False --> A * B.T + # Case 2: A , B, sparse_lhs = True --> B * A.T + # + # The mapping would be as below: + # TF Case 1: A , B , adjoint_a=False, adjoint_b=False + # --> sparse_dense(transpose(B), A, sparse_lhs=True) + # + # TF Case 2: A , B , adjoint_a=True, adjoint_b=False + # --> sparse_dense(transpose(B), transpose(A), sparse_lhs=True) + # + # TF Case 3: A , B , adjoint_a=False, adjoint_b=True + # --> sparse_dense(B, A, sparse_lhs=True) + # + # TF Case 4: A , B , adjoint_a=True, adjoint_b=True + # --> transpose(sparse_dense(B, transpose(A), sparse_lhs=False)) + + # By default, in tensorflow the first input ,i.e., data is sparse + sparse_lhs = True + + # If both are true means First input was dense and second was sparse + if attr.get("adjoint_a") and attr.get("adjoint_b"): + sparse_lhs = False + if sparse_lhs: + # TF Case 2: if attr.get("adjoint_a"): weight_sp = csr_matrix(weight_sp.transpose()) + # TF Case 1 & TF Case 2: if not attr.get("adjoint_b"): data = _op.transpose(data) else: + # TF Case 4 weight_sp = csr_matrix(weight_sp.transpose()) weight_data = _expr.const(weight_sp.data, weight_sp.data.dtype) @@ -956,6 +985,7 @@ def _impl(inputs, attr, params, mod): ret = _op.nn.sparse_dense(data, [weight_data, weight_indices, weight_indptrs], sparse_lhs) if not sparse_lhs: + # TF Case 4 ret = _op.transpose(ret) return ret From 47d211ec2ae1fc77e83504be39bbdcf3ad210ac0 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Sun, 24 Jan 2021 19:01:57 +0530 Subject: [PATCH 3/4] [2] Review comments handled --- python/tvm/relay/frontend/tensorflow.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index b6d1fbc44c21..0e740ca2bf30 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -963,19 +963,19 @@ def _impl(inputs, attr, params, mod): # By default, in tensorflow the first input ,i.e., data is sparse sparse_lhs = True - # If both are true means First input was dense and second was sparse - if attr.get("adjoint_a") and attr.get("adjoint_b"): - sparse_lhs = False - - if sparse_lhs: - # TF Case 2: - if attr.get("adjoint_a"): - weight_sp = csr_matrix(weight_sp.transpose()) - # TF Case 1 & TF Case 2: - if not attr.get("adjoint_b"): - data = _op.transpose(data) + # TF Case 1: + if not attr.get("adjoint_a") and not attr.get("adjoint_b"): + data = _op.transpose(data) + # TF Case 2: + elif attr.get("adjoint_a") and not attr.get("adjoint_b"): + data = _op.transpose(data) + weight_sp = csr_matrix(weight_sp.transpose()) + # TF Case 3: + elif not attr.get("adjoint_a") and attr.get("adjoint_b"): + pass + # TF Case 4: + # attr.get("adjoint_a") and attr.get("adjoint_b"): else: - # TF Case 4 weight_sp = csr_matrix(weight_sp.transpose()) weight_data = _expr.const(weight_sp.data, weight_sp.data.dtype) From 6d26ce5a7c8b69c0f6eb663f2734db37ff0a48e4 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Wed, 27 Jan 2021 09:13:26 +0530 Subject: [PATCH 4/4] [3] Review comments handled --- python/tvm/relay/frontend/tensorflow.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 0e740ca2bf30..c6f3b47419cd 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -939,7 +939,7 @@ def _impl(inputs, attr, params, mod): # Case 1: A , B , adjoint_a=False, adjoint_b=False --> A * B # Case 2: A , B , adjoint_a=True, adjoint_b=False --> A.T * B # Case 3: A , B , adjoint_a=False, adjoint_b=True --> A * B.T - # Case 4: A , B , adjoint_a=True, adjoint_b=True --> (A.T * B.T).T + # Case 4: A , B , adjoint_a=True, adjoint_b=True --> A.T * B.T # # Topi implementation for sparse_dense(matmul) has 2 possible input # combination where first input(A) is always dense @@ -949,15 +949,19 @@ def _impl(inputs, attr, params, mod): # # The mapping would be as below: # TF Case 1: A , B , adjoint_a=False, adjoint_b=False + # --> In TF: A * B --> In Topi: A * B.T.T # --> sparse_dense(transpose(B), A, sparse_lhs=True) # # TF Case 2: A , B , adjoint_a=True, adjoint_b=False + # --> In TF: A.T * B --> In Topi: A.T * B.T.T # --> sparse_dense(transpose(B), transpose(A), sparse_lhs=True) # # TF Case 3: A , B , adjoint_a=False, adjoint_b=True + # --> In TF: A * B.T --> In Topi: A * B # --> sparse_dense(B, A, sparse_lhs=True) # # TF Case 4: A , B , adjoint_a=True, adjoint_b=True + # --> In TF: A.T * B.T --> In Topi: (B * A.T).T # --> transpose(sparse_dense(B, transpose(A), sparse_lhs=False)) # By default, in tensorflow the first input ,i.e., data is sparse @@ -976,6 +980,7 @@ def _impl(inputs, attr, params, mod): # TF Case 4: # attr.get("adjoint_a") and attr.get("adjoint_b"): else: + sparse_lhs = False weight_sp = csr_matrix(weight_sp.transpose()) weight_data = _expr.const(weight_sp.data, weight_sp.data.dtype)