feat: add distributed ops AllGatherMatmul and MatmulReduceScatter (MC2)-hyper-parallel-AtomGit

Ahyper_parallel/core/shard/ops/parallel_all_gather_matmul.py+238-0

@@ -0,0 +1,238 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Distributed implementation for AllGatherMatmul (MC2 fusion) operator."""
	16	+import copy
	17	+from typing import Tuple
	18	+
	19	+from hyper_parallel.core.dtensor.layout import Layout
	20	+from .parallel_ops import DistributedOp
	21	+
	22	+
	23	+def _normalize_agm_args(
	24	+ x1,
	25	+ x2,
	26	+ group,
	27	+ world_size,
	28	+ bias=None,
	29	+ gather_index=0,
	30	+ gather_output=True,
	31	+ comm_turn=0,
	32	+ trans_input=False,
	33	+ trans_x2=False,
	34	+) -> Tuple[tuple, dict]:
	35	+ """Normalize positional and keyword arguments into a canonical positional tuple.
	36	+
	37	+ Args:
	38	+ x1: Input tensor. Physical shape (m, k); trans_input=False only (MS constraint).
	39	+ x2: Weight tensor with shape (k, n) if trans_x2=False, or (n, k) if trans_x2=True.
	40	+ group: Communication group name string.
	41	+ world_size: Number of ranks in the communication group.
	42	+ bias: Must be None (MS constraint; bias is not supported).
	43	+ gather_index: Index for gather operation (only 0 is supported).
	44	+ gather_output: Whether to return the gathered intermediate tensor.
	45	+ comm_turn: Communication turn (only 0 is supported).
	46	+ trans_input: Must be False (MS constraint).
	47	+ trans_x2: If True, x2 physical shape is (n, k); MindSpore transposes x2 before CANN call.
	48	+
	49	+ Returns:
	50	+ tuple: (positional_args_tuple, empty_kwargs_dict)
	51	+ """
	52	+ return (x1, x2, group, world_size, bias, gather_index, gather_output, comm_turn, trans_input, trans_x2), {}
	53	+
	54	+
	55	+class AllGatherMatmulDistributedOp(DistributedOp):
	56	+ """Distributed operator for mindspore.ops.all_gather_matmul (MC2 fusion).
	57	+
	58	+ The CANN AllGatherMatmul kernel handles communication (AllGather) internally.
	59	+ HyperParallel's role is to:
	60	+ 1. Extract local tensors from DTensor inputs via to_local().
	61	+ 2. Infer output DTensor layouts so downstream operators can correctly
	62	+ understand the distribution state of the output.
	63	+
	64	+ Shape transformation (logical, after any transpose):
	65	+ x1 (m_local, k_local) —[CANN AllGather on m]→ (m_global, k_local) —[matmul]→ output (m_global, n_local)
	66	+ gather_out (m_global, k_local) — valid only when gather_output=True
	67	+
	68	+ Sharding constraints:
	69	+ - x1 physical (m, k): dim 0 is m (AllGather consumes m); dim 1 (k) may be Replicate or Shard.
	70	+ - x1's m-dim tensor_map must not be a tuple (joint sharding across multiple mesh dims unsupported).
	71	+ - x1 k-dim and x2 k-dim must share the same placement (both Replicate, or both sharded on the
	72	+ same mesh axis). trans_x2=False: x2 k is dim 0; trans_x2=True: x2 k is dim 1.
	73	+ - When k is sharded, output carries Partial(sum) status on the k-dim mesh axis; the caller is
	74	+ responsible for applying AllReduce to obtain the correct full result.
	75	+ - gather_out k-dim follows x1's k-dim placement.
	76	+ - Partial inputs are not allowed.
	77	+ - gather_index=0, trans_input=False, bias=None only (current MS constraints).
	78	+ """
	79	+
	80	+ def preprocess(self, args: tuple, kwargs: dict) -> tuple:
	81	+ """Extract local tensors and build the layout cache.
	82	+
	83	+ Args:
	84	+ args: Positional arguments (DTensors for x1 and x2).
	85	+ kwargs: Keyword arguments.
	86	+
	87	+ Returns:
	88	+ tuple: (local_args, local_kwargs, cache_values) where
	89	+ local_args contains extracted local tensors,
	90	+ local_kwargs contains keyword arguments,
	91	+ cache_values = [x1_layout, x2_layout, trans_x2, gather_output].
	92	+ """
	93	+ norm_args, _ = _normalize_agm_args(args, *kwargs)
	94	+ x1 = norm_args[0]
	95	+ x2 = norm_args[1]
	96	+ gather_output = norm_args[6]
	97	+ trans_x2 = norm_args[9]
	98	+
	99	+ # MindSpore all_gather_matmul: only (input, x2, group, world_size) are
	100	+ # positional; bias, gather_index, gather_output, comm_turn, trans_input,
	101	+ # trans_x2 are keyword-only (after the '*' separator).
	102	+ local_args = (
	103	+ x1.to_local(),
	104	+ x2.to_local(),
	105	+ norm_args[2], # group
	106	+ norm_args[3], # world_size
	107	+ )
	108	+ local_kwargs = {
	109	+ 'bias': norm_args[4],
	110	+ 'gather_index': norm_args[5],
	111	+ 'gather_output': gather_output,
	112	+ 'comm_turn': norm_args[7],
	113	+ 'trans_input': norm_args[8],
	114	+ 'trans_x2': trans_x2,
	115	+ }
	116	+
	117	+ cache_values = [
	118	+ x1.layout,
	119	+ x2.layout,
	120	+ trans_x2,
	121	+ gather_output,
	122	+ ]
	123	+ return local_args, local_kwargs, cache_values
	124	+
	125	+ @staticmethod
	126	+ def _set_partial_from_k(output_layout: Layout, k_placement, op: str = 'sum') -> None:
	127	+ """Set Partial on output_layout for the mesh axes corresponding to k_placement.
	128	+
	129	+ Args:
	130	+ output_layout: Layout to mark as Partial.
	131	+ k_placement: Tensor_map value for the k dimension (integer or tuple of integers).
	132	+ op: Reduction operation, default 'sum'.
	133	+ """
	134	+ alias = output_layout.alias_name
	135	+ n = len(alias)
	136	+ if isinstance(k_placement, tuple):
	137	+ for v in k_placement:
	138	+ output_layout.set_partial_by_dev_axis(alias[n - 1 - v], op)
	139	+ else:
	140	+ output_layout.set_partial_by_dev_axis(alias[n - 1 - k_placement], op)
	141	+
	142	+ @staticmethod
	143	+ def _validate_input_layouts(
	144	+ x1_layout: Layout,
	145	+ x2_layout: Layout,
	146	+ trans_x2: bool,
	147	+ ) -> None:
	148	+ """Validate sharding constraints for AllGatherMatmul inputs.
	149	+
	150	+ Args:
	151	+ x1_layout: Layout of x1. Physical (m, k); trans_input=False only.
	152	+ x2_layout: Layout of x2 (k, n) if trans_x2=False, or (n, k) if trans_x2=True.
	153	+ trans_x2: Whether x2 is transposed.
	154	+
	155	+ Raises:
	156	+ ValueError: If x1's m-dim tensor_map is a tuple, the k-dim placements of x1 and x2
	157	+ do not match, or any input has Partial status.
	158	+ """
	159	+ op = "all_gather_matmul"
	160	+ x1_tm = x1_layout.tensor_map
	161	+ x2_tm = x2_layout.tensor_map
	162	+
	163	+ # trans_input=False only: x1 physical (m, k) — k is dim 1, m is dim 0.
	164	+ x1_m_dim = 0
	165	+
	166	+ if isinstance(x1_tm[x1_m_dim], tuple):
	167	+ raise ValueError(
	168	+ f"For {op}, x1 m-dim (dim {x1_m_dim}) "
	169	+ f"with tensor_map={x1_tm[x1_m_dim]} is jointly sharded across multiple "
	170	+ f"mesh dims, which is not supported in this version."
	171	+ )
	172	+
	173	+ # k-dim placement must match between x1 and x2.
	174	+ x2_k_dim = 1 if trans_x2 else 0
	175	+ if x1_tm[1] != x2_tm[x2_k_dim]:
	176	+ raise ValueError(
	177	+ f"For {op}, x1 k-dim (dim 1) placement {x1_tm[1]} must match "
	178	+ f"x2 k-dim (dim {x2_k_dim}) placement {x2_tm[x2_k_dim]} "
	179	+ f"(trans_x2={trans_x2})."
	180	+ )
	181	+
	182	+ def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
	183	+ """Infer output layouts for (output, gather_out).
	184	+
	185	+ AllGather on m dim: output dim 0 is always -1 (Replicate), because
	186	+ AllGather unconditionally makes the m dimension global.
	187	+
	188	+ n dim: follows x2's n placement.
	189	+ - trans_x2=False: n is x2 dim 1 → output_tm[1] = x2_tm[1]
	190	+ - trans_x2=True: n is x2 dim 0 → output_tm[1] = x2_tm[0]
	191	+
	192	+ k dim (contraction): when k is sharded, output carries Partial(sum) on the
	193	+ k-dim mesh axis; the caller must apply AllReduce to get the correct result.
	194	+
	195	+ gather_out layout: m is Replicate (-1); k follows x1's k-dim placement.
	196	+
	197	+ Args:
	198	+ cache_values: [x1_layout, x2_layout, trans_x2, gather_output]
	199	+
	200	+ Returns:
	201	+ tuple: ((output_layout, gather_out_layout), None)
	202	+
	203	+ Raises:
	204	+ ValueError: If any input has Partial status or sharding constraints are violated.
	205	+ """
	206	+ x1_layout = cache_values[0]
	207	+ x2_layout = cache_values[1]
	208	+ trans_x2 = cache_values[2]
	209	+ gather_output = cache_values[3]
	210	+
	211	+ self._check_partial_inputs([x1_layout, x2_layout])
	212	+ self._validate_input_layouts(x1_layout, x2_layout, trans_x2)
	213	+
	214	+ x1_tm = x1_layout.tensor_map
	215	+ x2_tm = x2_layout.tensor_map
	216	+ k_placement = x1_tm[1]
	217	+ n_placement = x2_tm[0] if trans_x2 else x2_tm[1]
	218	+
	219	+ # output: m is Replicate (-1) because AllGather consumed the m sharding;
	220	+ # n inherits from x2's n dim.
	221	+ output_layout = Layout.from_device_mesh(x1_layout.mesh)
	222	+ output_layout.set_tensor_map((-1, n_placement))
	223	+ output_layout.tensor_map_to_placement()
	224	+
	225	+ # When k is sharded, output is a partial sum; mark Partial so the framework
	226	+ # can insert AllReduce downstream.
	227	+ if k_placement != -1:
	228	+ self._set_partial_from_k(output_layout, k_placement)
	229	+
	230	+ # gather_out: gather_output=True → m Replicate (-1), k follows x1's k placement.
	231	+ # gather_output=False → CANN returns a 1-D empty tensor; force all-Replicate so
	232	+ # the layout is compatible with any tensor rank returned by the kernel.
	233	+ gather_out_layout = Layout.from_device_mesh(x1_layout.mesh)
	234	+ gather_k = k_placement if gather_output else -1
	235	+ gather_out_layout.set_tensor_map((-1, gather_k))
	236	+ gather_out_layout.tensor_map_to_placement()
	237	+
	238	+ return (copy.deepcopy(output_layout), copy.deepcopy(gather_out_layout)), None

Ahyper_parallel/core/shard/ops/parallel_matmul_reduce_scatter.py+211-0

@@ -0,0 +1,211 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Distributed implementation for MatmulReduceScatter (MC2 fusion) operator."""
	16	+import copy
	17	+from typing import Tuple
	18	+
	19	+from hyper_parallel.core.dtensor.layout import Layout
	20	+from .parallel_ops import DistributedOp
	21	+
	22	+
	23	+def _normalize_mrs_args(
	24	+ x1,
	25	+ x2,
	26	+ group,
	27	+ world_size,
	28	+ reduce_op='sum',
	29	+ bias=None,
	30	+ comm_turn=0,
	31	+ trans_input=False,
	32	+ trans_x2=False,
	33	+) -> Tuple[tuple, dict]:
	34	+ """Normalize positional and keyword arguments into a canonical positional tuple.
	35	+
	36	+ Args:
	37	+ x1: Input tensor. Physical shape (m, k); trans_input=False only (MS constraint).
	38	+ x2: Weight tensor with shape (k, n) if trans_x2=False, or (n, k) if trans_x2=True.
	39	+ group: Communication group name string.
	40	+ world_size: Number of ranks in the communication group.
	41	+ reduce_op: Reduce operation string (only 'sum' is supported).
	42	+ bias: Must be None (MS constraint; bias is not supported).
	43	+ comm_turn: Communication turn (only 0 is supported).
	44	+ trans_input: Must be False (MS constraint).
	45	+ trans_x2: If True, x2 physical shape is (n, k); MindSpore transposes x2 before CANN call.
	46	+
	47	+ Returns:
	48	+ tuple: (positional_args_tuple, empty_kwargs_dict)
	49	+ """
	50	+ return (x1, x2, group, world_size, reduce_op, bias, comm_turn, trans_input, trans_x2), {}
	51	+
	52	+
	53	+class MatmulReduceScatterDistributedOp(DistributedOp):
	54	+ """Distributed operator for mindspore.ops.matmul_reduce_scatter (MC2 fusion).
	55	+
	56	+ The CANN MatmulReduceScatter kernel handles communication (ReduceScatter)
	57	+ internally. HyperParallel's role is to:
	58	+ 1. Extract local tensors from DTensor inputs via to_local().
	59	+ 2. Infer output DTensor layouts so downstream operators can correctly
	60	+ understand the distribution state of the output.
	61	+
	62	+ No Partial state is needed because CANN internally completes the ReduceScatter;
	63	+ each rank receives the correctly reduced and scattered result slice.
	64	+
	65	+ Shape transformation (logical, after any transpose):
	66	+ x1 (m, k_local), x2 (k_local, n) —[matmul]→ partial (m, n)
	67	+ —[CANN ReduceScatter on m]→ output (m_local = m / world_size, n)
	68	+
	69	+ Sharding constraints:
	70	+ - x1 physical (m, k): dim 1 (k) must be Shard (TP/comm axis);
	71	+ dim 0 (m) may be Replicate or Shard (DP axis).
	72	+ - x2's k-dim must be Shard and match x1's k-dim placement exactly.
	73	+ - trans_x2=False: x2 k-dim is dim 0; trans_x2=True: x2 k-dim is dim 1.
	74	+ - trans_input=False, bias=None only (current MS constraints).
	75	+ - Partial inputs are not allowed.
	76	+ """
	77	+
	78	+ def preprocess(self, args: tuple, kwargs: dict) -> tuple:
	79	+ """Extract local tensors and build the layout cache.
	80	+
	81	+ Args:
	82	+ args: Positional arguments (DTensors for x1 and x2).
	83	+ kwargs: Keyword arguments.
	84	+
	85	+ Returns:
	86	+ tuple: (local_args, local_kwargs, cache_values) where
	87	+ local_args contains extracted local tensors,
	88	+ local_kwargs contains keyword arguments,
	89	+ cache_values = [x1_layout, x2_layout, trans_x2].
	90	+ """
	91	+ norm_args, _ = _normalize_mrs_args(args, *kwargs)
	92	+ x1 = norm_args[0]
	93	+ x2 = norm_args[1]
	94	+ trans_x2 = norm_args[8]
	95	+
	96	+ # MindSpore matmul_reduce_scatter: only (input, x2, group, world_size) are
	97	+ # positional; reduce_op, bias, comm_turn, trans_input, trans_x2 are
	98	+ # keyword-only (after the '*' separator).
	99	+ local_args = (
	100	+ x1.to_local(),
	101	+ x2.to_local(),
	102	+ norm_args[2], # group
	103	+ norm_args[3], # world_size
	104	+ )
	105	+ local_kwargs = {
	106	+ 'reduce_op': norm_args[4],
	107	+ 'bias': norm_args[5],
	108	+ 'comm_turn': norm_args[6],
	109	+ 'trans_input': norm_args[7],
	110	+ 'trans_x2': trans_x2,
	111	+ }
	112	+
	113	+ cache_values = [
	114	+ x1.layout,
	115	+ x2.layout,
	116	+ trans_x2,
	117	+ ]
	118	+ return local_args, local_kwargs, cache_values
	119	+
	120	+ @staticmethod
	121	+ def _validate_input_layouts(
	122	+ x1_layout: Layout,
	123	+ x2_layout: Layout,
	124	+ trans_x2: bool,
	125	+ ) -> None:
	126	+ """Validate sharding constraints for MatmulReduceScatter inputs.
	127	+
	128	+ Args:
	129	+ x1_layout: Layout of x1. Physical (m, k); trans_input=False only.
	130	+ x2_layout: Layout of x2 (k, n) if trans_x2=False, or (n, k) if trans_x2=True.
	131	+ trans_x2: Whether x2 is transposed.
	132	+
	133	+ Raises:
	134	+ ValueError: If x1's k dimension is Replicate, x2's k dimension layout does not
	135	+ match x1's k dimension layout, or any input has Partial status.
	136	+ """
	137	+ op = "matmul_reduce_scatter"
	138	+ x1_tm = x1_layout.tensor_map
	139	+ x2_tm = x2_layout.tensor_map
	140	+
	141	+ # trans_input=False only: x1 physical (m, k) — k is dim 1.
	142	+ x1_k_dim = 1
	143	+
	144	+ if x1_tm[x1_k_dim] == -1:
	145	+ raise ValueError(
	146	+ f"For {op}, x1 k-dim (dim {x1_k_dim}) must be "
	147	+ f"Shard (not Replicate), because ReduceScatter requires k to be sharded. "
	148	+ f"Got tensor_map={x1_tm}"
	149	+ )
	150	+
	151	+ x2_k_dim = 1 if trans_x2 else 0
	152	+ if x2_tm[x2_k_dim] != x1_tm[x1_k_dim]:
	153	+ raise ValueError(
	154	+ f"For {op}, x2 dim {x2_k_dim} (k) layout must match x1 k-dim (dim {x1_k_dim}) "
	155	+ f"layout (trans_x2={trans_x2}), "
	156	+ f"but got x1_k={x1_tm[x1_k_dim]}, x2_k={x2_tm[x2_k_dim]}"
	157	+ )
	158	+
	159	+ def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
	160	+ """Infer output layout for MatmulReduceScatter.
	161	+
	162	+ ReduceScatter converts the k-dim sharding (TP comm axis) into m-dim sharding.
	163	+ trans_input=False only: x1 physical (m, k).
	164	+
	165	+ comm_mesh_dim = x1_tm[1] (k → TP / ReduceScatter scatter axis)
	166	+ dp_mesh_dim = x1_tm[0] (m → DP, or -1 if Replicate)
	167	+
	168	+ If dp_mesh_dim == -1 (pure TP):
	169	+ output_tm[0] = comm_mesh_dim
	170	+ Else (TP + DP joint sharding):
	171	+ output_tm[0] = (comm_mesh_dim, dp_mesh_dim)
	172	+
	173	+ output_tm[1] = x2's n dimension placement
	174	+ - trans_x2=False: x2 dim 1 (n) → x2_tm[1]
	175	+ - trans_x2=True: x2 dim 0 (n) → x2_tm[0]
	176	+
	177	+ Args:
	178	+ cache_values: [x1_layout, x2_layout, trans_x2]
	179	+
	180	+ Returns:
	181	+ tuple: (output_layout, None)
	182	+
	183	+ Raises:
	184	+ ValueError: If any input has Partial status or sharding constraints are violated.
	185	+ """
	186	+ x1_layout = cache_values[0]
	187	+ x2_layout = cache_values[1]
	188	+ trans_x2 = cache_values[2]
	189	+
	190	+ self._check_partial_inputs([x1_layout, x2_layout])
	191	+ self._validate_input_layouts(x1_layout, x2_layout, trans_x2)
	192	+
	193	+ x1_tm = x1_layout.tensor_map
	194	+ x2_tm = x2_layout.tensor_map
	195	+
	196	+ # trans_input=False: k is dim 1 (comm axis), m is dim 0 (DP axis or Replicate).
	197	+ comm_mesh_dim = x1_tm[1]
	198	+ dp_mesh_dim = x1_tm[0]
	199	+
	200	+ if dp_mesh_dim == -1:
	201	+ output_m = comm_mesh_dim
	202	+ else:
	203	+ output_m = (comm_mesh_dim, dp_mesh_dim)
	204	+
	205	+ n_placement = x2_tm[0] if trans_x2 else x2_tm[1]
	206	+
	207	+ output_layout = Layout.from_device_mesh(x1_layout.mesh)
	208	+ output_layout.set_tensor_map((output_m, n_placement))
	209	+ output_layout.tensor_map_to_placement()
	210	+
	211	+ return copy.deepcopy(output_layout), None

