From 9651ef58e925f42eb0740e8410e93c73c4b01df2 Mon Sep 17 00:00:00 2001 From: henry Date: Mon, 6 Jul 2026 13:15:01 -0700 Subject: [PATCH] Add missing tests for the pas fix --- coreai_torch/_utils.py | 7 +++- tests/ops/test_ops.py | 90 ++++++++++++++++++++++++++++++++++++++++ tests/ops/test_ops_ir.py | 84 +++++++++++++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 2 deletions(-) diff --git a/coreai_torch/_utils.py b/coreai_torch/_utils.py index 0e2ac88..a4fbb73 100644 --- a/coreai_torch/_utils.py +++ b/coreai_torch/_utils.py @@ -46,6 +46,7 @@ from ._composite_declaration import generate_composite_decl from ._type_mapping import ( TORCH_TO_COREAI_DTYPE, + _get_coreai_to_numpy_dtype, _get_coreai_to_torch_dtype, ) @@ -1055,11 +1056,13 @@ def replace_pad_with_mode( padding[2 * dim] = inverted_padding[i] padding[2 * dim + 1] = inverted_padding[i + 1] - # padding_value is ignored for non-constant modes, but the op requires one. + # padding_value is ignored for non-constant modes, but the op requires it to + # be a constant of the input dtype (a cast op is rejected by the backend). + np_dtype = _get_coreai_to_numpy_dtype()[x.type.element_type] return coreai.pad( x, np.array(padding, dtype=np.uint32), - coreai.cast(0.0, x.type.element_type), + coreai.constant(np.array(0.0, dtype=np_dtype)), padding_mode=padding_mode, ) diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index a8dfdb6..c20d17e 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -7331,3 +7331,93 @@ def forward( remove_decomps=[torch.ops.aten.scaled_dot_product_attention.default], **kwargs, ) + + +# ndim (number of padded spatial dims) -> the aten op that must be preserved +# so the reflect/replicate lowering (coreai.pad) is exercised end to end. +_REFLECT_PAD_OP = { + 1: torch.ops.aten.reflection_pad1d.default, + 2: torch.ops.aten.reflection_pad2d.default, + 3: torch.ops.aten.reflection_pad3d.default, +} +_REPLICATE_PAD_OP = { + 1: torch.ops.aten.replication_pad1d.default, + 2: torch.ops.aten.replication_pad2d.default, + 3: torch.ops.aten.replication_pad3d.default, +} + + +class _PadModel(nn.Module): + def __init__(self, pad: tuple[int, ...], mode: str) -> None: + super().__init__() + self._pad = pad + self._mode = mode + + def forward(self, x: Tensor) -> Tensor: + return torch.nn.functional.pad(x, self._pad, mode=self._mode) + + +# (pad, input_shape) pairs valid for BOTH reflect and replicate. +# For reflect, torch requires every pad entry < the corresponding dim size. +_PAD_SHARED_CASES = [ + # 1D (pad = (left, right)), input (N, C, W) + ((2, 2), (1, 3, 8)), # symmetric + ((1, 3), (2, 4, 10)), # asymmetric + ((0, 2), (1, 1, 6)), # one-sided + # 2D (pad = (left, right, top, bottom)), input (N, C, H, W) + ((2, 2, 2, 2), (1, 3, 8, 8)), # symmetric + ((1, 2, 3, 1), (2, 3, 10, 12)), # asymmetric, per-side + ((0, 1, 1, 0), (1, 4, 7, 9)), # mixed zero/one-sided + # 3D (pad = (l, r, t, b, front, back)), input (N, C, D, H, W) + ((1, 1, 1, 1, 1, 1), (1, 2, 4, 6, 6)), # symmetric + ((2, 1, 0, 2, 1, 1), (1, 2, 5, 7, 7)), # asymmetric +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("mode", ["reflect", "replicate"]) +@pytest.mark.parametrize("pad, input_shape", _PAD_SHARED_CASES) +async def test_reflect_replicate_pad( + pad: tuple[int, ...], + input_shape: tuple[int, ...], + mode: str, + dtype: torch.dtype, +) -> None: + ndim = len(pad) // 2 + aten_op = (_REFLECT_PAD_OP if mode == "reflect" else _REPLICATE_PAD_OP)[ndim] + await validate_numerical_output( + model=_PadModel(pad, mode).eval(), + x=torch.rand(*input_shape, dtype=dtype), + remove_decomps=[aten_op], + ) + + +@pytest.mark.parametrize("mode", ["reflect", "replicate"]) +async def test_reflect_replicate_pad_dynamic_batch(mode: str) -> None: + """Padding amounts and spatial dims are static; only the batch is dynamic.""" + aten_op = (_REFLECT_PAD_OP if mode == "reflect" else _REPLICATE_PAD_OP)[2] + await validate_numerical_output( + model=_PadModel((2, 2, 2, 2), mode).eval(), + x=torch.rand(2, 3, 8, 8), + dynamic_shapes={"x": {0: torch.export.Dim("batch", min=1)}}, + remove_decomps=[aten_op], + ) + + +@pytest.mark.parametrize( + "pad, input_shape", + [ + ((4, 4), (1, 2, 3)), # 1D: pad exceeds the padded dim + ((5, 5, 5, 5), (1, 2, 3, 3)), # 2D: pad exceeds both spatial dims + ], +) +async def test_replicate_pad_larger_than_dim( + pad: tuple[int, ...], input_shape: tuple[int, ...] +) -> None: + """Replicate (unlike reflect) allows padding wider than the input dim.""" + ndim = len(pad) // 2 + await validate_numerical_output( + model=_PadModel(pad, "replicate").eval(), + x=torch.rand(*input_shape), + remove_decomps=[_REPLICATE_PAD_OP[ndim]], + ) diff --git a/tests/ops/test_ops_ir.py b/tests/ops/test_ops_ir.py index 59005b4..95e5b20 100644 --- a/tests/ops/test_ops_ir.py +++ b/tests/ops/test_ops_ir.py @@ -10,6 +10,8 @@ import torch.nn as nn from torch import Tensor +import coreai_torch + from ..utils import _all_dims_dynamic, filecheck_pattern, get_ir @@ -7238,3 +7240,85 @@ def body_fn(x): // CHECK-NEXT: } """, ) + + +class _PadModel(nn.Module): + def __init__(self, pad: tuple[int, ...], mode: str) -> None: + super().__init__() + self._pad = pad + self._mode = mode + + def forward(self, x: Tensor) -> Tensor: + return torch.nn.functional.pad(x, self._pad, mode=self._mode) + + +class TestPadModesIR: + """reflect/replicate padding must lower to coreai.pad, not coreai.gather_nd. + + ``remove_decomps`` keeps the pad op in the graph exactly as + ``get_decomp_table()`` does in production, so the direct lowering is + exercised (otherwise torch decomposes it into index/gather ops). + """ + + @pytest.mark.parametrize( + "mode, aten_op, coreai_mode", + [ + ("reflect", torch.ops.aten.reflection_pad2d.default, "reflect"), + ("replicate", torch.ops.aten.replication_pad2d.default, "replicate"), + ], + ) + def test_lowers_to_coreai_pad_not_gather( + self, mode: str, aten_op, coreai_mode: str + ) -> None: + ir = get_ir( + _PadModel((2, 2, 2, 2), mode).eval(), + x=torch.rand(1, 3, 8, 8), + remove_decomps=[aten_op], + ) + assert "coreai.pad" in ir, ir + assert f"mode = <{coreai_mode}>" in ir, ir + assert "gather_nd" not in ir, ir + + def test_reflect_pad_ir_structure(self) -> None: + ir = get_ir( + _PadModel((2, 2, 2, 2), "reflect").eval(), + x=torch.rand(1, 3, 8, 8), + remove_decomps=[torch.ops.aten.reflection_pad2d.default], + ) + filecheck_pattern( + ir, + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK: %[[PAD:.*]] = coreai.pad %[[X:.*]], %{{.*}}, %{{.*}} mode = + // CHECK-SAME: -> tensor<1x3x12x12xf32> + // CHECK: coreai.output %[[PAD]] + """, + ) + + def test_reflection_pad1d(self) -> None: + ir = get_ir( + _PadModel((2, 2), "reflect").eval(), + x=torch.rand(1, 3, 8), + remove_decomps=[torch.ops.aten.reflection_pad1d.default], + ) + assert "coreai.pad" in ir, ir + assert "mode = " in ir, ir + assert "gather_nd" not in ir, ir + + +class TestPadDecompTable: + """The decomposition table must preserve the pad ops so the handler fires.""" + + @pytest.mark.parametrize( + "aten_op", + [ + torch.ops.aten.reflection_pad1d.default, + torch.ops.aten.reflection_pad2d.default, + torch.ops.aten.reflection_pad3d.default, + torch.ops.aten.replication_pad1d.default, + torch.ops.aten.replication_pad2d.default, + torch.ops.aten.replication_pad3d.default, + ], + ) + def test_pad_ops_preserved(self, aten_op) -> None: + assert aten_op not in coreai_torch.get_decomp_table()