Diff Coverage

    Returns:
        HCCL communicator handle name expected by ``torch_npu`` MC2 kernels.
    """
    rank = torch.distributed.get_rank(group)
    if torch.__version__ > "2.0":
        global_rank = torch.distributed.get_global_rank(group, rank)
        # torch.distributed ProcessGroup exposes HCCL via a private backend API.
        return group._get_backend(torch.device("npu")).get_hccl_comm_name(  # pylint: disable=protected-access
            global_rank
        )
    return group.get_hccl_comm_name(rank)


def _require_torch_npu():
    try:
        import torch_npu  # pylint: disable=import-outside-toplevel
    except ImportError as exc:
        raise RuntimeError(
            "MC2 fused kernels require torch_npu "
            "(npu_all_gather_base_mm / npu_mm_reduce_scatter_base)."
        ) from exc
    return torch_npu


class AllGatherMatmulFunction(platform.Function):
    """Column-parallel fused all-gather + matmul with custom backward.

    @staticmethod
    def forward(ctx, x, w, group, world_size, bias):  # pylint: disable=arguments-differ
        """Run fused all-gather + matmul and stash tensors for backward."""
        torch_npu = _require_torch_npu()
        hcom = get_hcomm_info(group)
        # x2 = w.T -> physical (k, n_local); kernel computes AG(x) @ x2
        out, gathered = torch_npu.npu_all_gather_base_mm(
            x,
            w.t(),
            hcom,
            world_size,

            bias=None,
            gather_index=0,
            gather_output=True,
        )
        if bias is not None:
            out = out + bias
        ctx.save_for_backward(gathered, w)
        ctx.group = group
        ctx.world_size = world_size
        ctx.hcom = hcom
        ctx.has_bias = bias is not None
        return out

    @staticmethod
    def backward(ctx, grad_out):  # pylint: disable=arguments-differ
        """Gradient: dx via fused matmul_reduce_scatter, dw via gathered-input matmul.

    @staticmethod
    def forward(ctx, x, w, group, world_size, bias):  # pylint: disable=arguments-differ
        """Run fused matmul + reduce-scatter and stash tensors for backward."""
        torch_npu = _require_torch_npu()
        hcom = get_hcomm_info(group)
        out = torch_npu.npu_mm_reduce_scatter_base(
            x,
            w.t(),
            hcom,
            world_size,

            world_size,
            reduce_op="sum",
            bias=None,
        )
        if bias is not None:
            out = out + bias
        ctx.save_for_backward(x, w)
        ctx.group = group
        ctx.world_size = world_size
        ctx.hcom = hcom
        ctx.has_bias = bias is not None
        return out

    @staticmethod
    def backward(ctx, grad_out):  # pylint: disable=arguments-differ
        """Gradient: dx via fused all-gather+matmul, dw via gathered-grad matmul."""
        torch_npu = _require_torch_npu()
        x, w = ctx.saved_tensors
        # AG(dY) @ W : pass W (n, k) so the kernel uses transposed-x2 semantics.
        grad_x, grad_out_full = torch_npu.npu_all_gather_base_mm(
            grad_out,
            w,
            ctx.hcom,
            ctx.world_size,

            bias=None,
            gather_index=0,
            gather_output=True,
        )
        grad_w = grad_out_full.t().matmul(x)
        grad_bias = grad_out_full.sum(dim=0) if ctx.has_bias else None
        return grad_x, grad_w, None, None, grad_bias


class MC2Linear(nn.Linear):
    """``nn.Linear`` that uses fused matmul + TP communication kernels."""

        return linear

    def _mc2_forward(self, input_: DTensor, weight: DTensor) -> DTensor:
        """Run the configured fused kernel on local tensors."""
        input_local = input_.to_local()
        weight_local = weight.to_local()
        leading_global = tuple(int(s) for s in input_.shape[:-1])
        input_2d = input_local.reshape(-1, input_local.shape[-1])

        bias_local = None
        if self.bias is not None:
            bias = self.bias
            bias_local = bias.to_local() if isinstance(bias, DTensor) else bias

        if self.mc2_mode == "all_gather":
            output_2d = AllGatherMatmulFunction.apply(
                input_2d,
                weight_local,
                self.mc2_group,
                self.mc2_world_size,

                self.mc2_world_size,
                bias_local,
            )
            # After AG on flattened leading dims, reshape with global leading shape.
            output = output_2d.reshape(*leading_global, output_2d.shape[-1])
            return DTensor.from_local(output, input_.device_mesh, (Shard(-1),))

        output_2d = MatmulReduceScatterFunction.apply(
            input_2d,
            weight_local,
            self.mc2_group,
            self.mc2_world_size,

            self.mc2_group,
            self.mc2_world_size,
            bias_local,
        )
        leading_local = list(leading_global)
        seq_dim = self.mc2_sequence_dim
        if seq_dim < 0:
            seq_dim += len(leading_local)
        if seq_dim < 0 or seq_dim >= len(leading_local):
            raise RuntimeError(
                f"MC2Linear sequence_dim={self.mc2_sequence_dim} is out of range "
                f"for input rank {input_.dim()}."
            )
        if leading_local[seq_dim] % self.mc2_world_size != 0:
            raise RuntimeError(
                f"MC2Linear reduce_scatter requires sequence dim {seq_dim} "
                f"(size {leading_local[seq_dim]}) divisible by world_size "
                f"{self.mc2_world_size}."
            )
        leading_local[seq_dim] = leading_local[seq_dim] // self.mc2_world_size
        output = output_2d.reshape(*leading_local, output_2d.shape[-1])
        return DTensor.from_local(output, input_.device_mesh, (Shard(seq_dim),))

    def forward(self, input_: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Forward using MC2 for DTensor inputs and ``nn.Linear`` otherwise."""
        if not isinstance(input_, DTensor):
            return super().forward(input_)
        if not hasattr(self, "mc2_mode"):
            raise RuntimeError(
                "MC2Linear must be configured by an MC2 parallel style before use."
            )

        if weight is None:
            weight = self.weight
        if not isinstance(weight, DTensor):
            raise TypeError(
                "MC2Linear expects a DTensor weight after tensor-parallel sharding, "
                f"but got {type(weight).__name__}."
            )
        return self._mc2_forward(input_, weight)