From 5131d6e4aa91c6bae08f67fd2c523710593322ff Mon Sep 17 00:00:00 2001 From: gokulkrishna98 Date: Fri, 12 Jun 2026 08:30:32 -0700 Subject: [PATCH] _aten_to_core: implement masked_scatter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Register a lowering for aten.masked_scatter.default. Semantics: out = self.clone() out[mask] = source.flatten()[: mask.sum()] Implementation flattens self/mask/source, then for each flat position picks from source[cumsum(mask) - 1] if mask is True else self. False positions select away the gathered values via where(), so out-of-range indices there are harmless. Indices for True positions are guaranteed in [0, mask.sum() - 1] by construction. Supports dynamic shapes by reshaping with the runtime shape vector (coreai.get_shape) instead of the type-level shape, which carries a sentinel for unknown dims, and by using -1 in the flat reshapes to let coreai.reshape infer the numel. Mask broadcast also routes through the runtime shape when self is dynamic. Adds TestMaskedScatter covering: - Static and dynamic IR FileCheck (full operand chain pinned for static; op order asserted for dynamic). - Numerical validation across {f32, f16, i32} × {static, dynamic}. - Corner cases: all-False mask (output == self), all-True mask (output == src reshape), src larger than mask.sum() (extras unused), and a lower-rank mask that broadcasts right-aligned onto self. --- coreai_torch/_aten_to_core.py | 41 +++++++++ tests/ops/test_ops.py | 159 ++++++++++++++++++++++++++++++++++ 2 files changed, 200 insertions(+) diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 74e18a6..ba0c4b1 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -1985,6 +1985,46 @@ def replace_logical_not( return coreai.not_(x) +def replace_masked_scatter( + values_map: dict[str, Value], node: fx.Node, loc: Location +) -> Value: + """Lowers aten.masked_scatter.default to: + + out_flat = where(mask, source[cumsum(mask) - 1], self) + + Out-of-range indices at False positions are harmless — where() + discards them. + """ + self_, mask, source = _get_operands(values_map, node, [0, 1, 2]) + + self_flat = coreai.reshape(self_, [-1]) + if mask.type.shape != self_.type.shape: + target_shape = ( + coreai.get_shape(self_) + if any(d < 0 for d in self_.type.shape) + else list(self_.type.shape) + ) + mask = coreai.broadcast_to(mask, target_shape) + mask_flat = coreai.reshape(mask, [-1]) + source_flat = coreai.reshape(source, [-1]) + + if source_flat.type.element_type != self_flat.type.element_type: + source_flat = coreai.cast(source_flat, self_flat.type.element_type) + + mask_int = coreai.cast(mask_flat, np.int32) + idx = coreai.scan(mask_int, np.uint32(0), False, combiner="sum") + idx = coreai.broadcasting_sub(idx, coreai.constant(1, dtype=np.int32)) + gathered = coreai.gather_along_axis(source_flat, idx, 0) + + out_flat = coreai.broadcasting_where(mask_flat, gathered, self_flat) + out_shape = ( + coreai.get_shape(self_) + if any(d < 0 for d in self_.type.shape) + else list(self_.type.shape) + ) + return coreai.reshape(out_flat, out_shape) + + def replace_maxpool2d_with_indices( values_map: dict[str, Value], node: fx.Node, loc: Location ) -> OpResultList: @@ -3408,6 +3448,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value: "max.default": replace_max_default, "max.dim": replace_max_dim, "maximum.default": replace_binary_ops, + "masked_scatter.default": replace_masked_scatter, "mean.default": replace_mean_default, "mean.dim": replace_mean_dim, "min.default": replace_min_default, diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index ddfd4ac..d1f9230 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -4066,6 +4066,165 @@ def forward(self, x: Tensor, index: Tensor, src: Tensor) -> Tensor: ) +class TestMaskedScatter: + """Tests for aten.masked_scatter.default → coreai flatten / cumsum-1 / + gather / where lowering. + """ + + class _Model(nn.Module): + def forward(self, self_: Tensor, mask: Tensor, src: Tensor) -> Tensor: + return torch.masked_scatter(self_, mask, src) + + @staticmethod + def _shared_dynamic_shapes(self_: Tensor, src: Tensor) -> dict: + """Build dynamic_shapes where self_ and mask share dim objects + (torch.export emits an equality guard when traced shapes match) + and src has its own independent dim. + """ + dims = {i: torch.export.Dim(f"d{i}", min=1) for i in range(self_.dim())} + return { + "self_": dims, + "mask": dims, + "src": {0: torch.export.Dim("s0", min=1)} + if src.dim() == 1 + else _all_dims_dynamic(src, "s"), + } + + @pytest.mark.ir + def test_lowers_ir_static(self) -> None: + """Pin the full operand chain on a static shape.""" + self_ = torch.zeros(2, 3) + mask = torch.tensor([[True, False, True], [False, True, False]]) + src = torch.arange(1.0, 7.0) + program = torch.export.export( + self._Model(), args=(self_, mask, src) + ).run_decompositions() + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + filecheck_pattern( + str(coreai_program), + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK-SAME: %arg0: tensor<2x3xf32> + // CHECK-SAME: %arg1: tensor<2x3xi1> + // CHECK-SAME: %arg2: tensor<6xf32> + // CHECK: %[[SF:.+]] = coreai.reshape %arg0, %{{.+}} : (tensor<2x3xf32>, tensor<1xsi32>) -> tensor<6xf32> + // CHECK: %[[MF:.+]] = coreai.reshape %arg1, %{{.+}} : (tensor<2x3xi1>, tensor<1xsi32>) -> tensor<6xi1> + // CHECK: %[[VF:.+]] = coreai.reshape %arg2, %{{.+}} : (tensor<6xf32>, tensor<1xsi32>) -> tensor<6xf32> + // CHECK: %[[MI:.+]] = coreai.cast %[[MF]] : tensor<6xi1> to tensor<6xsi32> + // CHECK: %[[CS:.+]] = coreai.scan %[[MI]], %{{.+}}, %{{.+}} combiner = : (tensor<6xsi32>, tensor, tensor) -> tensor<6xsi32> + // CHECK: %[[IDX:.+]] = coreai.decomposable.broadcasting_sub %[[CS]], %{{.+}} : (tensor<6xsi32>, tensor) -> tensor<6xsi32> + // CHECK: %[[G:.+]] = coreai.gather_along_axis %[[VF]] at %[[IDX]] along %{{.+}} : (tensor<6xf32>, tensor<6xsi32>, tensor) to tensor<6xf32> + // CHECK: %[[W:.+]] = coreai.decomposable.broadcasting_where %[[MF]], %[[G]], %[[SF]] : (tensor<6xi1>, tensor<6xf32>, tensor<6xf32>) -> tensor<6xf32> + // CHECK: %[[OUT:.+]] = coreai.reshape %[[W]], %{{.+}} : (tensor<6xf32>, tensor<2xsi32>) -> tensor<2x3xf32> + // CHECK: coreai.output %[[OUT]] + """, + ) + + @pytest.mark.ir + def test_lowers_ir_dynamic(self) -> None: + """Confirm the same op chain appears under dynamic shapes — shapes + become symbolic, but cast/scan/sub/gather/where/reshape order is + preserved. + """ + self_ = torch.zeros(2, 3) + mask = torch.zeros(2, 3, dtype=torch.bool) + src = torch.arange(6.0) + dynamic_shapes = self._shared_dynamic_shapes(self_, src) + program = torch.export.export( + self._Model(), args=(self_, mask, src), dynamic_shapes=dynamic_shapes + ).run_decompositions() + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + filecheck_pattern( + str(coreai_program), + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK: coreai.reshape + // CHECK: coreai.reshape + // CHECK: coreai.reshape + // CHECK: coreai.cast {{.+}} to tensor<{{.*}}xsi32> + // CHECK: coreai.scan {{.+}} combiner = + // CHECK: coreai.decomposable.broadcasting_sub + // CHECK: coreai.gather_along_axis + // CHECK: coreai.decomposable.broadcasting_where + // CHECK: coreai.reshape + // CHECK: coreai.output + """, + ) + + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.int32]) + async def test_basic(self, dtype: torch.dtype, dynamic: bool) -> None: + """self.shape == mask.shape, src is 1-D of length mask.numel().""" + self_ = torch.zeros(2, 3, dtype=dtype) + mask = torch.tensor([[True, False, True], [False, True, False]]) + src = torch.arange(1, 7, dtype=dtype) + dynamic_shapes = self._shared_dynamic_shapes(self_, src) if dynamic else None + await validate_numerical_output( + model=self._Model().eval(), + self_=self_, + mask=mask, + src=src, + dynamic_shapes=dynamic_shapes, + ) + + @pytest.mark.parametrize("dynamic", [False, True]) + async def test_all_false_mask(self, dynamic: bool) -> None: + """No positions selected — output must equal self unchanged.""" + self_ = torch.arange(6.0).reshape(2, 3) + mask = torch.zeros(2, 3, dtype=torch.bool) + src = torch.full((6,), -1.0) + dynamic_shapes = self._shared_dynamic_shapes(self_, src) if dynamic else None + await validate_numerical_output( + model=self._Model().eval(), + self_=self_, + mask=mask, + src=src, + dynamic_shapes=dynamic_shapes, + ) + + @pytest.mark.parametrize("dynamic", [False, True]) + async def test_all_true_mask(self, dynamic: bool) -> None: + """All positions selected — output must equal src reshaped to + self.shape, with self values fully overwritten. + """ + self_ = torch.zeros(2, 3) + mask = torch.ones(2, 3, dtype=torch.bool) + src = torch.arange(1.0, 7.0) + dynamic_shapes = self._shared_dynamic_shapes(self_, src) if dynamic else None + await validate_numerical_output( + model=self._Model().eval(), + self_=self_, + mask=mask, + src=src, + dynamic_shapes=dynamic_shapes, + ) + + async def test_src_larger_than_mask_sum(self) -> None: + """src may carry more elements than mask.sum(); only the leading + mask.sum() are consumed (read flat). The rest are unused. + """ + self_ = torch.zeros(2, 3) + mask = torch.tensor([[True, False, True], [False, True, False]]) + # 10 elements, only first 3 are needed (mask has 3 True positions) + src = torch.arange(1.0, 11.0) + await validate_numerical_output( + model=self._Model().eval(), self_=self_, mask=mask, src=src + ) + + async def test_broadcast_mask_lower_rank(self) -> None: + """mask is rank-1 (3,) and self is rank-2 (2, 3) — mask is + right-aligned and broadcasts across the leading dim. Exercises + the broadcast_to branch in the lowering. + """ + self_ = torch.zeros(2, 3) + mask = torch.tensor([True, False, True]) + # mask broadcasts to (2, 3) with 4 True positions + src = torch.arange(1.0, 5.0) + await validate_numerical_output( + model=self._Model().eval(), self_=self_, mask=mask, src=src + ) + + @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "x,dim,index",