| 17 | |
| 18 | @register_lower_rule(mops.Convolution) |
| 19 | def convolution_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 20 | assert isinstance(ctx.op, mops.Convolution) |
| 21 | assert len(args) == 2, "convolution requires 2 arguments" |
| 22 | assert len(ctx.vars_in) == 2, "convolution requires 2 input variables" |
| 23 | assert len(ctx.vars_out) == 1, "convolution requires 1 output variable" |
| 24 | |
| 25 | opr = ctx.op |
| 26 | inp, weight = args[0], args[1] |
| 27 | |
| 28 | if opr.format == mops.AdaptivePooling.Format.NCHW: |
| 29 | inp_spec, weight_spec, out_spec = (0, 1, 2, 3), (0, 1, 2, 3), (0, 1, 2, 3) |
| 30 | dnums = hlo.ConvDimensionNumbers.get( |
| 31 | input_batch_dimension=inp_spec[0], |
| 32 | input_feature_dimension=inp_spec[1], |
| 33 | input_spatial_dimensions=list(inp_spec[2:]), |
| 34 | kernel_output_feature_dimension=weight_spec[0], |
| 35 | kernel_input_feature_dimension=weight_spec[1], |
| 36 | kernel_spatial_dimensions=list(weight_spec[2:]), |
| 37 | output_batch_dimension=out_spec[0], |
| 38 | output_feature_dimension=out_spec[1], |
| 39 | output_spatial_dimensions=list(out_spec[2:]), |
| 40 | ) |
| 41 | ic = inp.shape[1] # NCHW |
| 42 | oc = weight.shape[0] # OIHW or O11HW for dwconv |
| 43 | else: |
| 44 | assert False, "only nchw supported" |
| 45 | |
| 46 | num_spatial_dims = len(weight_spec) - 2 |
| 47 | window_reversal = ir_utils.dense_bool_elements([False] * num_spatial_dims) |
| 48 | |
| 49 | if opr.sparse == mops.BatchConvBias.Sparse.DENSE: |
| 50 | feature_group_count, batch_group_count = 1, 1 |
| 51 | else: |
| 52 | assert len(weight.shape) == 5, "mge dpconv weight dim is 5" |
| 53 | feature_group_count, batch_group_count = weight.shape[0], 1 |
| 54 | |
| 55 | if opr.format == mops.AdaptivePooling.Format.NCHW: |
| 56 | xla_weight_shape = xla_weight_shape = [ |
| 57 | weight.shape[0] * weight.shape[1], |
| 58 | weight.shape[2], |
| 59 | weight.shape[3], |
| 60 | weight.shape[4], |
| 61 | ] |
| 62 | weight = reshape(weight, xla_weight_shape) |
| 63 | |
| 64 | feature_group_count = ir_utils.i64_attr(feature_group_count) |
| 65 | batch_group_count = ir_utils.i64_attr(batch_group_count) |
| 66 | |
| 67 | window_strides = (opr.stride_h, opr.stride_w) |
| 68 | window_strides = ir_utils.dense_int_elements(window_strides) |
| 69 | |
| 70 | padding = ((opr.pad_h, opr.pad_h), (opr.pad_w, opr.pad_w)) |
| 71 | padding = ir_utils.dense_int_elements(padding) |
| 72 | |
| 73 | assert opr.dilate_h == 1 and opr.dilate_w == 1, "dilate_conv is not support now" |
| 74 | inp_dilation = (opr.dilate_h, opr.dilate_w) |
| 75 | weight_dilation = (opr.dilate_h, opr.dilate_w) |
| 76 | inp_dilation = ir_utils.dense_int_elements(inp_dilation) |