You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/07/09 22:56:51 UTC
[GitHub] [tvm] comaniac commented on a change in pull request #8440: [Relay][ONNX] Batch_matmul to dense optimization
comaniac commented on a change in pull request #8440:
URL: https://github.com/apache/tvm/pull/8440#discussion_r667245328
##########
File path: python/tvm/relay/frontend/onnx.py
##########
@@ -670,68 +670,67 @@ class MatMul(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
assert len(inputs) == 2, "MatMul op take 2 inputs, {} given".format(len(inputs))
+ inputs_0 = inputs[0]
+ inputs_1 = inputs[1]
+
# Need to check input shape as batch matmul must be supported.
- a_shape = shape_of(inputs[0])
- a_rank = infer_shape(a_shape)[0]
- b_shape = shape_of(inputs[1])
- b_rank = infer_shape(b_shape)[0]
- # When performing a batch matmul, we need to properly handle N-dim shapes.
- if a_rank > 2 or b_rank > 2:
+ a_shape = infer_shape(inputs_0)
+ b_shape = infer_shape(inputs_1)
- def flatten_to_3d(x, x_shape):
- ndims = infer_shape(x_shape)[0]
- newshape = _op.concatenate(
- [
- _expr.const([-1], dtype=infer_type(x_shape).checked_type.dtype),
- _op.strided_slice(x_shape, [ndims - 2], [ndims]),
- ],
- 0,
- )
- out = _op.reshape(x, fold_constant(newshape))
- return out
+ # When performing a batch matmul, we need to properly handle N-dim shapes.
+ if len(a_shape) > 2 and len(b_shape) > 2:
+ # Convert a into a 3 dimensional tensors.
+ need_reshape_output = False
+ if len(a_shape) != 3:
+ a = _op.reshape(inputs_0, [-1, a_shape[-2], a_shape[-1]])
+ need_reshape_output = True
+ else:
+ a = inputs_0
- # Convert a and b into 3 dimensional tensors.
- a = flatten_to_3d(inputs[0], a_shape)
- b = flatten_to_3d(inputs[1], b_shape)
# Transpose matrix dimensions of b.
- b = _op.transpose(b, [0, 2, 1])
+ trans_axes = list(range(len(b_shape)))
+ trans_axes[-2], trans_axes[-1] = trans_axes[-1], trans_axes[-2]
+ b = _op.transpose(inputs_1, trans_axes)
+
+ # Convert b into a 3 dimensional tensor. Note that the last two dimensions
+ # are transposed.
+ if len(b_shape) != 3:
+ b = _op.reshape(b, [-1, b_shape[-1], b_shape[-2]])
+
# Perform a batch matmul.
output = _op.nn.batch_matmul(a, b)
- # Determine the output batch dimension.
- if a_rank > b_rank:
- out_batch = _op.strided_slice(a_shape, [0], [a_rank - 2])
- elif a_rank < b_rank:
- out_batch = _op.strided_slice(b_shape, [0], [b_rank - 2])
- # If its unclear how broadcasting should be applied, the output
- # shape is determined by choosing the maximum value from each input.
- else:
- out_batch = _op.concatenate(
- [
- _op.maximum(
- _op.strided_slice(a_shape, [i], [i + 1]),
- _op.strided_slice(b_shape, [i], [i + 1]),
- )
- for i in range(a_rank - 2)
- ],
- 0,
- )
+
# Reshape output to original dimensions.
- final_shape = _op.concatenate(
- [
- out_batch,
- _op.strided_slice(
- a_shape, [infer_shape(a_shape)[0] - 2], [infer_shape(a_shape)[0] - 1]
- ),
- _op.strided_slice(
- b_shape, [infer_shape(b_shape)[0] - 1], [infer_shape(b_shape)[0]]
- ),
- ],
- 0,
+ if need_reshape_output:
+ return _op.reshape(output, [*a_shape[:-2], a_shape[-2], b_shape[-1]])
+ return output
+ elif len(a_shape) > 2:
+ inputs_0 = _op.reshape(inputs_0, [-1, a_shape[-1]])
Review comment:
Then you can get rid of the no-else-return error.
```suggestion
if len(a_shape) > 2:
inputs_0 = _op.reshape(inputs_0, [-1, a_shape[-1]])
```
##########
File path: python/tvm/relay/frontend/onnx.py
##########
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, import-self, len-as-condition, unused-argument, too-many-lines
-# pylint: disable=import-outside-toplevel
+# pylint: disable=import-outside-toplevel, no-else-return
Review comment:
Do not disable lint rule like this one.
##########
File path: python/tvm/relay/frontend/onnx.py
##########
@@ -670,68 +670,67 @@ class MatMul(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
assert len(inputs) == 2, "MatMul op take 2 inputs, {} given".format(len(inputs))
+ inputs_0 = inputs[0]
+ inputs_1 = inputs[1]
+
# Need to check input shape as batch matmul must be supported.
- a_shape = shape_of(inputs[0])
- a_rank = infer_shape(a_shape)[0]
- b_shape = shape_of(inputs[1])
- b_rank = infer_shape(b_shape)[0]
- # When performing a batch matmul, we need to properly handle N-dim shapes.
- if a_rank > 2 or b_rank > 2:
+ a_shape = infer_shape(inputs_0)
+ b_shape = infer_shape(inputs_1)
- def flatten_to_3d(x, x_shape):
- ndims = infer_shape(x_shape)[0]
- newshape = _op.concatenate(
- [
- _expr.const([-1], dtype=infer_type(x_shape).checked_type.dtype),
- _op.strided_slice(x_shape, [ndims - 2], [ndims]),
- ],
- 0,
- )
- out = _op.reshape(x, fold_constant(newshape))
- return out
+ # When performing a batch matmul, we need to properly handle N-dim shapes.
+ if len(a_shape) > 2 and len(b_shape) > 2:
+ # Convert a into a 3 dimensional tensors.
+ need_reshape_output = False
+ if len(a_shape) != 3:
+ a = _op.reshape(inputs_0, [-1, a_shape[-2], a_shape[-1]])
+ need_reshape_output = True
+ else:
+ a = inputs_0
- # Convert a and b into 3 dimensional tensors.
- a = flatten_to_3d(inputs[0], a_shape)
- b = flatten_to_3d(inputs[1], b_shape)
# Transpose matrix dimensions of b.
- b = _op.transpose(b, [0, 2, 1])
+ trans_axes = list(range(len(b_shape)))
+ trans_axes[-2], trans_axes[-1] = trans_axes[-1], trans_axes[-2]
+ b = _op.transpose(inputs_1, trans_axes)
+
+ # Convert b into a 3 dimensional tensor. Note that the last two dimensions
+ # are transposed.
+ if len(b_shape) != 3:
+ b = _op.reshape(b, [-1, b_shape[-1], b_shape[-2]])
+
# Perform a batch matmul.
output = _op.nn.batch_matmul(a, b)
- # Determine the output batch dimension.
- if a_rank > b_rank:
- out_batch = _op.strided_slice(a_shape, [0], [a_rank - 2])
- elif a_rank < b_rank:
- out_batch = _op.strided_slice(b_shape, [0], [b_rank - 2])
- # If its unclear how broadcasting should be applied, the output
- # shape is determined by choosing the maximum value from each input.
- else:
- out_batch = _op.concatenate(
- [
- _op.maximum(
- _op.strided_slice(a_shape, [i], [i + 1]),
- _op.strided_slice(b_shape, [i], [i + 1]),
- )
- for i in range(a_rank - 2)
- ],
- 0,
- )
+
# Reshape output to original dimensions.
- final_shape = _op.concatenate(
- [
- out_batch,
- _op.strided_slice(
- a_shape, [infer_shape(a_shape)[0] - 2], [infer_shape(a_shape)[0] - 1]
- ),
- _op.strided_slice(
- b_shape, [infer_shape(b_shape)[0] - 1], [infer_shape(b_shape)[0]]
- ),
- ],
- 0,
+ if need_reshape_output:
+ return _op.reshape(output, [*a_shape[:-2], a_shape[-2], b_shape[-1]])
+ return output
+ elif len(a_shape) > 2:
+ inputs_0 = _op.reshape(inputs_0, [-1, a_shape[-1]])
+
+ if len(b_shape) > 2:
+ trans_axes = list(range(len(b_shape)))
+ trans_axes[-2], trans_axes[-1] = trans_axes[-1], trans_axes[-2]
+ input_1 = _op.reshape(_op.transpose(inputs_1, trans_axes), [-1, b_shape[-2]])
+ elif len(b_shape) == 2:
+ input_1 = _op.transpose(inputs_1, axes=(1, 0))
+ elif len(b_shape) == 1:
+ input_1 = _op.expand_dims(inputs_1, 0, 1)
+
+ out = _op.nn.dense(inputs_0, input_1)
+
+ if len(b_shape) == 1:
+ out = _op.squeeze(out, axis=[-1])
+
+ # Reshape output into a N dimensional tensor when a or b dim > 2
+ if len(a_shape) > 2:
+ out = _op.reshape(out, [*a_shape[:-1], b_shape[-1]])
+ elif len(b_shape) > 2:
+ out = _op.reshape(out, [a_shape[-2], -1, b_shape[-1]])
+ out = _op.reshape(
+ _op.transpose(out, [1, 0, 2]), [*b_shape[:-2], a_shape[-2], b_shape[-1]]
)
Review comment:
The pattern `reshape - transpose - reshape` cannot be optimized by follow-up passes. Can we use just one `reshape` and one `transpose`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org