| 126 | |
| 127 | @register_lower_rule(mops.Reduce) |
| 128 | def reduce_lower(ctx, *args: Union[ir.Value, Sequence[ir.Value]]): |
| 129 | assert ctx.op.data_type == mops.Reduce.DataType.DEFAULT |
| 130 | |
| 131 | opr = ctx.op |
| 132 | keepdims = opr.keepdim |
| 133 | if len(args) == 1: |
| 134 | assert isinstance(opr.axis, int) |
| 135 | if opr.axis < 0: |
| 136 | axes = opr.axis + args[0].ndim |
| 137 | else: |
| 138 | axes = (opr.axis,) |
| 139 | if opr.axis > 7: |
| 140 | axes = tuple(np.arange(args[0].ndim)) |
| 141 | keepdims = False |
| 142 | else: |
| 143 | assert len(args) == 2 |
| 144 | src_shape = args[0].shape |
| 145 | if src_shape == ctx.vars_out[0].shape: |
| 146 | return args[0] |
| 147 | tgt_shape = list(ctx.vars_out[0].shape) |
| 148 | tgt_shape = [1,] * (len(src_shape) - len(tgt_shape)) + tgt_shape |
| 149 | src_idx, tgt_idx, axes = 0, 0, [] |
| 150 | while src_idx < len(src_shape) and tgt_idx < len(tgt_shape): |
| 151 | if src_shape[src_idx] != 1 and tgt_shape[tgt_idx] == 1: |
| 152 | axes.append(src_idx) |
| 153 | src_idx = src_idx + 1 |
| 154 | tgt_idx = tgt_idx + 1 |
| 155 | elif src_shape[src_idx] != tgt_shape[tgt_idx]: |
| 156 | axes.append(src_idx) |
| 157 | src_idx = src_idx + 1 |
| 158 | else: |
| 159 | src_idx = src_idx + 1 |
| 160 | tgt_idx = tgt_idx + 1 |
| 161 | assert tgt_idx == len( |
| 162 | tgt_shape |
| 163 | ), f"src_shape: {src_shape}, tgt_shape: {tgt_shape}" |
| 164 | axes = axes + list(range(src_idx, len(src_shape))) |
| 165 | |
| 166 | if opr.mode == mops.Reduce.Mode.SUM: |
| 167 | ret = sum(args[0], axes, keepdims) |
| 168 | elif opr.mode == mops.Reduce.Mode.MEAN: |
| 169 | ret = mean(args[0], axes, keepdims) |
| 170 | elif opr.mode == mops.Reduce.Mode.PRODUCT: |
| 171 | ret = prod(args[0], axes, keepdims) |
| 172 | elif opr.mode == mops.Reduce.Mode.MAX: |
| 173 | ret = max(args[0], axes, keepdims) |
| 174 | elif opr.mode == mops.Reduce.Mode.MIN: |
| 175 | ret = min(args[0], axes, keepdims) |
| 176 | else: |
| 177 | assert False, f"no support reduce mode {opr.mode}" |
| 178 | |
| 179 | if not _shape_equal(ret.shape, ctx.vars_out[0].shape): |
| 180 | ret = ret.reshape(ctx.vars_out[0].shape) |
| 181 | return ret |