diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 62378f54c07..d17c811df65 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -2756,6 +2756,246 @@ def _native_layer_norm_handler(P: MLXProgramBuilder, n: Node) -> Slot: return output_slots +@REGISTRY.register(target=[torch.ops.aten.native_group_norm.default]) +def _native_group_norm_handler(P: MLXProgramBuilder, n: Node) -> Slot: + """Handle native_group_norm which returns (output, mean, rstd). + + Group norm normalizes each group of ``C / group`` channels together with all + of their spatial positions, so reshaping the input to + ``(N * group, (C / group) * HxW)`` puts exactly that set on the last axis and + lets fast::layer_norm compute the normalization as a single fused kernel. + + The affine parameters are applied afterwards on the original shape rather + than being passed to layer_norm: group norm's weight and bias are per + channel, while layer_norm's are per normalized element, so the two only + coincide when every group holds a single channel. + + Only the normalized output (index 0) is computed; mean and rstd (indices 1 + and 2) are needed only for backward. + """ + unsupported = used_getitem_indices(n) & {1, 2} + if unsupported: + raise ValueError( + f"native_group_norm outputs {unsupported} (mean/rstd) are used, " + "but only the normalized output (index 0) is supported" + ) + + args = P.args(n) + require_args(args, 8, 8, "aten.native_group_norm") + require_kwargs(P.kwargs(n), set(), "aten.native_group_norm") + x, weight, bias, N, C, HxW, group, eps = args + + for name, value in (("N", N), ("C", C), ("HxW", HxW), ("group", group)): + if not isinstance(value, int): + raise ValueError( + f"aten.native_group_norm requires a static {name}, got {value!r}" + ) + + x_meta = n.args[0].meta.get("val") + if x_meta is None: + raise ValueError("aten.native_group_norm requires input shape metadata") + if any(not isinstance(d, int) for d in x_meta.shape): + raise ValueError( + f"aten.native_group_norm requires a static input shape, " + f"got {tuple(x_meta.shape)}" + ) + x_ndim = len(x_meta.shape) + + # native_group_norm returns (output, mean, rstd) -- allocate all 3 slots + output_slots = P.make_or_get_slots(n) + out = output_slots[0] + + _, flat = P.make_tmp_slot() + P.emit( + ReshapeNode( + x=P.slot_to_tid(x), + out=P.slot_to_tid(flat), + shape=[ + IntOrVid.from_literal(N * group), + IntOrVid.from_literal((C // group) * HxW), + ], + ) + ) + + _, normed = P.make_tmp_slot() + P.emit( + LayerNormNode( + x=P.slot_to_tid(flat), + out=P.slot_to_tid(normed), + weight=None, + bias=None, + eps=eps, + ) + ) + + orig_shape = [IntOrVid.from_literal(int(d)) for d in x_meta.shape] + if weight is None and bias is None: + P.emit( + ReshapeNode( + x=P.slot_to_tid(normed), + out=P.slot_to_tid(out), + shape=orig_shape, + ) + ) + return output_slots + + _, restored = P.make_tmp_slot() + P.emit( + ReshapeNode( + x=P.slot_to_tid(normed), + out=P.slot_to_tid(restored), + shape=orig_shape, + ) + ) + + # Broadcast the per-channel affine over the batch and spatial dimensions. + affine_shape = [IntOrVid.from_literal(1), IntOrVid.from_literal(C)] + [ + IntOrVid.from_literal(1) + ] * (x_ndim - 2) + + cur = restored + if weight is not None: + _, weight_bc = P.make_tmp_slot() + P.emit( + ReshapeNode( + x=P.slot_to_tid(weight), + out=P.slot_to_tid(weight_bc), + shape=affine_shape, + ) + ) + if bias is None: + scaled = out + else: + _, scaled = P.make_tmp_slot() + P.emit( + MultiplyNode( + a=P.slot_to_tid(cur), + b=P.slot_to_tid(weight_bc), + out=P.slot_to_tid(scaled), + ) + ) + cur = scaled + + if bias is not None: + _, bias_bc = P.make_tmp_slot() + P.emit( + ReshapeNode( + x=P.slot_to_tid(bias), + out=P.slot_to_tid(bias_bc), + shape=affine_shape, + ) + ) + P.emit( + AddNode( + a=P.slot_to_tid(cur), + b=P.slot_to_tid(bias_bc), + out=P.slot_to_tid(out), + ) + ) + + return output_slots + + +def _nearest_source_indices(in_size: int, out_size: int, scale: Optional[float]): + """Source index per output position for nearest-neighbour resampling. + + Mirrors aten's compute_scales_value + nearest_neighbor_compute_source_index: + the step is ``1 / scale`` when an explicit scale factor was given and + ``in_size / out_size`` otherwise, and the index is truncated then clamped to + the last input position. + """ + step = (1.0 / scale) if scale is not None else (in_size / out_size) + idx = (torch.arange(out_size, dtype=torch.float64) * step).to(torch.int64) + return torch.clamp(idx, max=in_size - 1).to(torch.int32) + + +@REGISTRY.register( + target=[ + torch.ops.aten.upsample_nearest2d.vec, + torch.ops.aten.upsample_nearest2d.default, + ] +) +def _upsample_nearest2d_handler(P: MLXProgramBuilder, n: Node) -> Slot: + """Nearest-neighbour 2D resampling as two gathers. + + Each output position reads a source position that depends only on the static + input and output sizes, so both index vectors are constants and the op is + ``take(take(x, idx_h, -2), idx_w, -1)``. Expressing it as a gather rather + than a repeat also covers non-integer scale factors and downsampling. + """ + args = P.args(n) + kwargs = P.kwargs(n) + x = args[0] + + if ".vec" in str(n.target): + require_args(args, 2, 3, "aten.upsample_nearest2d.vec") + require_kwargs(kwargs, set(), "aten.upsample_nearest2d.vec") + scale_factors = args[2] if len(args) > 2 else None + if scale_factors is None: + scales = (None, None) + else: + if len(scale_factors) != 2: + raise ValueError( + f"aten.upsample_nearest2d.vec expects 2 scale factors, " + f"got {len(scale_factors)}" + ) + scales = (scale_factors[0], scale_factors[1]) + else: + require_args(args, 2, 4, "aten.upsample_nearest2d") + require_kwargs(kwargs, set(), "aten.upsample_nearest2d") + scales = ( + args[2] if len(args) > 2 else None, + args[3] if len(args) > 3 else None, + ) + + x_meta = n.args[0].meta.get("val") + out_meta = n.meta.get("val") + if x_meta is None or out_meta is None: + raise ValueError("aten.upsample_nearest2d requires input and output shapes") + + sizes = (x_meta.shape[-2], x_meta.shape[-1], out_meta.shape[-2], out_meta.shape[-1]) + if any(not isinstance(d, int) for d in sizes): + raise ValueError( + f"aten.upsample_nearest2d requires static spatial sizes, got " + f"{tuple(x_meta.shape)} -> {tuple(out_meta.shape)}" + ) + in_h, in_w, out_h, out_w = sizes + + index_slots = [] + for axis_name, in_size, out_size, scale in ( + ("h", in_h, out_h, scales[0]), + ("w", in_w, out_w, scales[1]), + ): + indices = _nearest_source_indices(in_size, out_size, scale) + index_slots.append( + P.make_or_get_constant( + f"_upsample_nearest2d_{axis_name}_{in_size}_{out_size}_{scale}", + indices, + ) + ) + + _, rows = P.make_tmp_slot() + P.emit( + TakeNode( + x=P.slot_to_tid(x), + out=P.slot_to_tid(rows), + index=IntOrVidOrTid.from_tid(P.slot_to_tid(index_slots[0])), + axis=-2, + ) + ) + + out = P.make_or_get_slot(n) + P.emit( + TakeNode( + x=P.slot_to_tid(rows), + out=P.slot_to_tid(out), + index=IntOrVidOrTid.from_tid(P.slot_to_tid(index_slots[1])), + axis=-1, + ) + ) + return out + + @REGISTRY.register(target=[torch.ops.aten.arange.default]) def _arange_handler(P: MLXProgramBuilder, n: Node) -> Slot: """Handle arange with just stop, or (start, stop) or (start, stop, step). diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 91d6f9b130b..e7af2efb7e2 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -3332,6 +3332,136 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (x,) +class GroupNormModel(nn.Module): + """Simple model using GroupNorm.""" + + def __init__( + self, + num_groups: int = 8, + num_channels: int = 32, + eps: float = 1e-5, + affine: bool = True, + ): + super().__init__() + self.group_norm = nn.GroupNorm(num_groups, num_channels, eps=eps, affine=affine) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.group_norm(x) + + +@register_test +class GroupNormTest(OpTestCase): + """Test case for nn.GroupNorm (aten.native_group_norm).""" + + name = "group_norm" + rtol = 1e-4 + atol = 1e-4 + + def __init__( + self, + num_groups: int = 8, + num_channels: int = 32, + shape: Tuple[int, ...] = (2, 32, 8, 8), + eps: float = 1e-5, + affine: bool = True, + suffix: str = "", + ): + self.num_groups = num_groups + self.num_channels = num_channels + self.shape = shape + self.eps = eps + self.affine = affine + self.name = f"group_norm{suffix}" + + @classmethod + def get_test_configs(cls) -> List["GroupNormTest"]: + return [ + cls(), + # affine=False exercises the no-weight/no-bias path + cls(affine=False, suffix="_no_affine"), + # one channel per group (instance norm) and one group (all channels) + cls(num_groups=32, suffix="_per_channel_groups"), + cls(num_groups=1, suffix="_single_group"), + # non-square spatial extent, and a 3D (N, C, L) input + cls(num_groups=4, num_channels=16, shape=(1, 16, 5, 7), suffix="_odd"), + cls(num_groups=4, num_channels=12, shape=(2, 12, 7), suffix="_3d"), + ] + + def create_model(self) -> nn.Module: + return GroupNormModel(self.num_groups, self.num_channels, self.eps, self.affine) + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + return (torch.randn(*self.shape),) + + +class UpsampleNearest2dModel(nn.Module): + """Nearest-neighbour resize, by scale factor or by explicit output size.""" + + def __init__( + self, + scale_factor: Optional[Tuple[float, float]] = None, + size: Optional[Tuple[int, int]] = None, + ): + super().__init__() + self.scale_factor = scale_factor + self.size = size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.interpolate( + x, size=self.size, scale_factor=self.scale_factor, mode="nearest" + ) + + +@register_test +class UpsampleNearest2dTest(OpTestCase): + """Test case for aten.upsample_nearest2d.""" + + name = "upsample_nearest2d" + rtol = 0 + atol = 0 + + def __init__( + self, + shape: Tuple[int, ...] = (1, 3, 4, 4), + scale_factor: Optional[Tuple[float, float]] = (2.0, 2.0), + size: Optional[Tuple[int, int]] = None, + suffix: str = "", + ): + self.shape = shape + self.scale_factor = scale_factor + self.size = size + self.name = f"upsample_nearest2d{suffix}" + + @classmethod + def get_test_configs(cls) -> List["UpsampleNearest2dTest"]: + return [ + cls(), + # different scale per axis + cls(shape=(2, 5, 3, 7), scale_factor=(3.0, 2.0), suffix="_anisotropic"), + # non-integer ratios, which a repeat-based lowering could not express + cls(shape=(1, 3, 6, 6), scale_factor=(1.5, 2.5), suffix="_fractional"), + cls( + shape=(1, 2, 5, 5), + scale_factor=None, + size=(12, 12), + suffix="_explicit_size", + ), + # output smaller than input + cls( + shape=(1, 4, 8, 8), + scale_factor=None, + size=(4, 4), + suffix="_downsample", + ), + ] + + def create_model(self) -> nn.Module: + return UpsampleNearest2dModel(self.scale_factor, self.size) + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + return (torch.randn(*self.shape),) + + class Conv1dModel(nn.Module): """Simple model using Conv1d."""