We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Fix conv2d grad for strided cases under the new conv2d_transpose def
https://ci.tvm.ai/blue/organizations/jenkins/tvm/detail/PR-6173/3/pipeline



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `python/tvm/topi/cuda/conv2d_transpose_nchw.py`
Content:
```
1 # Licensed to the Apache Software Foundation (ASF) under one
2 # or more contributor license agreements.  See the NOTICE file
3 # distributed with this work for additional information
4 # regarding copyright ownership.  The ASF licenses this file
5 # to you under the Apache License, Version 2.0 (the
6 # "License"); you may not use this file except in compliance
7 # with the License.  You may obtain a copy of the License at
8 #
9 #   http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing,
12 # software distributed under the License is distributed on an
13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 # KIND, either express or implied.  See the License for the
15 # specific language governing permissions and limitations
16 # under the License.
17 # pylint: disable=invalid-name
18 """Conv2d transpose template for cuda backend"""
19 
20 import tvm
21 from tvm import te
22 from tvm import autotvm
23 from tvm.autotvm.task.space import SplitEntity, OtherOptionEntity
24 from .. import nn
25 from ..util import get_const_tuple, traverse_inline
26 
27 
28 
29 @autotvm.register_topi_compute("conv2d_transpose_nchw.cuda")
30 def conv2d_transpose_nchw(cfg, data, kernel, stride, padding, out_dtype,
31                           output_padding):
32     """Transposed 2D convolution nchw forward operator.
33 
34     Parameters
35     ----------
36     cfg: ConfigEntity
37         The config for this template
38     Input : tvm.te.Tensor
39         4-D with shape [batch, in_channel, in_height, in_width]
40     Filter : tvm.te.Tensor
41         4-D with shape [in_channel, num_filter, filter_height, filter_width]
42     strides : tuple of two ints
43         The spatial stride along height and width
44     padding : int or str
45         Padding size, or ['VALID', 'SAME']
46     out_dtype: str
47         The output type. This is used in mixed precision
48     output_padding : tuple of two ints
49         Used to disambiguate output shape.
50 
51     Returns
52     -------
53     Output : tvm.te.Tensor
54         4-D with shape [batch, out_channel, out_height, out_width]
55     """
56     batch, inp_channels, inp_height, inp_width = get_const_tuple(data.shape)
57     _, out_channels, kernel_height, kernel_width = get_const_tuple(kernel.shape)
58     stride_height, stride_width = stride
59     outpad_height, outpad_width = output_padding
60     assert outpad_height < stride_height and outpad_width < stride_width
61     cfg.stride = stride
62     pad_top, pad_left, pad_bottom, pad_right = nn.get_pad_tuple(
63         padding, (kernel_height, kernel_width))
64 
65     out_width = (inp_width - 1) * stride_width + \
66         kernel_width - pad_left - pad_right + outpad_width
67     pad_left = kernel_width - 1 - pad_left
68     pad_right = kernel_width - 1 - pad_right
69     dilated_width = stride_width * (inp_width - 1) + 1
70 
71     out_height = (inp_height - 1) * stride_height + \
72         kernel_height - pad_top - pad_bottom + outpad_height
73     pad_top = kernel_height - 1 - pad_top
74     pad_bottom = kernel_height - 1 - pad_bottom
75     dilated_height = stride_height * (inp_height - 1) + 1
76 
77     # compute pad
78     data = te.compute(
79         (batch, inp_channels,
80          pad_top + dilated_height + pad_bottom,
81          pad_left + dilated_width + pad_right),
82         lambda n, c, y, x: tvm.tir.if_then_else(
83             tvm.tir.all(x >= pad_left,
84                         x < pad_left + dilated_width,
85                         tvm.tir.indexmod(x - pad_left, stride_width).equal(0),
86                         y >= pad_top,
87                         y < pad_top + dilated_height,
88                         tvm.tir.indexmod(y - pad_top, stride_height).equal(0)),
89             data[n, c,
90                  tvm.tir.indexdiv(y - pad_top, stride_height),
91                  tvm.tir.indexdiv(x - pad_left, stride_width)],
92             tvm.tir.const(0., "float32")),
93         name='data_pad')
94 
95     # compute transposed conv
96     dc = te.reduce_axis((0, inp_channels), name='dc')
97     dh = te.reduce_axis((0, kernel_height), name='dh')
98     dw = te.reduce_axis((0, kernel_width), name='dw')
99     data_out = te.compute(
100         (batch, out_channels, out_height, out_width),
101         lambda b, c, h, w: te.sum(
102             data[b, dc, h + dh, w + dw].astype(out_dtype) *
103             kernel[dc,
104                    c,
105                    kernel_height - 1 - dh,
106                    kernel_width - 1 - dw].astype(out_dtype),
107             axis=[dc, dh, dw]), tag="conv2d_transpose_nchw")
108 
109     return data_out
110 
111 @autotvm.register_topi_schedule("conv2d_transpose_nchw.cuda")
112 def schedule_conv2d_transpose_nchw(cfg, outs):
113     """TOPI Schedule callback for conv2d transpose operator.
114 
115     Parameters
116     ----------
117     cfg: ConfigEntity
118         The parameters for this template
119 
120     outs: Array of Tensor
121         The computation graph description of conv2d transpose
122         in the format of an array of tensors.
123 
124     Returns
125     -------
126     s: Schedule
127         The computation schedule for conv2d transpose.
128     """
129     outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
130     s = te.create_schedule([x.op for x in outs])
131 
132     def _fallback_schedule(N, F, Y, X):
133         # pylint: disable=unused-argument
134         # split N (batch dimension)
135         if N > 1:
136             cfg["tile_n"] = SplitEntity([-1, 1, 1, 4])
137         else:
138             cfg["tile_n"] = SplitEntity([1, 1, 1, 1])
139         # split F (output channel dimension)
140         if F > 1:
141             cfg["tile_f"] = SplitEntity([-1, 1, 64, 1])
142         # split Y (height dimension)
143         y_split_factor = 1
144         for candidate in range(5, 17):
145             if Y % candidate == 0:
146                 y_split_factor = candidate
147                 break
148         cfg["tile_y"] = SplitEntity([-1, 1, 1, y_split_factor])
149         # split X (width dimension)
150         x_split_factor = 1
151         for candidate in range(5, 17):
152             if X % candidate == 0:
153                 x_split_factor = candidate
154                 break
155         cfg["tile_x"] = SplitEntity([-1, x_split_factor, 1, 1])
156         # split RC (input channel dimension, which is a reduction axis)
157         cfg["tile_rc"] = SplitEntity([-1, 1, 16])
158         # other configurations
159         cfg["fuse_yx"] = OtherOptionEntity(False)
160         cfg["unroll_explicit"] = OtherOptionEntity(True)
161         cfg["auto_unroll_max_step"] = OtherOptionEntity(1500)
162 
163     def _callback(op):
164         if op.tag == 'conv2d_transpose_nchw':
165             pad_data = op.input_tensors[0]
166             kernel = op.input_tensors[1]
167             conv = op.output(0)
168 
169             ##### space definition begin #####
170             n, f, y, x = s[conv].op.axis
171             rc = s[conv].op.reduce_axis[0]
172             cfg.define_split("tile_n", cfg.axis(n), num_outputs=4)
173             cfg.define_split("tile_f", cfg.axis(f), num_outputs=4)
174             cfg.define_split("tile_y", cfg.axis(y), num_outputs=4)
175             cfg.define_split("tile_x", cfg.axis(x), num_outputs=4)
176             cfg.define_split("tile_rc", cfg.axis(rc), num_outputs=3)
177             cfg.define_knob("auto_unroll_max_step", [64, 512, 1500])
178 
179             target = tvm.target.Target.current()
180             if target.kind.name in ['nvptx', 'rocm']:
181                 cfg.define_knob("unroll_explicit", [1])
182             else:
183                 cfg.define_knob("unroll_explicit", [0, 1])
184 
185             if cfg.is_fallback:
186                 N, F, Y, X = get_const_tuple(conv.shape)
187                 _fallback_schedule(N, F, Y, X)
188 
189             ##### space definition end #####
190 
191             if isinstance(kernel.op, tvm.te.ComputeOp) and 'dilate' in kernel.op.tag:
192                 s[kernel].compute_inline()
193 
194             if conv.op in s.outputs:
195                 output = conv
196                 OL = s.cache_write(conv, 'local')
197             else:
198                 output = s.outputs[0].output(0)
199                 s[conv].set_scope('local')
200                 OL = conv
201 
202             # create cache stage
203             s[pad_data].set_scope('shared')
204             AA = pad_data
205             WW = s.cache_read(kernel, 'shared', [OL])
206 
207             # tile and bind spatial axes
208             n, f, y, x = s[output].op.axis
209             kernel_scope, n = s[output].split(n, nparts=1)
210             bn, vn, tn, ni = cfg["tile_n"].apply(s, output, n)
211             bf, vf, tf, fi = cfg["tile_f"].apply(s, output, f)
212             by, vy, ty, yi = cfg["tile_y"].apply(s, output, y)
213             bx, vx, tx, xi = cfg["tile_x"].apply(s, output, x)
214 
215             s[output].reorder(bn, bf, by, bx, vn, vf, vy, vx, tn, tf, ty, tx, ni, fi, yi, xi)
216             s[output].bind(bn, te.thread_axis("blockIdx.z"))
217             s[output].bind(bf, te.thread_axis("blockIdx.y"))
218             s[output].bind(s[output].fuse(by, bx), te.thread_axis("blockIdx.x"))
219             s[output].bind(vn, te.thread_axis("vthread"))
220             s[output].bind(vf, te.thread_axis("vthread"))
221             s[output].bind(vy, te.thread_axis("vthread"))
222             s[output].bind(vx, te.thread_axis("vthread"))
223 
224             cfg.define_knob("fuse_yx", [0, 1]) # fuse ty,tx or tn,tf
225 
226             if cfg["fuse_yx"].val:
227                 s[output].bind(tn, te.thread_axis("threadIdx.z"))
228                 s[output].bind(tf, te.thread_axis("threadIdx.y"))
229                 tyx = s[output].fuse(ty, tx)
230                 s[output].bind(s[output].fuse(ty, tx), te.thread_axis("threadIdx.x"))
231                 s[OL].compute_at(s[output], tyx)
232 
233                 # number of threads
234                 n_tz = cfg["tile_n"].size[2]
235                 n_ty = cfg["tile_f"].size[2]
236                 n_tx = cfg["tile_y"].size[2] * cfg["tile_x"].size[2]
237             else:
238                 s[output].bind(s[output].fuse(tn, tf), te.thread_axis("threadIdx.z"))
239                 s[output].bind(ty, te.thread_axis("threadIdx.y"))
240                 s[output].bind(tx, te.thread_axis("threadIdx.x"))
241                 s[OL].compute_at(s[output], tx)
242 
243                 # number of threads
244                 n_tz = cfg["tile_n"].size[2] * cfg["tile_f"].size[2]
245                 n_ty = cfg["tile_y"].size[2]
246                 n_tx = cfg["tile_x"].size[2]
247 
248             # tile reduction axes
249             n, f, y, x = s[OL].op.axis
250             rc, ry, rx = s[OL].op.reduce_axis
251             rco, rcm, rci = cfg['tile_rc'].apply(s, OL, rc)
252             s[OL].reorder(rco, rcm, ry, rx, rci, n, f, y, x)
253 
254             s[AA].compute_at(s[OL], rx)
255             s[WW].compute_at(s[OL], rx)
256 
257             # cooperative fetching
258             for load in [AA, WW]:
259                 n, f, y, x = s[load].op.axis
260                 fused = s[load].fuse(f, y, x)
261                 tz, fused = s[load].split(fused, nparts=n_tz)
262                 ty, fused = s[load].split(fused, nparts=n_ty)
263                 tx, fused = s[load].split(fused, nparts=n_tx)
264                 s[load].bind(tz, te.thread_axis("threadIdx.z"))
265                 s[load].bind(ty, te.thread_axis("threadIdx.y"))
266                 s[load].bind(tx, te.thread_axis("threadIdx.x"))
267 
268             s[output].pragma(kernel_scope, 'auto_unroll_max_step', cfg['auto_unroll_max_step'].val)
269             s[output].pragma(kernel_scope, 'unroll_explicit', cfg['unroll_explicit'].val)
270 
271     traverse_inline(s, outs[0].op, _callback)
272 
273     return s
274
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/python/tvm/topi/cuda/conv2d_transpose_nchw.py b/python/tvm/topi/cuda/conv2d_transpose_nchw.py
--- a/python/tvm/topi/cuda/conv2d_transpose_nchw.py
+++ b/python/tvm/topi/cuda/conv2d_transpose_nchw.py
@@ -65,13 +65,13 @@
     out_width = (inp_width - 1) * stride_width + \
         kernel_width - pad_left - pad_right + outpad_width
     pad_left = kernel_width - 1 - pad_left
-    pad_right = kernel_width - 1 - pad_right
+    pad_right = kernel_width - 1 - pad_right + outpad_width
     dilated_width = stride_width * (inp_width - 1) + 1
 
     out_height = (inp_height - 1) * stride_height + \
         kernel_height - pad_top - pad_bottom + outpad_height
     pad_top = kernel_height - 1 - pad_top
-    pad_bottom = kernel_height - 1 - pad_bottom
+    pad_bottom = kernel_height - 1 - pad_bottom + outpad_height
     dilated_height = stride_height * (inp_height - 1) + 1
 
     # compute pad

{"golden_diff": "diff --git a/python/tvm/topi/cuda/conv2d_transpose_nchw.py b/python/tvm/topi/cuda/conv2d_transpose_nchw.py\n--- a/python/tvm/topi/cuda/conv2d_transpose_nchw.py\n+++ b/python/tvm/topi/cuda/conv2d_transpose_nchw.py\n@@ -65,13 +65,13 @@\n     out_width = (inp_width - 1) * stride_width + \\\n         kernel_width - pad_left - pad_right + outpad_width\n     pad_left = kernel_width - 1 - pad_left\n-    pad_right = kernel_width - 1 - pad_right\n+    pad_right = kernel_width - 1 - pad_right + outpad_width\n     dilated_width = stride_width * (inp_width - 1) + 1\n \n     out_height = (inp_height - 1) * stride_height + \\\n         kernel_height - pad_top - pad_bottom + outpad_height\n     pad_top = kernel_height - 1 - pad_top\n-    pad_bottom = kernel_height - 1 - pad_bottom\n+    pad_bottom = kernel_height - 1 - pad_bottom + outpad_height\n     dilated_height = stride_height * (inp_height - 1) + 1\n \n     # compute pad\n", "issue": "Fix conv2d grad for strided cases under the new conv2d_transpose def\nhttps://ci.tvm.ai/blue/organizations/jenkins/tvm/detail/PR-6173/3/pipeline\r\n\r\n\n", "before_files": [{"content": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements.  See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership.  The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License.  You may obtain a copy of the License at\n#\n#   http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied.  See the License for the\n# specific language governing permissions and limitations\n# under the License.\n# pylint: disable=invalid-name\n\"\"\"Conv2d transpose template for cuda backend\"\"\"\n\nimport tvm\nfrom tvm import te\nfrom tvm import autotvm\nfrom tvm.autotvm.task.space import SplitEntity, OtherOptionEntity\nfrom .. import nn\nfrom ..util import get_const_tuple, traverse_inline\n\n\n\n@autotvm.register_topi_compute(\"conv2d_transpose_nchw.cuda\")\ndef conv2d_transpose_nchw(cfg, data, kernel, stride, padding, out_dtype,\n                          output_padding):\n    \"\"\"Transposed 2D convolution nchw forward operator.\n\n    Parameters\n    ----------\n    cfg: ConfigEntity\n        The config for this template\n    Input : tvm.te.Tensor\n        4-D with shape [batch, in_channel, in_height, in_width]\n    Filter : tvm.te.Tensor\n        4-D with shape [in_channel, num_filter, filter_height, filter_width]\n    strides : tuple of two ints\n        The spatial stride along height and width\n    padding : int or str\n        Padding size, or ['VALID', 'SAME']\n    out_dtype: str\n        The output type. This is used in mixed precision\n    output_padding : tuple of two ints\n        Used to disambiguate output shape.\n\n    Returns\n    -------\n    Output : tvm.te.Tensor\n        4-D with shape [batch, out_channel, out_height, out_width]\n    \"\"\"\n    batch, inp_channels, inp_height, inp_width = get_const_tuple(data.shape)\n    _, out_channels, kernel_height, kernel_width = get_const_tuple(kernel.shape)\n    stride_height, stride_width = stride\n    outpad_height, outpad_width = output_padding\n    assert outpad_height < stride_height and outpad_width < stride_width\n    cfg.stride = stride\n    pad_top, pad_left, pad_bottom, pad_right = nn.get_pad_tuple(\n        padding, (kernel_height, kernel_width))\n\n    out_width = (inp_width - 1) * stride_width + \\\n        kernel_width - pad_left - pad_right + outpad_width\n    pad_left = kernel_width - 1 - pad_left\n    pad_right = kernel_width - 1 - pad_right\n    dilated_width = stride_width * (inp_width - 1) + 1\n\n    out_height = (inp_height - 1) * stride_height + \\\n        kernel_height - pad_top - pad_bottom + outpad_height\n    pad_top = kernel_height - 1 - pad_top\n    pad_bottom = kernel_height - 1 - pad_bottom\n    dilated_height = stride_height * (inp_height - 1) + 1\n\n    # compute pad\n    data = te.compute(\n        (batch, inp_channels,\n         pad_top + dilated_height + pad_bottom,\n         pad_left + dilated_width + pad_right),\n        lambda n, c, y, x: tvm.tir.if_then_else(\n            tvm.tir.all(x >= pad_left,\n                        x < pad_left + dilated_width,\n                        tvm.tir.indexmod(x - pad_left, stride_width).equal(0),\n                        y >= pad_top,\n                        y < pad_top + dilated_height,\n                        tvm.tir.indexmod(y - pad_top, stride_height).equal(0)),\n            data[n, c,\n                 tvm.tir.indexdiv(y - pad_top, stride_height),\n                 tvm.tir.indexdiv(x - pad_left, stride_width)],\n            tvm.tir.const(0., \"float32\")),\n        name='data_pad')\n\n    # compute transposed conv\n    dc = te.reduce_axis((0, inp_channels), name='dc')\n    dh = te.reduce_axis((0, kernel_height), name='dh')\n    dw = te.reduce_axis((0, kernel_width), name='dw')\n    data_out = te.compute(\n        (batch, out_channels, out_height, out_width),\n        lambda b, c, h, w: te.sum(\n            data[b, dc, h + dh, w + dw].astype(out_dtype) *\n            kernel[dc,\n                   c,\n                   kernel_height - 1 - dh,\n                   kernel_width - 1 - dw].astype(out_dtype),\n            axis=[dc, dh, dw]), tag=\"conv2d_transpose_nchw\")\n\n    return data_out\n\n@autotvm.register_topi_schedule(\"conv2d_transpose_nchw.cuda\")\ndef schedule_conv2d_transpose_nchw(cfg, outs):\n    \"\"\"TOPI Schedule callback for conv2d transpose operator.\n\n    Parameters\n    ----------\n    cfg: ConfigEntity\n        The parameters for this template\n\n    outs: Array of Tensor\n        The computation graph description of conv2d transpose\n        in the format of an array of tensors.\n\n    Returns\n    -------\n    s: Schedule\n        The computation schedule for conv2d transpose.\n    \"\"\"\n    outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs\n    s = te.create_schedule([x.op for x in outs])\n\n    def _fallback_schedule(N, F, Y, X):\n        # pylint: disable=unused-argument\n        # split N (batch dimension)\n        if N > 1:\n            cfg[\"tile_n\"] = SplitEntity([-1, 1, 1, 4])\n        else:\n            cfg[\"tile_n\"] = SplitEntity([1, 1, 1, 1])\n        # split F (output channel dimension)\n        if F > 1:\n            cfg[\"tile_f\"] = SplitEntity([-1, 1, 64, 1])\n        # split Y (height dimension)\n        y_split_factor = 1\n        for candidate in range(5, 17):\n            if Y % candidate == 0:\n                y_split_factor = candidate\n                break\n        cfg[\"tile_y\"] = SplitEntity([-1, 1, 1, y_split_factor])\n        # split X (width dimension)\n        x_split_factor = 1\n        for candidate in range(5, 17):\n            if X % candidate == 0:\n                x_split_factor = candidate\n                break\n        cfg[\"tile_x\"] = SplitEntity([-1, x_split_factor, 1, 1])\n        # split RC (input channel dimension, which is a reduction axis)\n        cfg[\"tile_rc\"] = SplitEntity([-1, 1, 16])\n        # other configurations\n        cfg[\"fuse_yx\"] = OtherOptionEntity(False)\n        cfg[\"unroll_explicit\"] = OtherOptionEntity(True)\n        cfg[\"auto_unroll_max_step\"] = OtherOptionEntity(1500)\n\n    def _callback(op):\n        if op.tag == 'conv2d_transpose_nchw':\n            pad_data = op.input_tensors[0]\n            kernel = op.input_tensors[1]\n            conv = op.output(0)\n\n            ##### space definition begin #####\n            n, f, y, x = s[conv].op.axis\n            rc = s[conv].op.reduce_axis[0]\n            cfg.define_split(\"tile_n\", cfg.axis(n), num_outputs=4)\n            cfg.define_split(\"tile_f\", cfg.axis(f), num_outputs=4)\n            cfg.define_split(\"tile_y\", cfg.axis(y), num_outputs=4)\n            cfg.define_split(\"tile_x\", cfg.axis(x), num_outputs=4)\n            cfg.define_split(\"tile_rc\", cfg.axis(rc), num_outputs=3)\n            cfg.define_knob(\"auto_unroll_max_step\", [64, 512, 1500])\n\n            target = tvm.target.Target.current()\n            if target.kind.name in ['nvptx', 'rocm']:\n                cfg.define_knob(\"unroll_explicit\", [1])\n            else:\n                cfg.define_knob(\"unroll_explicit\", [0, 1])\n\n            if cfg.is_fallback:\n                N, F, Y, X = get_const_tuple(conv.shape)\n                _fallback_schedule(N, F, Y, X)\n\n            ##### space definition end #####\n\n            if isinstance(kernel.op, tvm.te.ComputeOp) and 'dilate' in kernel.op.tag:\n                s[kernel].compute_inline()\n\n            if conv.op in s.outputs:\n                output = conv\n                OL = s.cache_write(conv, 'local')\n            else:\n                output = s.outputs[0].output(0)\n                s[conv].set_scope('local')\n                OL = conv\n\n            # create cache stage\n            s[pad_data].set_scope('shared')\n            AA = pad_data\n            WW = s.cache_read(kernel, 'shared', [OL])\n\n            # tile and bind spatial axes\n            n, f, y, x = s[output].op.axis\n            kernel_scope, n = s[output].split(n, nparts=1)\n            bn, vn, tn, ni = cfg[\"tile_n\"].apply(s, output, n)\n            bf, vf, tf, fi = cfg[\"tile_f\"].apply(s, output, f)\n            by, vy, ty, yi = cfg[\"tile_y\"].apply(s, output, y)\n            bx, vx, tx, xi = cfg[\"tile_x\"].apply(s, output, x)\n\n            s[output].reorder(bn, bf, by, bx, vn, vf, vy, vx, tn, tf, ty, tx, ni, fi, yi, xi)\n            s[output].bind(bn, te.thread_axis(\"blockIdx.z\"))\n            s[output].bind(bf, te.thread_axis(\"blockIdx.y\"))\n            s[output].bind(s[output].fuse(by, bx), te.thread_axis(\"blockIdx.x\"))\n            s[output].bind(vn, te.thread_axis(\"vthread\"))\n            s[output].bind(vf, te.thread_axis(\"vthread\"))\n            s[output].bind(vy, te.thread_axis(\"vthread\"))\n            s[output].bind(vx, te.thread_axis(\"vthread\"))\n\n            cfg.define_knob(\"fuse_yx\", [0, 1]) # fuse ty,tx or tn,tf\n\n            if cfg[\"fuse_yx\"].val:\n                s[output].bind(tn, te.thread_axis(\"threadIdx.z\"))\n                s[output].bind(tf, te.thread_axis(\"threadIdx.y\"))\n                tyx = s[output].fuse(ty, tx)\n                s[output].bind(s[output].fuse(ty, tx), te.thread_axis(\"threadIdx.x\"))\n                s[OL].compute_at(s[output], tyx)\n\n                # number of threads\n                n_tz = cfg[\"tile_n\"].size[2]\n                n_ty = cfg[\"tile_f\"].size[2]\n                n_tx = cfg[\"tile_y\"].size[2] * cfg[\"tile_x\"].size[2]\n            else:\n                s[output].bind(s[output].fuse(tn, tf), te.thread_axis(\"threadIdx.z\"))\n                s[output].bind(ty, te.thread_axis(\"threadIdx.y\"))\n                s[output].bind(tx, te.thread_axis(\"threadIdx.x\"))\n                s[OL].compute_at(s[output], tx)\n\n                # number of threads\n                n_tz = cfg[\"tile_n\"].size[2] * cfg[\"tile_f\"].size[2]\n                n_ty = cfg[\"tile_y\"].size[2]\n                n_tx = cfg[\"tile_x\"].size[2]\n\n            # tile reduction axes\n            n, f, y, x = s[OL].op.axis\n            rc, ry, rx = s[OL].op.reduce_axis\n            rco, rcm, rci = cfg['tile_rc'].apply(s, OL, rc)\n            s[OL].reorder(rco, rcm, ry, rx, rci, n, f, y, x)\n\n            s[AA].compute_at(s[OL], rx)\n            s[WW].compute_at(s[OL], rx)\n\n            # cooperative fetching\n            for load in [AA, WW]:\n                n, f, y, x = s[load].op.axis\n                fused = s[load].fuse(f, y, x)\n                tz, fused = s[load].split(fused, nparts=n_tz)\n                ty, fused = s[load].split(fused, nparts=n_ty)\n                tx, fused = s[load].split(fused, nparts=n_tx)\n                s[load].bind(tz, te.thread_axis(\"threadIdx.z\"))\n                s[load].bind(ty, te.thread_axis(\"threadIdx.y\"))\n                s[load].bind(tx, te.thread_axis(\"threadIdx.x\"))\n\n            s[output].pragma(kernel_scope, 'auto_unroll_max_step', cfg['auto_unroll_max_step'].val)\n            s[output].pragma(kernel_scope, 'unroll_explicit', cfg['unroll_explicit'].val)\n\n    traverse_inline(s, outs[0].op, _callback)\n\n    return s\n", "path": "python/tvm/topi/cuda/conv2d_transpose_nchw.py"}], "after_files": [{"content": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements.  See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership.  The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License.  You may obtain a copy of the License at\n#\n#   http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied.  See the License for the\n# specific language governing permissions and limitations\n# under the License.\n# pylint: disable=invalid-name\n\"\"\"Conv2d transpose template for cuda backend\"\"\"\n\nimport tvm\nfrom tvm import te\nfrom tvm import autotvm\nfrom tvm.autotvm.task.space import SplitEntity, OtherOptionEntity\nfrom .. import nn\nfrom ..util import get_const_tuple, traverse_inline\n\n\n\n@autotvm.register_topi_compute(\"conv2d_transpose_nchw.cuda\")\ndef conv2d_transpose_nchw(cfg, data, kernel, stride, padding, out_dtype,\n                          output_padding):\n    \"\"\"Transposed 2D convolution nchw forward operator.\n\n    Parameters\n    ----------\n    cfg: ConfigEntity\n        The config for this template\n    Input : tvm.te.Tensor\n        4-D with shape [batch, in_channel, in_height, in_width]\n    Filter : tvm.te.Tensor\n        4-D with shape [in_channel, num_filter, filter_height, filter_width]\n    strides : tuple of two ints\n        The spatial stride along height and width\n    padding : int or str\n        Padding size, or ['VALID', 'SAME']\n    out_dtype: str\n        The output type. This is used in mixed precision\n    output_padding : tuple of two ints\n        Used to disambiguate output shape.\n\n    Returns\n    -------\n    Output : tvm.te.Tensor\n        4-D with shape [batch, out_channel, out_height, out_width]\n    \"\"\"\n    batch, inp_channels, inp_height, inp_width = get_const_tuple(data.shape)\n    _, out_channels, kernel_height, kernel_width = get_const_tuple(kernel.shape)\n    stride_height, stride_width = stride\n    outpad_height, outpad_width = output_padding\n    assert outpad_height < stride_height and outpad_width < stride_width\n    cfg.stride = stride\n    pad_top, pad_left, pad_bottom, pad_right = nn.get_pad_tuple(\n        padding, (kernel_height, kernel_width))\n\n    out_width = (inp_width - 1) * stride_width + \\\n        kernel_width - pad_left - pad_right + outpad_width\n    pad_left = kernel_width - 1 - pad_left\n    pad_right = kernel_width - 1 - pad_right + outpad_width\n    dilated_width = stride_width * (inp_width - 1) + 1\n\n    out_height = (inp_height - 1) * stride_height + \\\n        kernel_height - pad_top - pad_bottom + outpad_height\n    pad_top = kernel_height - 1 - pad_top\n    pad_bottom = kernel_height - 1 - pad_bottom + outpad_height\n    dilated_height = stride_height * (inp_height - 1) + 1\n\n    # compute pad\n    data = te.compute(\n        (batch, inp_channels,\n         pad_top + dilated_height + pad_bottom,\n         pad_left + dilated_width + pad_right),\n        lambda n, c, y, x: tvm.tir.if_then_else(\n            tvm.tir.all(x >= pad_left,\n                        x < pad_left + dilated_width,\n                        tvm.tir.indexmod(x - pad_left, stride_width).equal(0),\n                        y >= pad_top,\n                        y < pad_top + dilated_height,\n                        tvm.tir.indexmod(y - pad_top, stride_height).equal(0)),\n            data[n, c,\n                 tvm.tir.indexdiv(y - pad_top, stride_height),\n                 tvm.tir.indexdiv(x - pad_left, stride_width)],\n            tvm.tir.const(0., \"float32\")),\n        name='data_pad')\n\n    # compute transposed conv\n    dc = te.reduce_axis((0, inp_channels), name='dc')\n    dh = te.reduce_axis((0, kernel_height), name='dh')\n    dw = te.reduce_axis((0, kernel_width), name='dw')\n    data_out = te.compute(\n        (batch, out_channels, out_height, out_width),\n        lambda b, c, h, w: te.sum(\n            data[b, dc, h + dh, w + dw].astype(out_dtype) *\n            kernel[dc,\n                   c,\n                   kernel_height - 1 - dh,\n                   kernel_width - 1 - dw].astype(out_dtype),\n            axis=[dc, dh, dw]), tag=\"conv2d_transpose_nchw\")\n\n    return data_out\n\n@autotvm.register_topi_schedule(\"conv2d_transpose_nchw.cuda\")\ndef schedule_conv2d_transpose_nchw(cfg, outs):\n    \"\"\"TOPI Schedule callback for conv2d transpose operator.\n\n    Parameters\n    ----------\n    cfg: ConfigEntity\n        The parameters for this template\n\n    outs: Array of Tensor\n        The computation graph description of conv2d transpose\n        in the format of an array of tensors.\n\n    Returns\n    -------\n    s: Schedule\n        The computation schedule for conv2d transpose.\n    \"\"\"\n    outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs\n    s = te.create_schedule([x.op for x in outs])\n\n    def _fallback_schedule(N, F, Y, X):\n        # pylint: disable=unused-argument\n        # split N (batch dimension)\n        if N > 1:\n            cfg[\"tile_n\"] = SplitEntity([-1, 1, 1, 4])\n        else:\n            cfg[\"tile_n\"] = SplitEntity([1, 1, 1, 1])\n        # split F (output channel dimension)\n        if F > 1:\n            cfg[\"tile_f\"] = SplitEntity([-1, 1, 64, 1])\n        # split Y (height dimension)\n        y_split_factor = 1\n        for candidate in range(5, 17):\n            if Y % candidate == 0:\n                y_split_factor = candidate\n                break\n        cfg[\"tile_y\"] = SplitEntity([-1, 1, 1, y_split_factor])\n        # split X (width dimension)\n        x_split_factor = 1\n        for candidate in range(5, 17):\n            if X % candidate == 0:\n                x_split_factor = candidate\n                break\n        cfg[\"tile_x\"] = SplitEntity([-1, x_split_factor, 1, 1])\n        # split RC (input channel dimension, which is a reduction axis)\n        cfg[\"tile_rc\"] = SplitEntity([-1, 1, 16])\n        # other configurations\n        cfg[\"fuse_yx\"] = OtherOptionEntity(False)\n        cfg[\"unroll_explicit\"] = OtherOptionEntity(True)\n        cfg[\"auto_unroll_max_step\"] = OtherOptionEntity(1500)\n\n    def _callback(op):\n        if op.tag == 'conv2d_transpose_nchw':\n            pad_data = op.input_tensors[0]\n            kernel = op.input_tensors[1]\n            conv = op.output(0)\n\n            ##### space definition begin #####\n            n, f, y, x = s[conv].op.axis\n            rc = s[conv].op.reduce_axis[0]\n            cfg.define_split(\"tile_n\", cfg.axis(n), num_outputs=4)\n            cfg.define_split(\"tile_f\", cfg.axis(f), num_outputs=4)\n            cfg.define_split(\"tile_y\", cfg.axis(y), num_outputs=4)\n            cfg.define_split(\"tile_x\", cfg.axis(x), num_outputs=4)\n            cfg.define_split(\"tile_rc\", cfg.axis(rc), num_outputs=3)\n            cfg.define_knob(\"auto_unroll_max_step\", [64, 512, 1500])\n\n            target = tvm.target.Target.current()\n            if target.kind.name in ['nvptx', 'rocm']:\n                cfg.define_knob(\"unroll_explicit\", [1])\n            else:\n                cfg.define_knob(\"unroll_explicit\", [0, 1])\n\n            if cfg.is_fallback:\n                N, F, Y, X = get_const_tuple(conv.shape)\n                _fallback_schedule(N, F, Y, X)\n\n            ##### space definition end #####\n\n            if isinstance(kernel.op, tvm.te.ComputeOp) and 'dilate' in kernel.op.tag:\n                s[kernel].compute_inline()\n\n            if conv.op in s.outputs:\n                output = conv\n                OL = s.cache_write(conv, 'local')\n            else:\n                output = s.outputs[0].output(0)\n                s[conv].set_scope('local')\n                OL = conv\n\n            # create cache stage\n            s[pad_data].set_scope('shared')\n            AA = pad_data\n            WW = s.cache_read(kernel, 'shared', [OL])\n\n            # tile and bind spatial axes\n            n, f, y, x = s[output].op.axis\n            kernel_scope, n = s[output].split(n, nparts=1)\n            bn, vn, tn, ni = cfg[\"tile_n\"].apply(s, output, n)\n            bf, vf, tf, fi = cfg[\"tile_f\"].apply(s, output, f)\n            by, vy, ty, yi = cfg[\"tile_y\"].apply(s, output, y)\n            bx, vx, tx, xi = cfg[\"tile_x\"].apply(s, output, x)\n\n            s[output].reorder(bn, bf, by, bx, vn, vf, vy, vx, tn, tf, ty, tx, ni, fi, yi, xi)\n            s[output].bind(bn, te.thread_axis(\"blockIdx.z\"))\n            s[output].bind(bf, te.thread_axis(\"blockIdx.y\"))\n            s[output].bind(s[output].fuse(by, bx), te.thread_axis(\"blockIdx.x\"))\n            s[output].bind(vn, te.thread_axis(\"vthread\"))\n            s[output].bind(vf, te.thread_axis(\"vthread\"))\n            s[output].bind(vy, te.thread_axis(\"vthread\"))\n            s[output].bind(vx, te.thread_axis(\"vthread\"))\n\n            cfg.define_knob(\"fuse_yx\", [0, 1]) # fuse ty,tx or tn,tf\n\n            if cfg[\"fuse_yx\"].val:\n                s[output].bind(tn, te.thread_axis(\"threadIdx.z\"))\n                s[output].bind(tf, te.thread_axis(\"threadIdx.y\"))\n                tyx = s[output].fuse(ty, tx)\n                s[output].bind(s[output].fuse(ty, tx), te.thread_axis(\"threadIdx.x\"))\n                s[OL].compute_at(s[output], tyx)\n\n                # number of threads\n                n_tz = cfg[\"tile_n\"].size[2]\n                n_ty = cfg[\"tile_f\"].size[2]\n                n_tx = cfg[\"tile_y\"].size[2] * cfg[\"tile_x\"].size[2]\n            else:\n                s[output].bind(s[output].fuse(tn, tf), te.thread_axis(\"threadIdx.z\"))\n                s[output].bind(ty, te.thread_axis(\"threadIdx.y\"))\n                s[output].bind(tx, te.thread_axis(\"threadIdx.x\"))\n                s[OL].compute_at(s[output], tx)\n\n                # number of threads\n                n_tz = cfg[\"tile_n\"].size[2] * cfg[\"tile_f\"].size[2]\n                n_ty = cfg[\"tile_y\"].size[2]\n                n_tx = cfg[\"tile_x\"].size[2]\n\n            # tile reduction axes\n            n, f, y, x = s[OL].op.axis\n            rc, ry, rx = s[OL].op.reduce_axis\n            rco, rcm, rci = cfg['tile_rc'].apply(s, OL, rc)\n            s[OL].reorder(rco, rcm, ry, rx, rci, n, f, y, x)\n\n            s[AA].compute_at(s[OL], rx)\n            s[WW].compute_at(s[OL], rx)\n\n            # cooperative fetching\n            for load in [AA, WW]:\n                n, f, y, x = s[load].op.axis\n                fused = s[load].fuse(f, y, x)\n                tz, fused = s[load].split(fused, nparts=n_tz)\n                ty, fused = s[load].split(fused, nparts=n_ty)\n                tx, fused = s[load].split(fused, nparts=n_tx)\n                s[load].bind(tz, te.thread_axis(\"threadIdx.z\"))\n                s[load].bind(ty, te.thread_axis(\"threadIdx.y\"))\n                s[load].bind(tx, te.thread_axis(\"threadIdx.x\"))\n\n            s[output].pragma(kernel_scope, 'auto_unroll_max_step', cfg['auto_unroll_max_step'].val)\n            s[output].pragma(kernel_scope, 'unroll_explicit', cfg['unroll_explicit'].val)\n\n    traverse_inline(s, outs[0].op, _callback)\n\n    return s\n", "path": "python/tvm/topi/cuda/conv2d_transpose_nchw.py"}]}