Ahyper_parallel/core/shard/ops/yaml/all_gather_matmul_ops.yaml+4-0

@@ -0,0 +1,4 @@
	1	+AllGatherMatmul:
	2	+ dist_op_name: _all_gather_matmul_dist_op
	3	+ distributed_op_class: AllGatherMatmulDistributedOp
	4	+ distributed_op_file: parallel_all_gather_matmul

Ahyper_parallel/core/shard/ops/yaml/matmul_reduce_scatter_ops.yaml+4-0

@@ -0,0 +1,4 @@
	1	+MatmulReduceScatter:
	2	+ dist_op_name: _matmul_reduce_scatter_dist_op
	3	+ distributed_op_class: MatmulReduceScatterDistributedOp
	4	+ distributed_op_file: parallel_matmul_reduce_scatter

Atests/mindspore/st/shard/ops/all_gather_matmul_shard_in_python.py+328-0

@@ -0,0 +1,328 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""MindSpore distributed ST for mindspore.ops.all_gather_matmul (MC2 fusion).
	16	+
	17	+Reference computation uses primitive ops:
	18	+ AllGather(x1_local) → gathered_x1 → matmul(gathered_x1, x2_local)
	19	+
	20	+MC2 operator handles communication internally; HyperParallel is responsible for:
	21	+ - Extracting local tensors from DTensor inputs (preprocess / to_local).
	22	+ - Inferring output DTensor layouts (infer_layout).
	23	+
	24	+Numerical comparison uses float16 with atol=1e-2 / rtol=1e-2.
	25	+
	26	+Note on dtype choice: AGM reference is computed with float16 primitives (AllGather + matmul),
	27	+so MC2 output and reference share the same numerical path — errors cancel and max_diff < 1e-3.
	28	+float16 (10-bit mantissa) is preferred over bfloat16 for higher precision.
	29	+"""
	30	+import numpy as np
	31	+import mindspore as ms
	32	+import mindspore.communication.management as D
	33	+import mindspore.ops as ms_ops
	34	+from mindspore import Tensor
	35	+from mindspore.ops.function.comm_func import all_gather_into_tensor
	36	+from hyper_parallel import init_device_mesh
	37	+from hyper_parallel.core.dtensor.dtensor import distribute_tensor
	38	+from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
	39	+
	40	+
	41	+def setup_module():
	42	+ """Initialize MindSpore context and communication."""
	43	+ ms.context.set_context(mode=ms.context.PYNATIVE_MODE, device_target="Ascend")
	44	+ D.init()
	45	+
	46	+
	47	+np.random.seed(42)
	48	+ms.set_seed(42)
	49	+ms.set_deterministic(True)
	50	+
	51	+# Global tensor dimensions — satisfy CANN constraints: K ∈ [256, 65535), M is multiple of world_size.
	52	+M = 128
	53	+K = 256
	54	+N = 512
	55	+
	56	+
	57	+def _fp16(arr: np.ndarray) -> Tensor:
	58	+ """Convert float32 numpy array to MindSpore float16 Tensor."""
	59	+ return Tensor(arr.astype(np.float32), ms.float16)
	60	+
	61	+
	62	+def _to_f32np(t: Tensor) -> np.ndarray:
	63	+ """Convert MindSpore tensor (any dtype) to float32 numpy array."""
	64	+ return t.astype(ms.float32).asnumpy()
	65	+
	66	+
	67	+def _make_inputs():
	68	+ """Generate global float32 numpy arrays for x1 and x2."""
	69	+ x1_np = np.random.randn(M, K).astype(np.float32)
	70	+ x2_np = np.random.randn(K, N).astype(np.float32)
	71	+ return x1_np, x2_np
	72	+
	73	+
	74	+def _allclose(a: np.ndarray, b: np.ndarray, tag: str):
	75	+ assert np.allclose(a, b, atol=1e-2, rtol=1e-2), (
	76	+ f"{tag}: max_diff={np.abs(a - b).max():.4f}, "
	77	+ f"actual shape={a.shape}, ref shape={b.shape}"
	78	+ )
	79	+
	80	+
	81	+def test_agm_tp_input_shard0_x2_replicate():
	82	+ """
	83	+ Feature: AllGatherMatmul TP — x1 Shard(0) on tp, x2 Replicate, gather_output=True.
	84	+ Description:
	85	+ - 2-rank tp mesh; x1 global (128, 256), each rank holds (64, 256).
	86	+ - x2 global (256, 512), replicated on each rank.
	87	+ - Reference: AllGather(x1_local) → (128, 256); matmul → (128, 512).
	88	+ - Verifies output shape, values, and gather_out shape/values.
	89	+ Expectation: Output matches reference within float16 tolerance.
	90	+ """
	91	+ x1_np, x2_np = _make_inputs()
	92	+
	93	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	94	+ group = mesh.get_group("tp")
	95	+ world_size = 2
	96	+
	97	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0),))
	98	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(),))
	99	+
	100	+ out_dtensor, gather_dtensor = ms_ops.all_gather_matmul(
	101	+ x1_dtensor, x2_dtensor, group, world_size, gather_output=True
	102	+ )
	103	+
	104	+ x1_local = x1_dtensor.to_local()
	105	+ gathered_x1, _ = all_gather_into_tensor(None, x1_local, group=group)
	106	+ ref_output = ms_ops.matmul(gathered_x1, x2_dtensor.to_local())
	107	+
	108	+ out_full = _to_f32np(out_dtensor.full_tensor())
	109	+ ref_np = _to_f32np(ref_output)
	110	+ _allclose(out_full, ref_np, "agm_tp_input_shard0_x2_replicate output")
	111	+
	112	+ gather_full = _to_f32np(gather_dtensor.full_tensor())
	113	+ gathered_ref = _to_f32np(gathered_x1)
	114	+ _allclose(gather_full, gathered_ref, "agm_tp_input_shard0_x2_replicate gather_out")
	115	+
	116	+
	117	+def test_agm_tp_input_shard0_x2_shard1():
	118	+ """
	119	+ Feature: AllGatherMatmul TP — x1 Shard(0) on tp, x2 Shard(1) on tp (n-dim sharded).
	120	+ Description:
	121	+ - 2-rank tp mesh; x1 global (128, 256), x2 global (256, 512).
	122	+ - x1 Shard(0): each rank holds (64, 256).
	123	+ - x2 Shard(1): each rank holds (256, 256) — n-dim halved.
	124	+ - Reference: AllGather(x1_local); matmul → (128, 256) per rank.
	125	+ - Full output reconstructed by gathering along n-dim.
	126	+ Expectation: Gathered output matches reference within float16 tolerance.
	127	+ """
	128	+ x1_np, x2_np = _make_inputs()
	129	+
	130	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	131	+ group = mesh.get_group("tp")
	132	+ world_size = 2
	133	+
	134	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0),))
	135	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Shard(1),))
	136	+
	137	+ out_dtensor, _ = ms_ops.all_gather_matmul(
	138	+ x1_dtensor, x2_dtensor, group, world_size, gather_output=False
	139	+ )
	140	+
	141	+ x1_local = x1_dtensor.to_local()
	142	+ gathered_x1, _ = all_gather_into_tensor(None, x1_local, group=group)
	143	+ ref_output = ms_ops.matmul(gathered_x1, _fp16(x2_np))
	144	+
	145	+ out_full = _to_f32np(out_dtensor.full_tensor())
	146	+ ref_np = _to_f32np(ref_output)
	147	+ _allclose(out_full, ref_np, "agm_tp_input_shard0_x2_shard1 output")
	148	+
	149	+
	150	+def test_agm_trans_x2_true():
	151	+ """
	152	+ Feature: AllGatherMatmul with trans_x2=True.
	153	+ Description:
	154	+ - 2-rank tp mesh; x1 Shard(0), x2 physical shape (N, K) with Shard(0) on n-dim.
	155	+ - Reference: AllGather(x1_local); matmul(gathered_x1, x2_local.T).
	156	+ - C++ impl transposes x2 before calling CANN.
	157	+ Expectation: Output matches reference within float16 tolerance.
	158	+ """
	159	+ x1_np = np.random.randn(M, K).astype(np.float32)
	160	+ x2_t_np = np.random.randn(N, K).astype(np.float32) # physical shape (N, K) when trans_x2=True
	161	+
	162	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	163	+ group = mesh.get_group("tp")
	164	+ world_size = 2
	165	+
	166	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0),))
	167	+ # x2 shape (N, K): Shard(0) shards n-dim on tp
	168	+ x2_dtensor = distribute_tensor(_fp16(x2_t_np), mesh, (Shard(0),))
	169	+
	170	+ out_dtensor, _ = ms_ops.all_gather_matmul(
	171	+ x1_dtensor, x2_dtensor, group, world_size, gather_output=False, trans_x2=True
	172	+ )
	173	+
	174	+ x1_local = x1_dtensor.to_local()
	175	+ gathered_x1, _ = all_gather_into_tensor(None, x1_local, group=group)
	176	+ # trans_x2=True: C++ transposes x2 before matmul; reference uses x2_t.T = (K, N) sliced portion
	177	+ x2_t_full = _fp16(x2_t_np)
	178	+ ref_output = ms_ops.matmul(gathered_x1, x2_t_full.T)
	179	+
	180	+ out_full = _to_f32np(out_dtensor.full_tensor())
	181	+ ref_np = _to_f32np(ref_output)
	182	+ _allclose(out_full, ref_np, "agm_trans_x2_true output")
	183	+
	184	+
	185	+def test_agm_gather_output_false():
	186	+ """
	187	+ Feature: AllGatherMatmul with gather_output=False.
	188	+ Description:
	189	+ - 2-rank tp mesh; x1 Shard(0), x2 Replicate.
	190	+ - gather_output=False: CANN returns empty tensor for gather_out.
	191	+ - Verifies output values and that gather_out is empty (shape[0]==0).
	192	+ Expectation: Output matches reference within float16 tolerance; gather_out tensor is empty.
	193	+ """
	194	+ x1_np, x2_np = _make_inputs()
	195	+
	196	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	197	+ group = mesh.get_group("tp")
	198	+ world_size = 2
	199	+
	200	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0),))
	201	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(),))
	202	+
	203	+ out_dtensor, gather_dtensor = ms_ops.all_gather_matmul(
	204	+ x1_dtensor, x2_dtensor, group, world_size, gather_output=False
	205	+ )
	206	+
	207	+ x1_local = x1_dtensor.to_local()
	208	+ gathered_x1, _ = all_gather_into_tensor(None, x1_local, group=group)
	209	+ ref_output = ms_ops.matmul(gathered_x1, x2_dtensor.to_local())
	210	+
	211	+ out_full = _to_f32np(out_dtensor.full_tensor())
	212	+ ref_np = _to_f32np(ref_output)
	213	+ _allclose(out_full, ref_np, "agm_gather_output_false output")
	214	+
	215	+ gather_local = gather_dtensor.to_local()
	216	+ assert gather_local.shape[0] == 0, (
	217	+ f"gather_output=False: expected gather_out shape[0]==0, "
	218	+ f"got shape={gather_local.shape}"
	219	+ )
	220	+
	221	+
	222	+def test_agm_dp_tp_input_shard0():
	223	+ """
	224	+ Feature: AllGatherMatmul on 2D (dp=2, tp=2) mesh.
	225	+ Description:
	226	+ - 4-rank mesh (dp=2, tp=2); x1 Shard(0) on tp only; x2 Replicate.
	227	+ - x1 global (128, 256); each rank holds (64, 256) — halved by tp.
	228	+ - Reference: AllGather on tp axis → (128, 256); matmul → (128, 512).
	229	+ Expectation: Output matches reference within float16 tolerance.
	230	+ """
	231	+ x1_np, x2_np = _make_inputs()
	232	+
	233	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
	234	+ tp_group = mesh.get_group("tp")
	235	+ world_size = 2
	236	+
	237	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Replicate(), Shard(0)))
	238	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(), Replicate()))
	239	+
	240	+ out_dtensor, _ = ms_ops.all_gather_matmul(
	241	+ x1_dtensor, x2_dtensor, tp_group, world_size, gather_output=False
	242	+ )
	243	+
	244	+ x1_local = x1_dtensor.to_local()
	245	+ gathered_x1, _ = all_gather_into_tensor(None, x1_local, group=tp_group)
	246	+ ref_output = ms_ops.matmul(gathered_x1, x2_dtensor.to_local())
	247	+
	248	+ out_full = _to_f32np(out_dtensor.full_tensor())
	249	+ ref_np = _to_f32np(ref_output)
	250	+ _allclose(out_full, ref_np, "agm_dp_tp_input_shard0 output")
	251	+
	252	+
	253	+def test_agm_k_sharded_mp2_tp2():
	254	+ """
	255	+ Feature: AllGatherMatmul with m-dim sharded on mp and k-dim sharded on tp.
	256	+ Description:
	257	+ - 4-rank (mp=2, tp=2) mesh; x1 (Shard(0), Shard(1)), x2 (Replicate(), Shard(0)).
	258	+ - K_SHARDED=512 so k_local=256 per tp rank, satisfying CANN k>=256 constraint.
	259	+ - AllGather on mp_group gathers m-dim; k sharded on tp produces Partial output on tp.
	260	+ - full_tensor() resolves Partial via AllReduce on tp to yield correct (M, N) result.
	261	+ Expectation: Output matches float16 matmul reference within tolerance.
	262	+ """
	263	+ k_sharded = 512
	264	+ x1_np = np.random.randn(M, k_sharded).astype(np.float32)
	265	+ x2_np = np.random.randn(k_sharded, N).astype(np.float32)
	266	+
	267	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("mp", "tp"))
	268	+ mp_group = mesh.get_group("mp")
	269	+ world_size = 2
	270	+
	271	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0), Shard(1)))
	272	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(), Shard(0)))
	273	+
	274	+ out_dtensor, _ = ms_ops.all_gather_matmul(
	275	+ x1_dtensor, x2_dtensor, mp_group, world_size, gather_output=False
	276	+ )
	277	+
	278	+ # k-sharded path: CANN splits k across tp ranks and AllReduces partial sums.
	279	+ # Accumulated float16 rounding differs from a single full-k matmul, so use
	280	+ # float32 numpy reference with the same loose tolerance as MRS k-sharded tests.
	281	+ ref_np = np.matmul(x1_np, x2_np)
	282	+ out_full = _to_f32np(out_dtensor.full_tensor())
	283	+ assert np.allclose(out_full, ref_np, atol=0.1, rtol=0.02), (
	284	+ f"agm_k_sharded_mp2_tp2 output: max_diff={np.abs(out_full - ref_np).max():.4f}, "
	285	+ f"actual shape={out_full.shape}, ref shape={ref_np.shape}"
	286	+ )
	287	+
	288	+
	289	+def test_agm_k_sharded_mp2_np2_tp2():
	290	+ """
	291	+ Feature: AllGatherMatmul with m, n, k each sharded on a distinct mesh axis.
	292	+ Description:
	293	+ - 8-rank (mp=2, np=2, tp=2) mesh.
	294	+ - x1 (M, K_SHARDED) with placements (Shard(0), Replicate(), Shard(1)):
	295	+ m sharded on mp, k sharded on tp; each rank holds (M/2, K/2).
	296	+ - x2 (K_SHARDED, N_SHARDED) with placements (Replicate(), Shard(1), Shard(0)):
	297	+ n sharded on np, k sharded on tp; each rank holds (K/2, N/2).
	298	+ - K_SHARDED=512, N_SHARDED=512 so k_local=256 satisfying CANN k>=256 constraint.
	299	+ - AllGather on mp_group gathers m-dim; k on tp → Partial on tp; n on np → Shard on n.
	300	+ - full_tensor() resolves Partial (AllReduce on tp) and Shard (AllGather on np).
	301	+ Expectation: Output matches float16 matmul reference within tolerance.
	302	+ """
	303	+ k_sharded = 512
	304	+ n_sharded = 512
	305	+ x1_np = np.random.randn(M, k_sharded).astype(np.float32)
	306	+ x2_np = np.random.randn(k_sharded, n_sharded).astype(np.float32)
	307	+
	308	+ mesh = init_device_mesh(
	309	+ device_type="npu", mesh_shape=(2, 2, 2), mesh_dim_names=("mp", "np", "tp")
	310	+ )
	311	+ mp_group = mesh.get_group("mp")
	312	+ world_size = 2
	313	+
	314	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0), Replicate(), Shard(1)))
	315	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(), Shard(1), Shard(0)))
	316	+
	317	+ out_dtensor, _ = ms_ops.all_gather_matmul(
	318	+ x1_dtensor, x2_dtensor, mp_group, world_size, gather_output=False
	319	+ )
	320	+
	321	+ # k-sharded path: same rounding argument as test_agm_k_sharded_mp2_tp2;
	322	+ # use float32 numpy reference with loose tolerance.
	323	+ ref_np = np.matmul(x1_np, x2_np)
	324	+ out_full = _to_f32np(out_dtensor.full_tensor())
	325	+ assert np.allclose(out_full, ref_np, atol=0.1, rtol=0.02), (
	326	+ f"agm_k_sharded_mp2_np2_tp2 output: max_diff={np.abs(out_full - ref_np).max():.4f}, "
	327	+ f"actual shape={out_full.shape}, ref shape={ref_np.shape}"
	328	+ )

Atests/mindspore/st/shard/ops/matmul_reduce_scatter_shard_in_python.py+225-0

@@ -0,0 +1,225 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""MindSpore distributed ST for mindspore.ops.matmul_reduce_scatter (MC2 fusion).
	16	+
	17	+Reference computation uses primitive ops:
	18	+ local matmul(x1_local, x2_local) → partial_out → ReduceScatter → output
	19	+
	20	+MC2 operator handles communication internally; HyperParallel is responsible for:
	21	+ - Extracting local tensors from DTensor inputs (preprocess / to_local).
	22	+ - Inferring output DTensor layouts (infer_layout).
	23	+
	24	+Numerical comparison: full_tensor() of the output DTensor vs np.matmul of global inputs,
	25	+using float16 (10-bit mantissa, ~8x more precise than bfloat16) with atol=0.1 / rtol=0.02.
	26	+
	27	+Note on dtype choice: MRS reference is float32 numpy (unlike AGM which uses bf16 primitives
	28	+as reference). With bfloat16 (7-bit mantissa + ReduceScatter cross-rank sum), max_diff was
	29	+~0.5, requiring atol=0.5. Switching to float16 reduces expected max_diff to ~0.065,
	30	+allowing the tighter atol=0.1.
	31	+"""
	32	+import numpy as np
	33	+import mindspore as ms
	34	+import mindspore.communication.management as D
	35	+import mindspore.ops as ms_ops
	36	+from mindspore import Tensor
	37	+from hyper_parallel import init_device_mesh
	38	+from hyper_parallel.core.dtensor.dtensor import distribute_tensor
	39	+from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
	40	+
	41	+
	42	+def setup_module():
	43	+ """Initialize MindSpore context and communication."""
	44	+ ms.context.set_context(mode=ms.context.PYNATIVE_MODE, device_target="Ascend")
	45	+ D.init()
	46	+
	47	+
	48	+np.random.seed(42)
	49	+ms.set_seed(42)
	50	+ms.set_deterministic(True)
	51	+
	52	+# Global tensor dimensions — satisfy CANN constraints for MatmulReduceScatter:
	53	+# K is sharded by world_size, so k_local = K / world_size must be ≥ 256.
	54	+# With world_size=2: K ≥ 512. M must be a multiple of world_size.
	55	+M = 128
	56	+K = 512
	57	+N = 512
	58	+
	59	+
	60	+def _fp16(arr: np.ndarray) -> Tensor:
	61	+ """Convert float32 numpy array to MindSpore float16 Tensor."""
	62	+ return Tensor(arr.astype(np.float32), ms.float16)
	63	+
	64	+
	65	+def _to_f32np(t: Tensor) -> np.ndarray:
	66	+ """Convert MindSpore tensor (any dtype) to float32 numpy array."""
	67	+ return t.astype(ms.float32).asnumpy()
	68	+
	69	+
	70	+def _make_inputs():
	71	+ """Generate global float32 numpy arrays for x1 and x2."""
	72	+ x1_np = np.random.randn(M, K).astype(np.float32)
	73	+ x2_np = np.random.randn(K, N).astype(np.float32)
	74	+ return x1_np, x2_np
	75	+
	76	+
	77	+def _allclose(a: np.ndarray, b: np.ndarray, tag: str):
	78	+ # float16 has 10-bit mantissa (~8x more precise than bfloat16).
	79	+ # With k_local=256 products per rank + ReduceScatter sum, expected max_diff ~0.065
	80	+ # for output values ~20-40 → atol=0.1, rtol=0.02.
	81	+ assert np.allclose(a, b, atol=0.1, rtol=0.02), (
	82	+ f"{tag}: max_diff={np.abs(a - b).max():.4f}, "
	83	+ f"actual shape={a.shape}, ref shape={b.shape}"
	84	+ )
	85	+
	86	+
	87	+def test_mrs_tp_basic():
	88	+ """
	89	+ Feature: MatmulReduceScatter TP — x1 Shard(1) on tp (k-dim), x2 Shard(0) on tp (k-dim).
	90	+ Description:
	91	+ - 2-rank tp mesh; x1 global (128, 256), x2 global (256, 512).
	92	+ - x1 Shard(1): each rank holds (128, 128) — k-dim halved.
	93	+ - x2 Shard(0): each rank holds (128, 512) — k-dim halved.
	94	+ - Reference: full matmul(x1_global, x2_global) → (128, 512); each rank gets M/2 rows.
	95	+ - Verifies output full_tensor matches global reference within float16 tolerance.
	96	+ Expectation: Output matches reference within float16 tolerance.
	97	+ """
	98	+ x1_np, x2_np = _make_inputs()
	99	+
	100	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	101	+ group = mesh.get_group("tp")
	102	+ world_size = 2
	103	+
	104	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(1),))
	105	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Shard(0),))
	106	+
	107	+ out_dtensor = ms_ops.matmul_reduce_scatter(x1_dtensor, x2_dtensor, group, world_size)
	108	+
	109	+ ref_np = np.matmul(x1_np, x2_np)
	110	+ out_full = _to_f32np(out_dtensor.full_tensor())
	111	+ _allclose(out_full, ref_np, "mrs_tp_basic output")
	112	+
	113	+
	114	+def test_mrs_trans_x2_true():
	115	+ """
	116	+ Feature: MatmulReduceScatter with trans_x2=True.
	117	+ Description:
	118	+ - 2-rank tp mesh; x1 Shard(1) on k, x2 physical shape (N, K) with Shard(1) on k-dim.
	119	+ - Reference: full matmul(x1_global, x2_t_global.T) — x2_t physical (512, 256), transposed is (256, 512).
	120	+ - C++ impl transposes x2 before calling CANN.
	121	+ Expectation: Output matches reference within float16 tolerance.
	122	+ """
	123	+ x1_np = np.random.randn(M, K).astype(np.float32)
	124	+ x2_t_np = np.random.randn(N, K).astype(np.float32) # physical shape (N, K) when trans_x2=True
	125	+
	126	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	127	+ group = mesh.get_group("tp")
	128	+ world_size = 2
	129	+
	130	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(1),))
	131	+ # x2 physical (N, K): Shard(1) shards k-dim on tp
	132	+ x2_dtensor = distribute_tensor(_fp16(x2_t_np), mesh, (Shard(1),))
	133	+
	134	+ out_dtensor = ms_ops.matmul_reduce_scatter(
	135	+ x1_dtensor, x2_dtensor, group, world_size, trans_x2=True
	136	+ )
	137	+
	138	+ # trans_x2=True: C++ transposes x2 before matmul; reference is matmul(x1, x2_t.T)
	139	+ ref_np = np.matmul(x1_np, x2_t_np.T) # x2_t.T shape: (K, N) → output (M, N)
	140	+ out_full = _to_f32np(out_dtensor.full_tensor())
	141	+ _allclose(out_full, ref_np, "mrs_trans_x2_true output")
	142	+
	143	+
	144	+def test_mrs_large_m():
	145	+ """
	146	+ Feature: MatmulReduceScatter with larger M dimension.
	147	+ Description:
	148	+ - 2-rank tp mesh; x1 global (256, 512) Shard(1) on k, x2 global (512, 512) Shard(0) on k.
	149	+ - Verifies output correctness with M=256 (larger than test_mrs_tp_basic which uses M=128).
	150	+ Expectation: Output matches reference within float16 tolerance.
	151	+ """
	152	+ x1_np = np.random.randn(256, K).astype(np.float32)
	153	+ x2_np = np.random.randn(K, N).astype(np.float32)
	154	+
	155	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	156	+ group = mesh.get_group("tp")
	157	+ world_size = 2
	158	+
	159	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(1),))
	160	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Shard(0),))
	161	+
	162	+ out_dtensor = ms_ops.matmul_reduce_scatter(x1_dtensor, x2_dtensor, group, world_size)
	163	+
	164	+ ref_np = np.matmul(x1_np, x2_np)
	165	+ out_full = _to_f32np(out_dtensor.full_tensor())
	166	+ _allclose(out_full, ref_np, "mrs_large_m output")
	167	+
	168	+
	169	+def test_mrs_dp_tp_basic():
	170	+ """
	171	+ Feature: MatmulReduceScatter on 2D (dp=2, tp=2) mesh.
	172	+ Description:
	173	+ - 4-rank mesh (dp=2, tp=2); x1 Shard(1) on tp only; x2 Shard(0) on tp only.
	174	+ - x1 global (128, 256); each rank holds (128, 128) — halved by tp.
	175	+ - Reference: full matmul(x1_global, x2_global) → (128, 512); each tp-rank gets M/2 rows.
	176	+ - Verifies forward output correctness on 2D mesh.
	177	+ Expectation: Output matches reference within float16 tolerance.
	178	+ """
	179	+ x1_np, x2_np = _make_inputs()
	180	+
	181	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
	182	+ tp_group = mesh.get_group("tp")
	183	+ world_size = 2
	184	+
	185	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Replicate(), Shard(1)))
	186	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(), Shard(0)))
	187	+
	188	+ out_dtensor = ms_ops.matmul_reduce_scatter(x1_dtensor, x2_dtensor, tp_group, world_size)
	189	+
	190	+ ref_np = np.matmul(x1_np, x2_np)
	191	+ out_full = _to_f32np(out_dtensor.full_tensor())
	192	+ _allclose(out_full, ref_np, "mrs_dp_tp_basic output")
	193	+
	194	+
	195	+def test_mrs_mp2_np2_tp2():
	196	+ """
	197	+ Feature: MatmulReduceScatter with m, n, k each sharded on a distinct mesh axis.
	198	+ Description:
	199	+ - 8-rank (mp=2, np=2, tp=2) mesh.
	200	+ - x1 (M, K) with placements (Shard(0), Replicate(), Shard(1)):
	201	+ m sharded on mp, k sharded on tp; each rank holds (M/2, K/2).
	202	+ - x2 (K, N) with placements (Replicate(), Shard(1), Shard(0)):
	203	+ n sharded on np, k sharded on tp; each rank holds (K/2, N/2).
	204	+ - ReduceScatter on tp_group (world_size=2): sums k partial and scatters m;
	205	+ output m is jointly sharded on (tp, mp), n sharded on np;
	206	+ each rank holds (M/4, N/2).
	207	+ - full_tensor() reconstructs (M, N) via AllGather on both m and n dimensions.
	208	+ Expectation: Output matches reference within float16 tolerance.
	209	+ """
	210	+ x1_np, x2_np = _make_inputs()
	211	+
	212	+ mesh = init_device_mesh(
	213	+ device_type="npu", mesh_shape=(2, 2, 2), mesh_dim_names=("mp", "np", "tp")
	214	+ )
	215	+ tp_group = mesh.get_group("tp")
	216	+ world_size = 2
	217	+
	218	+ x1_dtensor = distribute_tensor(_fp16(x1_np), mesh, (Shard(0), Replicate(), Shard(1)))
	219	+ x2_dtensor = distribute_tensor(_fp16(x2_np), mesh, (Replicate(), Shard(1), Shard(0)))
	220	+
	221	+ out_dtensor = ms_ops.matmul_reduce_scatter(x1_dtensor, x2_dtensor, tp_group, world_size)
	222	+
	223	+ ref_np = np.matmul(x1_np, x2_np)
	224	+ out_full = _to_f32np(out_dtensor.full_tensor())
	225	+ _allclose(out_full, ref_np, "mrs_mp2_np2_tp2 output")

