| 162 | |
| 163 | @register_lower_rule("ConvolutionBackwardDataV2", mops.ConvolutionBackwardData) |
| 164 | def conv_backward_data_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 165 | assert ( |
| 166 | ctx.param["dilate_h"] == 1 and ctx.param["dilate_w"] == 1 |
| 167 | ), "dilate_conv is not support now" |
| 168 | |
| 169 | if len(args) == 3: |
| 170 | weight, dout, inp = args[0], args[1], args[2] |
| 171 | else: |
| 172 | weight, dout, inp = args[0], args[1], None |
| 173 | if ctx.param["format"] == mops.AdaptivePooling.Format.NCHW: |
| 174 | dnums = ((0, 1, 2, 3), (0, 1, 2, 3), (0, 1, 2, 3)) |
| 175 | inp_spec, weight_spec, out_spec = dnums |
| 176 | inp_hw, weight_hw, out_hw = map(lambda s: s[2:], dnums) |
| 177 | inp_dilation = (1, 1) |
| 178 | weight_dilation = (ctx.param["dilate_h"], ctx.param["dilate_w"]) |
| 179 | window_strides = (ctx.param["stride_h"], ctx.param["stride_w"]) |
| 180 | ph, pw = ctx.param["pad_h"], ctx.param["pad_w"] |
| 181 | padding = ((ph, ph), (pw, pw)) |
| 182 | weight_shape = weight.shape |
| 183 | inp_shape = inp.shape if inp else ctx.vars_out[0].shape |
| 184 | ic = inp_shape[1] # NCHW |
| 185 | oc = weight.shape[0] # OIHW or O11HW for dwconv |
| 186 | t_weight_spec = (weight_spec[1], weight_spec[0]) + weight_spec[2:] |
| 187 | dnums = hlo.ConvDimensionNumbers.get( |
| 188 | input_batch_dimension=out_spec[0], |
| 189 | input_feature_dimension=out_spec[1], |
| 190 | input_spatial_dimensions=list(out_spec[2:]), |
| 191 | kernel_output_feature_dimension=t_weight_spec[0], |
| 192 | kernel_input_feature_dimension=t_weight_spec[1], |
| 193 | kernel_spatial_dimensions=list(t_weight_spec[2:]), |
| 194 | output_batch_dimension=inp_spec[0], |
| 195 | output_feature_dimension=inp_spec[1], |
| 196 | output_spatial_dimensions=list(inp_spec[2:]), |
| 197 | ) |
| 198 | |
| 199 | if ctx.param["sparse"] == mops.BatchConvBias.Sparse.DENSE: |
| 200 | feature_group_count, batch_group_count = 1, 1 |
| 201 | else: |
| 202 | weight_shape = weight.shape |
| 203 | assert len(weight_shape) == 5, "mge dpconv weight dim is 5" |
| 204 | feature_group_count, batch_group_count = weight.shape[0], 1 |
| 205 | weight_shape = [ |
| 206 | weight.shape[1], |
| 207 | weight.shape[0] * weight.shape[2], |
| 208 | weight.shape[3], |
| 209 | weight.shape[4], |
| 210 | ] |
| 211 | weight = weight.transpose((1, 0, 2, 3, 4)) |
| 212 | weight = weight.reshape(weight_shape) |
| 213 | weight_shape = [ |
| 214 | weight_shape[1], |
| 215 | weight_shape[0], |
| 216 | weight_shape[2], |
| 217 | weight_shape[3], |
| 218 | ] |
| 219 | |
| 220 | padding = _conv_general_vjp_lhs_padding( |
| 221 | np.take(inp_shape, inp_hw), |