anijain2305 commented on a change in pull request #4399: [Relay][Legalize]
Legalize conv2d_transpose for NHWC
URL: https://github.com/apache/incubator-tvm/pull/4399#discussion_r349830244
##########
File path: topi/python/topi/nn/conv2d_transpose.py
##########
@@ -102,3 +103,61 @@ def declaration_conv2d_transpose_impl(data, kernel,
strides, padding, out_dtype)
axis=[dc, dh, dw]), tag="conv2d_transpose_nchw")
return Output
+
+
+@tvm.target.generic_func
+def conv2d_transpose_legalize(attrs, inputs, types):
+ """Legalizes Transposed 2D convolution op.
+
+ Parameters
+ ----------
+ attrs : tvm.attrs.Attrs
+ Attributes of current Transposed 2D convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
+ The legalized expr
+ """
+ if attrs['data_layout'] == 'NHWC':
+ data, kernel = inputs
+ kernel_layout = attrs['kernel_layout']
+ # Convert Kernel layout to IOHW
+ # kernel_layout is different from input kernel layout - IO is swapped
+ if kernel_layout == 'HWIO':
+ # input kernel layout is swapped to HWOI
+ # output kernel layout will be IOHW
+ kernel = relay.transpose(kernel, axes=(3, 2, 0, 1))
+ elif kernel_layout == 'HWOI':
+ # input kernel layout is swapped to HWIO
+ # output kernel layout will be IOHW
+ kernel = relay.transpose(kernel, axes=(2, 3, 0, 1))
+ elif kernel_layout == 'IOHW':
+ # input kernel layout is swapped to OIHW
+ # output kernel layout will be IOHW
+ kernel = relay.transpose(kernel, axes=(1, 0, 2, 3))
+ elif kernel_layout == 'OIHW':
+ # input kernel layout is swapped to IOHW
+ # output kernel layout will be IOHW
+ pass
+ else:
+ raise ValueError("Invalid kernel_layout {}".format(kernel_layout))
Review comment:
We can do `return None` here. That is if some topi schedule supports this
particular combination of layouts, it will fallback to that.
Return None basically means legalize does not transform anything.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