Atests/mindspore/st/shard/ops/test_all_gather_matmul_shard_in_python.py+70-0

@@ -0,0 +1,70 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Test runner for all_gather_matmul distributed ST (MindSpore)."""
	16	+from pathlib import Path
	17	+
	18	+from tests.common.mark_utils import arg_mark
	19	+from tests.common.parallel_case import parallel_run, MindSporeCase
	20	+
	21	+IMPL_FILE = str(
	22	+ Path(__file__).resolve().parent / "all_gather_matmul_shard_in_python.py"
	23	+)
	24	+
	25	+
	26	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	27	+def test_agm_ms_group1():
	28	+ """
	29	+ Feature: parallel run cases in all_gather_matmul_shard_in_python (2-card cases, 8 cards total)
	30	+ Description:
	31	+ 1. test_agm_tp_input_shard0_x2_replicate — x1 Shard(0) on tp, x2 Replicate, gather_output=True
	32	+ 2. test_agm_tp_input_shard0_x2_shard1 — x1 Shard(0), x2 Shard(1) on n-dim
	33	+ 3. test_agm_trans_x2_true — trans_x2=True, x2 physical (N,K) Shard(0) on n
	34	+ 4. test_agm_gather_output_false — gather_output=False; verify gather_out is empty
	35	+ Expectation: Run success.
	36	+ """
	37	+ parallel_run([
	38	+ MindSporeCase(IMPL_FILE, "test_agm_tp_input_shard0_x2_replicate", 20100, 2, 2, 2),
	39	+ MindSporeCase(IMPL_FILE, "test_agm_tp_input_shard0_x2_shard1", 20101, 2, 2, 2),
	40	+ MindSporeCase(IMPL_FILE, "test_agm_trans_x2_true", 20102, 2, 2, 2),
	41	+ MindSporeCase(IMPL_FILE, "test_agm_gather_output_false", 20103, 2, 2, 2),
	42	+ ])
	43	+
	44	+
	45	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	46	+def test_agm_ms_group2():
	47	+ """
	48	+ Feature: parallel run cases in all_gather_matmul_shard_in_python (4-card cases)
	49	+ Description:
	50	+ 1. test_agm_dp_tp_input_shard0 — 2D (dp=2, tp=2) mesh, x1 Shard(0) on tp; fwd+grad
	51	+ 2. test_agm_k_sharded_mp2_tp2 — 2D (mp=2, tp=2) mesh, m on mp + k on tp; Partial output
	52	+ Expectation: Run success.
	53	+ """
	54	+ parallel_run([
	55	+ MindSporeCase(IMPL_FILE, "test_agm_dp_tp_input_shard0", 20104, 4, 4, 2),
	56	+ MindSporeCase(IMPL_FILE, "test_agm_k_sharded_mp2_tp2", 20105, 4, 4, 2),
	57	+ ])
	58	+
	59	+
	60	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	61	+def test_agm_ms_group3():
	62	+ """
	63	+ Feature: parallel run cases in all_gather_matmul_shard_in_python (8-card case)
	64	+ Description:
	65	+ 1. test_agm_k_sharded_mp2_np2_tp2 — 3D (mp=2, np=2, tp=2) mesh, m/n/k sharded on distinct axes
	66	+ Expectation: Run success.
	67	+ """
	68	+ parallel_run([
	69	+ MindSporeCase(IMPL_FILE, "test_agm_k_sharded_mp2_np2_tp2", 20106, 8, 8, 2),
	70	+ ])

Atests/mindspore/st/shard/ops/test_matmul_reduce_scatter_shard_in_python.py+66-0

@@ -0,0 +1,66 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Test runner for matmul_reduce_scatter distributed ST (MindSpore)."""
	16	+from pathlib import Path
	17	+
	18	+from tests.common.mark_utils import arg_mark
	19	+from tests.common.parallel_case import parallel_run, MindSporeCase
	20	+
	21	+IMPL_FILE = str(
	22	+ Path(__file__).resolve().parent / "matmul_reduce_scatter_shard_in_python.py"
	23	+)
	24	+
	25	+
	26	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	27	+def test_mrs_ms_group1():
	28	+ """
	29	+ Feature: parallel run cases in matmul_reduce_scatter_shard_in_python (2-card cases, 6 cards total)
	30	+ Description:
	31	+ 1. test_mrs_tp_basic — x1 Shard(1) on k, x2 Shard(0) on k; reference = matmul + ReduceScatter
	32	+ 2. test_mrs_trans_x2_true — trans_x2=True, x2 physical (N,K) Shard(1) on k
	33	+ 3. test_mrs_large_m — larger M dimension (M=256)
	34	+ Expectation: Run success.
	35	+ """
	36	+ parallel_run([
	37	+ MindSporeCase(IMPL_FILE, "test_mrs_tp_basic", 20200, 2, 2, 2),
	38	+ MindSporeCase(IMPL_FILE, "test_mrs_trans_x2_true", 20201, 2, 2, 2),
	39	+ MindSporeCase(IMPL_FILE, "test_mrs_large_m", 20202, 2, 2, 2),
	40	+ ])
	41	+
	42	+
	43	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	44	+def test_mrs_ms_group2():
	45	+ """
	46	+ Feature: parallel run cases in matmul_reduce_scatter_shard_in_python (4-card case)
	47	+ Description:
	48	+ 1. test_mrs_dp_tp_basic — 2D (dp=2, tp=2) mesh, x1 Shard(1) on tp, x2 Shard(0) on tp
	49	+ Expectation: Run success.
	50	+ """
	51	+ parallel_run([
	52	+ MindSporeCase(IMPL_FILE, "test_mrs_dp_tp_basic", 20203, 4, 4, 2),
	53	+ ])
	54	+
	55	+
	56	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	57	+def test_mrs_ms_group3():
	58	+ """
	59	+ Feature: parallel run cases in matmul_reduce_scatter_shard_in_python (8-card case)
	60	+ Description:
	61	+ 1. test_mrs_mp2_np2_tp2 — 3D (mp=2, np=2, tp=2) mesh, m/n/k sharded on distinct axes
	62	+ Expectation: Run success.
	63	+ """
	64	+ parallel_run([
	65	+ MindSporeCase(IMPL_FILE, "test_mrs_mp2_np2_tp2", 20204, 8, 8, 2),
	66	+ ])

