[BUG] CuTeDSL >= 4.6 cannot compile NCCL GIN device ops — ptxas rejects the generated PTX (4.5.2 works)

#!/usr/bin/env python3
"""Minimal reproducer: CuTeDSL >= 4.6 cannot compile any NCCL GIN *remote*
operation, while local-poll operations compile fine.

Two independent defects, so the script runs ONE case per process -- declaring
the case B externs in the same module changes the failure reported for case A,
so they must not share a process:

    python nvidia_issue_repro.py --case control   # gin.wait_counter  -> OK
    python nvidia_issue_repro.py --case v2        # gin.put/signal    -> defect 1
    python nvidia_issue_repro.py --case v1        # v1 externs        -> defect 2

Dependencies are three public wheels, nothing else:

    pip install 'nvidia-nccl-cu12==2.30.7' 'nccl4py[cu12]==0.4.1' \
                'nvidia-cutlass-dsl==4.7.1' 'cuda-python==12.9.7'

A single rank on any CUDA GPU is enough; nothing is transferred, both failures
are at compile time.
"""
import argparse

import torch
import cutlass
import cutlass.cute as cute
import cutlass.torch as cutlass_torch
import cuda.bindings.driver as cuda

import nccl.core as nccl
import nccl.core.device.cute as nccl_cute
import nccl.core.interop.torch as nccl_torch
from nccl.core.device.cute._helpers import _to_ptr, _to_coop_value, _to_value
from nccl.core.device.cute._structs import _LLVMPtrType, ncclTeam, ncclCoopAny

SLOT = 1 << 20


def make_v1_externs():
    """Declared lazily: the mere presence of these stubs in the module changes
    which error the stock (_v2) path reports, so case A must never see them."""
    from nccl.core.device.cute._bindings import _BC

    # Same signatures as nccl4py 0.4.1's ncclGinPut_v2 / ncclGinSignal_v2 stubs
    # minus the trailing opt_flags parameter, i.e. the v1 C-ABI still exported
    # by libnccl_device.bc.
    @cute.extern(name="ncclGinPut", source=_BC)
    def put_v1(
        gin: _LLVMPtrType, team: ncclTeam, peer: cutlass.Int32,
        dst_win: _LLVMPtrType, dst_offset: cutlass.Int64,
        src_win: _LLVMPtrType, src_offset: cutlass.Int64,
        size: cutlass.Int64,
        is_signal: cutlass.Boolean, signal_id: cutlass.Int32,
        signal_op: cutlass.Int32, signal_op_arg: cutlass.Int64,
        is_counter: cutlass.Boolean, counter_id: cutlass.Int32,
        coop: ncclCoopAny,
        is_descriptor: cutlass.Boolean, descriptor_ptr: _LLVMPtrType,
        given_release: cutlass.Int32, required_release: cutlass.Int32,
    ) -> None: ...

    @cute.extern(name="ncclGinSignal", source=_BC)
    def signal_v1(
        gin: _LLVMPtrType, team: ncclTeam, peer: cutlass.Int32,
        is_signal: cutlass.Boolean, signal_id: cutlass.Int32,
        signal_op: cutlass.Int32, signal_op_arg: cutlass.Int64,
        coop: ncclCoopAny,
        is_descriptor: cutlass.Boolean, descriptor_ptr: _LLVMPtrType,
        given_release: cutlass.Int32, required_release: cutlass.Int32,
    ) -> None: ...

    return put_v1, signal_v1


def build(case, op):
    class K:
        @cute.jit
        def run(self, dev_comm, win, stream: cuda.CUstream):
            self.kern(dev_comm, win).launch(
                grid=[1, 1, 1], block=[128, 1, 1], stream=stream)

        @cute.kernel
        def kern(self, dev_comm, win):
            team = dev_comm.team_world
            coop = nccl_cute.cta()
            gin = dev_comm.gin(nccl_cute.GinBackendMask.PROXY, 0)
            st = win.tensor(cutlass.Int8, cute.make_layout(SLOT), 0)
            dt = win.tensor(cutlass.Int8, cute.make_layout(SLOT), cutlass.Int64(SLOT))

            if cutlass.const_expr(case == "control"):
                gin.wait_counter(coop, counter=0, least=1)

            elif cutlass.const_expr(case == "v2"):
                if cutlass.const_expr(op == "put"):
                    gin.put(team, 0, win, dt, win, st, coop,
                            is_signal=True, signal_id=0, signal_op=1,
                            signal_op_arg=1, is_counter=True, counter_id=0)
                else:
                    gin.signal(team, 0, True, 0, 1, 1, coop)

            else:  # v1 externs
                put_v1, signal_v1 = K._externs
                if cutlass.const_expr(op == "put"):
                    put_v1(gin.ptr, _to_value(team), cutlass.Int32(0),
                           win.ptr, cutlass.Int64(SLOT),
                           win.ptr, cutlass.Int64(0),
                           cutlass.Int64(SLOT),
                           cutlass.Boolean(True), cutlass.Int32(0),
                           cutlass.Int32(1), cutlass.Int64(1),
                           cutlass.Boolean(True), cutlass.Int32(0),
                           _to_coop_value(coop),
                           cutlass.Boolean(False), _to_ptr(0),
                           cutlass.Int32(0), cutlass.Int32(2))
                else:
                    signal_v1(gin.ptr, _to_value(team), cutlass.Int32(0),
                              cutlass.Boolean(True), cutlass.Int32(0),
                              cutlass.Int32(1), cutlass.Int64(1),
                              _to_coop_value(coop),
                              cutlass.Boolean(False), _to_ptr(0),
                              cutlass.Int32(0), cutlass.Int32(2))

    if case == "v1":
        K._externs = make_v1_externs()
    return K()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--case", default="control", choices=["control", "v2", "v1"])
    a = ap.parse_args()

    import importlib.metadata as md
    for p in ("nvidia-cutlass-dsl", "nccl4py", "nvidia-nccl-cu12", "cuda-python"):
        try:
            print(f"{p:22s} {md.version(p)}")
        except Exception:
            print(f"{p:22s} <not installed>")
    print(f"{'gpu':22s} {torch.cuda.get_device_name(0)}")
    print(f"{'case':22s} {a.case}\n")

    torch.cuda.set_device(0)
    dev = torch.device("cuda", 0)
    comm = nccl.Communicator.init(nranks=1, rank=0, unique_id=nccl.get_unique_id())
    reqs = nccl.NCCLDevCommRequirements(
        gin_connection_type=nccl.NcclGinConnectionType.FULL,
        gin_context_count=1, gin_signal_count=8, gin_counter_count=1)
    devcomm = comm.create_dev_comm(requirements=reqs)
    buf = nccl_torch.empty(SLOT * 2, dtype=torch.uint8, device=dev)
    win = comm.register_window(buf)
    stream = cutlass_torch.current_stream()

    ops = ["put"] if a.case == "control" else ["put", "signal"]
    for op in ops:
        label = "wait_counter" if a.case == "control" else f"{op} ({a.case})"
        try:
            cute.compile(build(a.case, op).run, devcomm, win, stream)
            print(f"{label:20s} COMPILED_OK")
        except BaseException as e:
            s = str(e)
            # the distinguishing line is inside the quoted ptxas log, not the
            # generic "ptxas rejected the PTX ..." wrapper
            hit = next((ln.strip() for ln in s.splitlines()
                        if "Parsing error near" in ln or "Modifier" in ln), None)
            print(f"{label:20s} FAILED: {hit or s.splitlines()[0][:120]}")


if __name__ == "__main__":
    main()