Mtests/mindspore/st/shard/ops/test_npu_mhc_post_custom_op.py+0-3

	@@ -23,9 +23,6 @@ IMPL_FILE = str(
23	)	23	)
24		24
25		25
26	-@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
27	-
28	-
29	@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")	26	@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
30	def test_mhc_post_group1():	27	def test_mhc_post_group1():
31	"""	28	"""

Mtests/mindspore/st/shard/ops/test_npu_mhc_pre_sinkhorn_custom_op.py+0-3

	@@ -23,9 +23,6 @@ IMPL_FILE = str(
23	)	23	)
24		24
25		25
26	-@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
27	-
28	-
29	@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")	26	@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
30	def test_mhc_pre_sinkhorn_group1():	27	def test_mhc_pre_sinkhorn_group1():
31	"""	28	"""

Atests/ut/core/shard/ops/test_parallel_all_gather_matmul.py+359-0

@@ -0,0 +1,359 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Unit tests for AllGatherMatmulDistributedOp."""
	16	+import os
	17	+import unittest
	18	+from unittest.mock import MagicMock, patch
	19	+
	20	+import numpy as np
	21	+
	22	+os.environ["HYPER_PARALLEL_PLATFORM"] = "torch"
	23	+
	24	+from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE
	25	+from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
	26	+from hyper_parallel.core.shard.ops.parallel_all_gather_matmul import (
	27	+ AllGatherMatmulDistributedOp,
	28	+ _normalize_agm_args,
	29	+)
	30	+from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
	31	+from hyper_parallel.core.dtensor.device_mesh import init_device_mesh, _DEVICE_MESH_MAP
	32	+from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS
	33	+
	34	+
	35	+class TestAllGatherMatmulDistributedOp(unittest.TestCase):
	36	+ """Unit tests for AllGatherMatmulDistributedOp."""
	37	+
	38	+ def setUp(self):
	39	+ """Clear global state before each test."""
	40	+ EXISTING_COMM_GROUPS.clear()
	41	+ _DEVICE_MESH_MAP.clear()
	42	+ _LAYOUT_CACHE.clear()
	43	+
	44	+ def tearDown(self):
	45	+ """Clear global state after each test."""
	46	+ EXISTING_COMM_GROUPS.clear()
	47	+ _DEVICE_MESH_MAP.clear()
	48	+ _LAYOUT_CACHE.clear()
	49	+
	50	+ def _setup_mock_platform(self, mock_platform, world_size=8):
	51	+ """Configure mock platform for mesh creation."""
	52	+ mock_platform.get_rank.return_value = 0
	53	+ mock_platform.get_world_size.return_value = world_size
	54	+ mock_platform.split_group.return_value = MagicMock()
	55	+ mock_platform.tensor_to_numpy.side_effect = (
	56	+ lambda t: t.asnumpy() if hasattr(t, "asnumpy") else np.array(t)
	57	+ )
	58	+
	59	+ def _make_1d_mesh(self, mock_platform, size=4, name="tp"):
	60	+ """Return a 1-D mesh of given size."""
	61	+ self._setup_mock_platform(mock_platform, world_size=size)
	62	+ return init_device_mesh(device_type="npu", mesh_shape=(size,), mesh_dim_names=(name,),
	63	+ init_backend=False)
	64	+
	65	+ def _make_2d_mesh(self, mock_platform, shape=(2, 4), names=("dp", "tp")):
	66	+ """Return a 2-D mesh."""
	67	+ self._setup_mock_platform(mock_platform, world_size=shape[0] * shape[1])
	68	+ return init_device_mesh(device_type="npu", mesh_shape=shape, mesh_dim_names=names,
	69	+ init_backend=False)
	70	+
	71	+ @staticmethod
	72	+ def _get_op():
	73	+ return get_distributed_op("AllGatherMatmul")
	74	+
	75	+ # ------------------------------------------------------------------
	76	+ # _normalize_agm_args
	77	+ # ------------------------------------------------------------------
	78	+
	79	+ def test_normalize_args_defaults_1(self):
	80	+ """
	81	+ Feature: _normalize_agm_args fills in default values.
	82	+ Description: Call with only 4 mandatory positional args.
	83	+ Expectation: Optional args get correct defaults.
	84	+ """
	85	+ x1, x2, group, world_size = object(), object(), "hccl_world_group", 2
	86	+ args, kwargs = _normalize_agm_args(x1, x2, group, world_size)
	87	+ self.assertIs(args[0], x1)
	88	+ self.assertIs(args[1], x2)
	89	+ self.assertEqual(args[2], "hccl_world_group")
	90	+ self.assertEqual(args[3], 2)
	91	+ self.assertIsNone(args[4], msg=f"bias default should be None, got {args[4]}")
	92	+ self.assertEqual(args[5], 0, msg=f"gather_index default should be 0, got {args[5]}")
	93	+ self.assertTrue(args[6], msg=f"gather_output default should be True, got {args[6]}")
	94	+ self.assertEqual(args[7], 0, msg=f"comm_turn default should be 0, got {args[7]}")
	95	+ self.assertFalse(args[8], msg=f"trans_input default should be False, got {args[8]}")
	96	+ self.assertFalse(args[9], msg=f"trans_x2 default should be False, got {args[9]}")
	97	+ self.assertEqual(kwargs, {}, msg=f"kwargs should be empty, got {kwargs}")
	98	+
	99	+ # ------------------------------------------------------------------
	100	+ # YAML registration
	101	+ # ------------------------------------------------------------------
	102	+
	103	+ def test_yaml_registration_mindspore_2(self):
	104	+ """
	105	+ Feature: YAML loader registers AllGatherMatmulDistributedOp.
	106	+ Description: Call get_distributed_op with MindSpore CamelCase op name.
	107	+ Expectation: Returns an AllGatherMatmulDistributedOp instance.
	108	+ """
	109	+ op = self._get_op()
	110	+ self.assertIsNotNone(op, msg="AllGatherMatmul should be registered")
	111	+ self.assertIsInstance(
	112	+ op, AllGatherMatmulDistributedOp,
	113	+ msg=f"Expected AllGatherMatmulDistributedOp, got {type(op)}"
	114	+ )
	115	+
	116	+ # ------------------------------------------------------------------
	117	+ # infer_layout — positive cases
	118	+ # cache_values format: [x1_layout, x2_layout, trans_x2, gather_output]
	119	+ # ------------------------------------------------------------------
	120	+
	121	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	122	+ def test_agm_replicated_3(self, mock_platform):
	123	+ """
	124	+ Feature: infer_layout all replicated.
	125	+ Description: Both inputs fully replicated.
	126	+ Expectation: output and gather_out tensor_maps are both (-1, -1).
	127	+ """
	128	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	129	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	130	+ x2_layout = _build_layout(mesh, (Replicate(),), 2)
	131	+
	132	+ op = self._get_op()
	133	+ (out_layout, gather_layout), extra = op.infer_layout(
	134	+ [x1_layout, x2_layout, False, True])
	135	+
	136	+ self.assertIsNone(extra)
	137	+ self.assertEqual(out_layout.tensor_map, (-1, -1),
	138	+ msg=f"Replicated output: expected (-1,-1), got {out_layout.tensor_map}")
	139	+ self.assertEqual(gather_layout.tensor_map, (-1, -1),
	140	+ msg=f"Replicated gather_out: expected (-1,-1), got {gather_layout.tensor_map}")
	141	+
	142	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	143	+ def test_agm_input_shard0_x2_replicate_4(self, mock_platform):
	144	+ """
	145	+ Feature: infer_layout x1 Shard(0) on tp, x2 Replicate.
	146	+ Description: x1 m-dim sharded by AllGather; x2 n-dim Replicate.
	147	+ Expectation: output tensor_map (-1, -1); gather_out tensor_map (-1, -1).
	148	+ """
	149	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	150	+ x1_layout = _build_layout(mesh, (Shard(0),), 2)
	151	+ x2_layout = _build_layout(mesh, (Replicate(),), 2)
	152	+
	153	+ op = self._get_op()
	154	+ (out_layout, gather_layout), extra = op.infer_layout(
	155	+ [x1_layout, x2_layout, False, True])
	156	+
	157	+ self.assertIsNone(extra)
	158	+ self.assertEqual(out_layout.tensor_map, (-1, -1),
	159	+ msg=f"Expected (-1,-1), got {out_layout.tensor_map}")
	160	+ self.assertEqual(gather_layout.tensor_map, (-1, -1),
	161	+ msg=f"Expected (-1,-1), got {gather_layout.tensor_map}")
	162	+
	163	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	164	+ def test_agm_input_shard0_x2_shard1_5(self, mock_platform):
	165	+ """
	166	+ Feature: infer_layout x1 Shard(0) on tp, x2 Shard(1) on n.
	167	+ Description: x2 n-dim sharded on tp mesh_dim (index 0).
	168	+ Expectation: output tensor_map (-1, 0); gather_out tensor_map (-1, -1).
	169	+ """
	170	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	171	+ x1_layout = _build_layout(mesh, (Shard(0),), 2)
	172	+ x2_layout = _build_layout(mesh, (Shard(1),), 2)
	173	+
	174	+ op = self._get_op()
	175	+ (out_layout, gather_layout), extra = op.infer_layout(
	176	+ [x1_layout, x2_layout, False, True])
	177	+
	178	+ self.assertIsNone(extra)
	179	+ self.assertEqual(out_layout.tensor_map, (-1, 0),
	180	+ msg=f"Expected (-1, 0), got {out_layout.tensor_map}")
	181	+ self.assertEqual(gather_layout.tensor_map, (-1, -1),
	182	+ msg=f"Expected (-1,-1), got {gather_layout.tensor_map}")
	183	+
	184	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	185	+ def test_agm_trans_x2_false_n_sharded_6(self, mock_platform):
	186	+ """
	187	+ Feature: infer_layout trans_x2=False, x2 Shard(1).
	188	+ Description: trans_x2=False: n is x2 dim 1; x2 dim 1 sharded on tp (index 0).
	189	+ Expectation: output tensor_map (-1, 0).
	190	+ """
	191	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	192	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	193	+ x2_layout = _build_layout(mesh, (Shard(1),), 2)
	194	+
	195	+ op = self._get_op()
	196	+ (out_layout, _), extra = op.infer_layout([x1_layout, x2_layout, False, True])
	197	+
	198	+ self.assertIsNone(extra)
	199	+ self.assertEqual(out_layout.tensor_map, (-1, 0),
	200	+ msg=f"trans_x2=False n sharded: expected (-1, 0), got {out_layout.tensor_map}")
	201	+
	202	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	203	+ def test_agm_trans_x2_true_n_sharded_7(self, mock_platform):
	204	+ """
	205	+ Feature: infer_layout trans_x2=True, x2 Shard(0) on n.
	206	+ Description: trans_x2=True: n is x2 dim 0; x2 physical shape is (n, k);
	207	+ x2 Shard(0) means n is sharded on tp mesh_dim (index 0).
	208	+ Expectation: output tensor_map (-1, 0).
	209	+ """
	210	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	211	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	212	+ # trans_x2=True: x2 shape is (n, k); Shard(0) shards n on tp (mesh_dim index 0)
	213	+ x2_layout = _build_layout(mesh, (Shard(0),), 2)
	214	+
	215	+ op = self._get_op()
	216	+ (out_layout, _), extra = op.infer_layout([x1_layout, x2_layout, True, True])
	217	+
	218	+ self.assertIsNone(extra)
	219	+ self.assertEqual(out_layout.tensor_map, (-1, 0),
	220	+ msg=f"trans_x2=True n sharded: expected (-1, 0), got {out_layout.tensor_map}")
	221	+
	222	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	223	+ def test_agm_gather_output_false_8(self, mock_platform):
	224	+ """
	225	+ Feature: infer_layout gather_output=False.
	226	+ Description: When gather_output is False, gather_out layout should still be (-1, -1).
	227	+ Expectation: gather_out tensor_map is (-1, -1).
	228	+ """
	229	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	230	+ x1_layout = _build_layout(mesh, (Shard(0),), 2)
	231	+ x2_layout = _build_layout(mesh, (Replicate(),), 2)
	232	+
	233	+ op = self._get_op()
	234	+ (_, gather_layout), extra = op.infer_layout([x1_layout, x2_layout, False, False])
	235	+
	236	+ self.assertIsNone(extra)
	237	+ self.assertEqual(gather_layout.tensor_map, (-1, -1),
	238	+ msg=f"gather_output=False: expected (-1,-1), got {gather_layout.tensor_map}")
	239	+
	240	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	241	+ def test_agm_2d_mesh_dp_tp_9(self, mock_platform):
	242	+ """
	243	+ Feature: infer_layout 2D (dp=2, tp=4) mesh.
	244	+ Description: x1 Shard(0) on tp only; x2 Replicate. AllGather on tp axis.
	245	+ Expectation: output tensor_map (-1, -1); gather_out tensor_map (-1, -1).
	246	+ """
	247	+ mesh = self._make_2d_mesh(mock_platform, shape=(2, 4), names=("dp", "tp"))
	248	+ # x1: m-dim sharded on tp (mesh_dim index 1), k-dim Replicate
	249	+ x1_layout = _build_layout(mesh, (Replicate(), Shard(0)), 2)
	250	+ x2_layout = _build_layout(mesh, (Replicate(), Replicate()), 2)
	251	+
	252	+ op = self._get_op()
	253	+ (out_layout, gather_layout), extra = op.infer_layout(
	254	+ [x1_layout, x2_layout, False, True])
	255	+
	256	+ self.assertIsNone(extra)
	257	+ self.assertEqual(out_layout.tensor_map, (-1, -1),
	258	+ msg=f"2D mesh output: expected (-1,-1), got {out_layout.tensor_map}")
	259	+ self.assertEqual(gather_layout.tensor_map, (-1, -1),
	260	+ msg=f"2D mesh gather_out: expected (-1,-1), got {gather_layout.tensor_map}")
	261	+
	262	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	263	+ def test_output_independent_copies_10(self, mock_platform):
	264	+ """
	265	+ Feature: infer_layout returns independent deep copies for output and gather_out.
	266	+ Description: Mutate output_layout.tensor_map; gather_out_layout must be unaffected.
	267	+ Expectation: The two layouts are separate objects with independent state.
	268	+ """
	269	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	270	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	271	+ x2_layout = _build_layout(mesh, (Shard(1),), 2)
	272	+
	273	+ op = self._get_op()
	274	+ (out_layout, gather_layout), _ = op.infer_layout(
	275	+ [x1_layout, x2_layout, False, True])
	276	+
	277	+ self.assertIsNot(out_layout, gather_layout,
	278	+ msg="output_layout and gather_out_layout must be distinct objects")
	279	+ self.assertEqual(out_layout.tensor_map, (-1, 0))
	280	+ self.assertEqual(gather_layout.tensor_map, (-1, -1))
	281	+
	282	+ # ------------------------------------------------------------------
	283	+ # infer_layout — error cases
	284	+ # ------------------------------------------------------------------
	285	+
	286	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	287	+ def test_partial_input_raises_11(self, mock_platform):
	288	+ """
	289	+ Feature: infer_layout rejects Partial inputs.
	290	+ Description: x1 layout has Partial status.
	291	+ Expectation: ValueError with 'Partial status' message.
	292	+ """
	293	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	294	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	295	+ x1_layout.set_partial_by_dev_axis("tp", 'sum')
	296	+ x2_layout = _build_layout(mesh, (Replicate(),), 2)
	297	+
	298	+ op = self._get_op()
	299	+ with self.assertRaisesRegex(ValueError, "Partial status"):
	300	+ op.infer_layout([x1_layout, x2_layout, False, True])
	301	+
	302	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	303	+ def test_agm_k_sharded_output_partial_12(self, mock_platform):
	304	+ """
	305	+ Feature: infer_layout k-dim sharded on both x1 and x2 — output is Partial.
	306	+ Description: x1 Shard(1) on tp (k-dim), x2 Shard(0) on tp (k-dim); placements match.
	307	+ Output carries Partial(sum) on tp; gather_out k-dim follows x1's k placement.
	308	+ Expectation: output is_partial() True, tensor_map (-1,-1); gather_out tensor_map (-1, 0).
	309	+ """
	310	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	311	+ # x1 Shard(1): tensor_dim 1 (k) sharded on tp → tensor_map = (-1, 0)
	312	+ x1_layout = _build_layout(mesh, (Shard(1),), 2)
	313	+ # x2 Shard(0): tensor_dim 0 (k, trans_x2=False) sharded on tp → tensor_map = (0, -1)
	314	+ x2_layout = _build_layout(mesh, (Shard(0),), 2)
	315	+
	316	+ op = self._get_op()
	317	+ (out_layout, gather_layout), extra = op.infer_layout(
	318	+ [x1_layout, x2_layout, False, True])
	319	+
	320	+ self.assertIsNone(extra)
	321	+ self.assertTrue(out_layout.is_partial(),
	322	+ msg="k sharded: output must carry Partial status")
	323	+ self.assertEqual(out_layout.tensor_map, (-1, -1),
	324	+ msg=f"k sharded: expected output tensor_map (-1,-1), got {out_layout.tensor_map}")
	325	+ self.assertEqual(gather_layout.tensor_map, (-1, 0),
	326	+ msg=f"k sharded: expected gather_out tensor_map (-1,0), got {gather_layout.tensor_map}")
	327	+
	328	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	329	+ def test_agm_k_mismatch_raises_13(self, mock_platform):
	330	+ """
	331	+ Feature: infer_layout rejects mismatched k-dim placements between x1 and x2.
	332	+ Description: x1 k-dim is Replicate but x2 k-dim is Shard on tp.
	333	+ Expectation: ValueError mentioning k-dim placement must match.
	334	+ """
	335	+ mesh = self._make_1d_mesh(mock_platform, size=4)
	336	+ x1_layout = _build_layout(mesh, (Replicate(),), 2) # k-dim = Replicate (-1)
	337	+ x2_layout = _build_layout(mesh, (Shard(0),), 2) # k-dim (dim 0) = tp (0)
	338	+
	339	+ op = self._get_op()
	340	+ with self.assertRaisesRegex(ValueError, "k-dim.placement.must match"):
	341	+ op.infer_layout([x1_layout, x2_layout, False, True])
	342	+
	343	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	344	+ def test_input_m_multi_mesh_raises_14(self, mock_platform):
	345	+ """
	346	+ Feature: infer_layout rejects x1 m-dim jointly sharded across multiple mesh dims.
	347	+ Description: trans_input=False: x1 m-dim is dim 0; jointly sharded across dp+tp.
	348	+ Expectation: ValueError mentioning jointly sharded is not supported.
	349	+ """
	350	+ mesh = self._make_2d_mesh(mock_platform, shape=(2, 4), names=("dp", "tp"))
	351	+ x1_layout = _build_layout(mesh, (Shard(0), Shard(0)), 2)
	352	+ x2_layout = _build_layout(mesh, (Replicate(), Replicate()), 2)
	353	+
	354	+ op = self._get_op()
	355	+ with self.assertRaisesRegex(ValueError, "jointly sharded"):
	356	+ op.infer_layout([x1_layout, x2_layout, False, True])
	357	+
	358	+if __name__ == "__main__":
	359	+ unittest.main()

Atests/ut/core/shard/ops/test_parallel_matmul_reduce_scatter.py+298-0

@@ -0,0 +1,298 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Unit tests for MatmulReduceScatterDistributedOp."""
	16	+import os
	17	+import unittest
	18	+from unittest.mock import MagicMock, patch
	19	+
	20	+import numpy as np
	21	+
	22	+os.environ["HYPER_PARALLEL_PLATFORM"] = "torch"
	23	+
	24	+from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE
	25	+from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
	26	+from hyper_parallel.core.shard.ops.parallel_matmul_reduce_scatter import (
	27	+ MatmulReduceScatterDistributedOp,
	28	+ _normalize_mrs_args,
	29	+)
	30	+from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
	31	+from hyper_parallel.core.dtensor.device_mesh import init_device_mesh, _DEVICE_MESH_MAP
	32	+from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS
	33	+
	34	+
	35	+class TestMatmulReduceScatterDistributedOp(unittest.TestCase):
	36	+ """Unit tests for MatmulReduceScatterDistributedOp."""
	37	+
	38	+ def setUp(self):
	39	+ """Clear global state before each test."""
	40	+ EXISTING_COMM_GROUPS.clear()
	41	+ _DEVICE_MESH_MAP.clear()
	42	+ _LAYOUT_CACHE.clear()
	43	+
	44	+ def tearDown(self):
	45	+ """Clear global state after each test."""
	46	+ EXISTING_COMM_GROUPS.clear()
	47	+ _DEVICE_MESH_MAP.clear()
	48	+ _LAYOUT_CACHE.clear()
	49	+
	50	+ def _setup_mock_platform(self, mock_platform, world_size=8):
	51	+ """Configure mock platform for mesh creation."""
	52	+ mock_platform.get_rank.return_value = 0
	53	+ mock_platform.get_world_size.return_value = world_size
	54	+ mock_platform.split_group.return_value = MagicMock()
	55	+ mock_platform.tensor_to_numpy.side_effect = (
	56	+ lambda t: t.asnumpy() if hasattr(t, "asnumpy") else np.array(t)
	57	+ )
	58	+
	59	+ def _make_1d_mesh(self, mock_platform, size=4, name="tp"):
	60	+ """Return a 1-D mesh of given size."""
	61	+ self._setup_mock_platform(mock_platform, world_size=size)
	62	+ return init_device_mesh(device_type="npu", mesh_shape=(size,), mesh_dim_names=(name,),
	63	+ init_backend=False)
	64	+
	65	+ def _make_2d_mesh(self, mock_platform, shape=(2, 4), names=("dp", "tp")):
	66	+ """Return a 2-D mesh."""
	67	+ self._setup_mock_platform(mock_platform, world_size=shape[0] * shape[1])
	68	+ return init_device_mesh(device_type="npu", mesh_shape=shape, mesh_dim_names=names,
	69	+ init_backend=False)
	70	+
	71	+ @staticmethod
	72	+ def _get_op():
	73	+ return get_distributed_op("MatmulReduceScatter")
	74	+
	75	+ # ------------------------------------------------------------------
	76	+ # _normalize_mrs_args
	77	+ # ------------------------------------------------------------------
	78	+
	79	+ def test_normalize_args_defaults_1(self):
	80	+ """
	81	+ Feature: _normalize_mrs_args fills in default values.
	82	+ Description: Call with only 4 mandatory positional args.
	83	+ Expectation: Optional args get correct defaults.
	84	+ """
	85	+ x1, x2, group, world_size = object(), object(), "hccl_world_group", 2
	86	+ args, kwargs = _normalize_mrs_args(x1, x2, group, world_size)
	87	+ self.assertIs(args[0], x1)
	88	+ self.assertIs(args[1], x2)
	89	+ self.assertEqual(args[2], "hccl_world_group")
	90	+ self.assertEqual(args[3], 2)
	91	+ self.assertEqual(args[4], 'sum', msg=f"reduce_op default should be 'sum', got {args[4]}")
	92	+ self.assertIsNone(args[5], msg=f"bias default should be None, got {args[5]}")
	93	+ self.assertEqual(args[6], 0, msg=f"comm_turn default should be 0, got {args[6]}")
	94	+ self.assertFalse(args[7], msg=f"trans_input default should be False, got {args[7]}")
	95	+ self.assertFalse(args[8], msg=f"trans_x2 default should be False, got {args[8]}")
	96	+ self.assertEqual(kwargs, {}, msg=f"kwargs should be empty, got {kwargs}")
	97	+
	98	+ # ------------------------------------------------------------------
	99	+ # YAML registration
	100	+ # ------------------------------------------------------------------
	101	+
	102	+ def test_yaml_registration_mindspore_2(self):
	103	+ """
	104	+ Feature: YAML loader registers MatmulReduceScatterDistributedOp.
	105	+ Description: Call get_distributed_op with MindSpore CamelCase op name.
	106	+ Expectation: Returns a MatmulReduceScatterDistributedOp instance.
	107	+ """
	108	+ op = self._get_op()
	109	+ self.assertIsNotNone(op, msg="MatmulReduceScatter should be registered")
	110	+ self.assertIsInstance(
	111	+ op, MatmulReduceScatterDistributedOp,
	112	+ msg=f"Expected MatmulReduceScatterDistributedOp, got {type(op)}"
	113	+ )
	114	+
	115	+ # ------------------------------------------------------------------
	116	+ # infer_layout — positive cases
	117	+ # cache_values format: [x1_layout, x2_layout, trans_x2]
	118	+ # ------------------------------------------------------------------
	119	+
	120	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	121	+ def test_mrs_basic_tp_3(self, mock_platform):
	122	+ """
	123	+ Feature: infer_layout basic TP case.
	124	+ Description: x1 k-dim (dim 1) sharded on tp (mesh_dim 0); x2 k-dim (dim 0) sharded on tp.
	125	+ Pure TP: x1 m-dim (dim 0) is Replicate.
	126	+ Expectation: output tensor_map (0, -1) — m-dim sharded by tp (ReduceScatter), n-dim Replicate.
	127	+ """
	128	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	129	+ # x1 (m, k): m Replicate, k Shard(1) on tp
	130	+ x1_layout = _build_layout(mesh, (Shard(1),), 2)
	131	+ # x2 (k, n): k Shard(0) on tp, n Replicate
	132	+ x2_layout = _build_layout(mesh, (Shard(0),), 2)
	133	+
	134	+ op = self._get_op()
	135	+ out_layout, extra = op.infer_layout([x1_layout, x2_layout, False])
	136	+
	137	+ self.assertIsNone(extra)
	138	+ self.assertEqual(out_layout.tensor_map, (0, -1),
	139	+ msg=f"Basic TP: expected (0, -1), got {out_layout.tensor_map}")
	140	+
	141	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	142	+ def test_mrs_trans_x2_true_4(self, mock_platform):
	143	+ """
	144	+ Feature: infer_layout trans_x2=True.
	145	+ Description: x2 physical shape is (n, k); trans_x2=True means x2 dim 1 is k.
	146	+ x2 Shard(1) means n-dim (dim 0) is Replicate, k-dim (dim 1) is sharded on tp.
	147	+ Expectation: output tensor_map (0, -1) — m-dim sharded by tp, n-dim Replicate.
	148	+ """
	149	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	150	+ # x1 (m, k): k-dim (tensor dim 1) sharded on tp (mesh_dim 0)
	151	+ x1_layout = _build_layout(mesh, (Shard(1),), 2)
	152	+ # x2 (n, k): k-dim is tensor dim 1; Shard(1) shards k on tp (mesh_dim 0)
	153	+ x2_layout = _build_layout(mesh, (Shard(1),), 2)
	154	+
	155	+ op = self._get_op()
	156	+ out_layout, extra = op.infer_layout([x1_layout, x2_layout, True])
	157	+
	158	+ self.assertIsNone(extra)
	159	+ self.assertEqual(out_layout.tensor_map, (0, -1),
	160	+ msg=f"trans_x2=True: expected (0, -1), got {out_layout.tensor_map}")
	161	+
	162	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	163	+ def test_mrs_dp_tp_mesh_5(self, mock_platform):
	164	+ """
	165	+ Feature: infer_layout 2D (dp=2, tp=4) mesh, x1 m sharded by dp, k sharded by tp.
	166	+ Description: x1: m-dim (tensor dim 0) sharded by dp (mesh_dim 0),
	167	+ k-dim (tensor dim 1) sharded by tp (mesh_dim 1).
	168	+ x2: k-dim (tensor dim 0) sharded by tp (mesh_dim 1), n-dim Replicate.
	169	+ Expectation: output tensor_map ((1, 0), -1) — m jointly sharded by tp+dp, n Replicate.
	170	+ """
	171	+ mesh = self._make_2d_mesh(mock_platform, shape=(2, 4), names=("dp", "tp"))
	172	+ # x1 (m, k): m sharded by dp (mesh_dim 0), k sharded by tp (mesh_dim 1)
	173	+ x1_layout = _build_layout(mesh, (Shard(0), Shard(1)), 2)
	174	+ # x2 (k, n): k sharded by tp (mesh_dim 1), n Replicate
	175	+ x2_layout = _build_layout(mesh, (Replicate(), Shard(0)), 2)
	176	+
	177	+ op = self._get_op()
	178	+ out_layout, extra = op.infer_layout([x1_layout, x2_layout, False])
	179	+
	180	+ self.assertIsNone(extra)
	181	+ # tensor_map uses reversed mesh_dim indexing: value 0 = tp (innermost), value 1 = dp (outermost).
	182	+ # ReduceScatter converts k-sharding (tp, value 0) to m-sharding;
	183	+ # combined with dp m-sharding (value 1) → output_m = (0, 1).
	184	+ self.assertEqual(out_layout.tensor_map, ((0, 1), -1),
	185	+ msg=f"DP+TP mesh: expected ((0, 1), -1), got {out_layout.tensor_map}")
	186	+
	187	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	188	+ def test_mrs_x2_n_sharded_6(self, mock_platform):
	189	+ """
	190	+ Feature: infer_layout x2 n-dim sharded.
	191	+ Description: 2D mesh: x2 k sharded by tp (mesh_dim 1), n sharded by dp (mesh_dim 0).
	192	+ Expectation: output tensor_map (0, 1) — m sharded by tp, n sharded by dp.
	193	+ """
	194	+ mesh = self._make_2d_mesh(mock_platform, shape=(2, 4), names=("dp", "tp"))
	195	+ # x1 (m, k): m Replicate, k sharded by tp (mesh_dim 1)
	196	+ x1_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2)
	197	+ # x2 (k, n): k sharded by tp (mesh_dim 1), n sharded by dp (mesh_dim 0)
	198	+ x2_layout = _build_layout(mesh, (Shard(1), Shard(0)), 2)
	199	+
	200	+ op = self._get_op()
	201	+ out_layout, extra = op.infer_layout([x1_layout, x2_layout, False])
	202	+
	203	+ self.assertIsNone(extra)
	204	+ # tensor_map value 0 = tp (innermost), value 1 = dp (outermost).
	205	+ # output_m = 0 (tp ReduceScatter on m), output_n = 1 (dp shards n).
	206	+ self.assertEqual(out_layout.tensor_map, (0, 1),
	207	+ msg=f"x2 n sharded: expected (0, 1), got {out_layout.tensor_map}")
	208	+
	209	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	210	+ def test_mrs_x2_n_replicate_7(self, mock_platform):
	211	+ """
	212	+ Feature: infer_layout x2 n-dim Replicate.
	213	+ Description: Pure TP with Replicate n-dim.
	214	+ Expectation: output tensor_map (0, -1).
	215	+ """
	216	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	217	+ x1_layout = _build_layout(mesh, (Shard(1),), 2)
	218	+ x2_layout = _build_layout(mesh, (Shard(0),), 2)
	219	+
	220	+ op = self._get_op()
	221	+ out_layout, extra = op.infer_layout([x1_layout, x2_layout, False])
	222	+
	223	+ self.assertIsNone(extra)
	224	+ self.assertEqual(out_layout.tensor_map, (0, -1),
	225	+ msg=f"n Replicate: expected (0, -1), got {out_layout.tensor_map}")
	226	+
	227	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	228	+ def test_output_independent_copy_8(self, mock_platform):
	229	+ """
	230	+ Feature: infer_layout returns an independent deep copy of output layout.
	231	+ Description: Verify the returned layout is a new object, not a reference to any input.
	232	+ Expectation: output_layout is a distinct object from x1_layout and x2_layout.
	233	+ """
	234	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	235	+ x1_layout = _build_layout(mesh, (Shard(1),), 2)
	236	+ x2_layout = _build_layout(mesh, (Shard(0),), 2)
	237	+
	238	+ op = self._get_op()
	239	+ out_layout, _ = op.infer_layout([x1_layout, x2_layout, False])
	240	+
	241	+ self.assertIsNot(out_layout, x1_layout,
	242	+ msg="output_layout must not be the same object as x1_layout")
	243	+ self.assertIsNot(out_layout, x2_layout,
	244	+ msg="output_layout must not be the same object as x2_layout")
	245	+
	246	+ # ------------------------------------------------------------------
	247	+ # infer_layout — error cases
	248	+ # ------------------------------------------------------------------
	249	+
	250	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	251	+ def test_partial_input_raises_9(self, mock_platform):
	252	+ """
	253	+ Feature: infer_layout rejects Partial inputs.
	254	+ Description: x1 layout has Partial status.
	255	+ Expectation: ValueError with 'Partial status' message.
	256	+ """
	257	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	258	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	259	+ x1_layout.set_partial_by_dev_axis("tp", 'sum')
	260	+ x2_layout = _build_layout(mesh, (Shard(0),), 2)
	261	+
	262	+ op = self._get_op()
	263	+ with self.assertRaisesRegex(ValueError, "Partial status"):
	264	+ op.infer_layout([x1_layout, x2_layout, False])
	265	+
	266	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	267	+ def test_k_replicate_raises_10(self, mock_platform):
	268	+ """
	269	+ Feature: infer_layout rejects x1 k-dim (dim 1) Replicate.
	270	+ Description: k must be sharded for ReduceScatter to be semantically correct.
	271	+ Expectation: ValueError mentioning x1 k-dim must be Shard.
	272	+ """
	273	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	274	+ x1_layout = _build_layout(mesh, (Replicate(),), 2)
	275	+ x2_layout = _build_layout(mesh, (Replicate(),), 2)
	276	+
	277	+ op = self._get_op()
	278	+ with self.assertRaisesRegex(ValueError, "x1 k-dim.*must be Shard"):
	279	+ op.infer_layout([x1_layout, x2_layout, False])
	280	+
	281	+ @patch("hyper_parallel.core.dtensor.device_mesh.platform")
	282	+ def test_k_mismatch_raises_11(self, mock_platform):
	283	+ """
	284	+ Feature: infer_layout rejects x2 k-dim layout mismatch with x1 k-dim.
	285	+ Description: x1 k-dim sharded on mesh_dim 0, x2 k-dim Replicate → mismatch.
	286	+ Expectation: ValueError mentioning k layout must match.
	287	+ """
	288	+ mesh = self._make_1d_mesh(mock_platform, size=4, name="tp")
	289	+ x1_layout = _build_layout(mesh, (Shard(1),), 2)
	290	+ x2_layout = _build_layout(mesh, (Replicate(),), 2)
	291	+
	292	+ op = self._get_op()
	293	+ with self.assertRaisesRegex(ValueError, "x2 dim.*\\(k\\) layout must match"):
	294	+ op.infer_layout([x1_layout, x2_layout, False])
	295	+
	296	+
	297	+if __name__ == "__main__":
	298	+ unittest.main()