# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import os

import torch
import torch.distributed as dist
from torch.distributed.tensor import distribute_tensor, Shard

import nvfuser_direct as nvfuser
from .benchmark_utils import get_benchmark_fns
from nvfuser_direct import DataType, FusionDefinition, CommunicatorBackend, TensorView
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

def row_parallel_linear_forward(
    h: int, num_devices: int, num_chunks: int
) -> FusionDefinition:
    with FusionDefinition() as fd:
        inp = fd.define_tensor(
            shape=[-1, h * 4], contiguity=True, dtype=DataType.BFloat16
        )
        weight = fd.define_tensor(
            shape=[h, h * 4], contiguity=True, dtype=DataType.BFloat16
        )
        out = fd.ops.linear(inp, weight)
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
        for tv in (inp, weight):
            tv.set_device_mesh(mesh)

        inp.outer_split(0, num_chunks)
        inp.axis(0).parallelize(nvfuser.ParallelType.stream)
        inp.outer_split(2, mesh.size)
        inp.axis(2).parallelize(nvfuser.ParallelType.mesh_x)
        weight.outer_split(1, mesh.size)
        weight.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    # Expected pre-segmentation IR:
    #
    #   [t, 4h]                                 [h, 4h]
    #   /\  /\                                      /\.
    #  s*  d                                       d
    #                      |
    #                      | linear
    #                      |
    #                          r{4h}
    #                          /  \.
    #                 [t, h, d, r{4h/d}]
    #                 /\.
    #                s
    #                     |
    #                     | sum
    #                     |
    #                  [t, h, r{d}]
    #                  /\.
    #                 s*

    # The host IR dumped with NVFUSER_DUMP=host_ir is similar to `row_parallel_linear_forward_reference`:
    #
    # %HostIrContainer { (T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1})) -> (T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1})) :
    #   T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}) = ALLOCATE(buffer=T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), mem_type=global, size=( i0 * 2 ), zero_init=false, resets_to_zero=false)
    #   Stream 0x174e5c80 = GetCurrentStream()
    #   FOR i535 from 0 to 3:
    #     SetCurrentStream(Stream i535)
    #     Synchronize(Stream 0x174e5c80)
    #     T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}) = ShardByStream(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), stream_index = i535)
    #     T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}) = ALLOCATE(buffer=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), mem_type=global, size=( ( ceilDiv(i0, 3) ) * 12 ), zero_init=false, resets_to_zero=false)
    #     T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1})
    #        = linear(T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}),
    #                 T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1})      )
    #     T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}) = ShardByStream(T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), stream_index = i535)
    #     Communication 272 (type=Allreduce, team=(0 1), input=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), output=T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}), backend=NCCL)
    #     Wait(Communication 272)
    #   SetCurrentStream(Stream 0x174e5c80)
    #   FOR i535 from 0 to 3:
    #     Synchronize(Stream i535)
    # } // %HostIrContainer

    return fd

@pytest.mark.mpi
def test_row_parallel_linear_forward(multidevice_test):
    # This is a port of CollectiveBasedOverlapTest.RowParallelLinear_Forward.
    h, s, t = 2, 3, 6
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    fd = row_parallel_linear_forward(h, d, s)

    inp_ref = torch.testing.make_tensor(t, h * 4, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])
    # nvfuser_direct.PythonProfiler failed with host IR lowering. The main
    # reason is that HostIrContainer doesn't keep segments while SegmentProfiler
    # is still expecting data.  It's unclear to me whether we should relax
    # SegmentProfiler's assumptions or stop creating them in the first place.
    with torch.profiler.profile(record_shapes=True) as prof:
        (out,) = fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    torch.testing.assert_close(out.cpu(), out_ref)

    matmul_events = [event for event in prof.events() if event.name == "aten::mm"]
    assert len(matmul_events) == s

    m = t // s
    n = h
    k = h * 4 // d
    for event in matmul_events:
        assert event.input_shapes == [[m, k], [k, n], [m, n]]

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize("s", [1, 2, 4])
def test_row_parallel_linear_forward_benchmark(multidevice_test, benchmark, s):
    # This is a port of CollectiveBasedOverlapTest.RowParallelLinear_Forward.
    h, t = 8192, 8192
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    fd = row_parallel_linear_forward(h, d, s)

    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16, device="cpu")

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

# The caching allocator in PyTorch can't cache buffers across streams, so we
# have to reuse streams to avoid repeated cudaMalloc. torch.cuda.Stream() is
# backed by a stream pool as well but I failed to find a way to set its size.
class StreamPool:
    def __init__(self):
        self._streams = {}

    def get(self, sid: int) -> torch.cuda.Stream:
        s = self._streams.get(sid)
        if s is None:
            s = torch.cuda.Stream()
            self._streams[sid] = s
        return s

def row_parallel_linear_forward_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    num_chunks: int,
    stream_pool: StreamPool,
) -> torch.Tensor:
    out = torch.empty(
        inp_shard.size(0),
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    inp_chunks = inp_shard.chunk(num_chunks)
    out_chunks = out.chunk(num_chunks)

    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i, (inp_chunk, out_chunk) in enumerate(zip(inp_chunks, out_chunks)):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            torch.matmul(inp_chunk, weight_shard.T, out=out_chunk)
            work = dist.all_reduce(out_chunk, op=dist.ReduceOp.SUM, async_op=True)
            work.wait()

    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)

    return out

@pytest.mark.mpi
def test_row_parallel_linear_forward_reference(setup_default_process_group):
    h, s, t = 2, 3, 6
    d = dist.get_world_size()
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    torch.manual_seed(0)
    inp_ref = torch.testing.make_tensor(t, h * 4, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_ref.cuda()).cpu()

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(-1)]).to_local()
    weight_shard = distribute_tensor(
        weight_ref, mesh, placements=[Shard(-1)]
    ).to_local()
    stream_pool = StreamPool()
    out = row_parallel_linear_forward_reference(inp_shard, weight_shard, s, stream_pool)

    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
@pytest.mark.benchmark
def test_row_parallel_linear_forward_reference_benchmark(
    setup_default_process_group, benchmark
):
    h, s, t = 8192, 2, 8192
    d = dist.get_world_size()
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    torch.manual_seed(0)
    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16)
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16)

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(-1)]).to_local()
    weight_shard = distribute_tensor(
        weight_ref, mesh, placements=[Shard(-1)]
    ).to_local()

    stream_pool = StreamPool()
    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: row_parallel_linear_forward_reference(
            inp_shard, weight_shard, s, stream_pool
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

def column_parallel_linear_forward_cyclic_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    stream_pool: StreamPool,
) -> torch.Tensor:
    d = dist.get_world_size()
    my_rank = dist.get_rank()
    out = torch.empty(
        inp_shard.size(0) * d,
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    out_chunks = out.chunk(d)
    main_stream = torch.cuda.current_stream()
    worker_streams = []
    buffers = [inp_shard] + [torch.empty_like(inp_shard) for _ in range(d - 1)]
    for i in range(d):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            if i > 0:
                send_op = dist.P2POp(dist.isend, buffers[i - 1], (my_rank + 1) % d)
                recv_op = dist.P2POp(dist.irecv, buffers[i], (my_rank - 1 + d) % d)
                reqs = dist.batch_isend_irecv([send_op, recv_op])
                for req in reqs:
                    req.wait()
            torch.matmul(
                buffers[i],
                weight_shard.T,
                out=out_chunks[(my_rank - i + d) % d],
            )
    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)
    return out

def column_parallel_linear_forward_uncyclic_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    stream_pool: StreamPool,
) -> torch.Tensor:
    d = dist.get_world_size()
    my_rank = dist.get_rank()
    out = torch.empty(
        inp_shard.size(0) * d,
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    out_chunks = out.chunk(d)
    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i in range(d):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            if i == 0:
                buffer = inp_shard
            else:
                buffer = torch.empty_like(inp_shard)
                send_op = dist.P2POp(dist.isend, inp_shard, (my_rank - i + d) % d)
                recv_op = dist.P2POp(dist.irecv, buffer, (my_rank + i) % d)
                reqs = dist.batch_isend_irecv([send_op, recv_op])
                for req in reqs:
                    req.wait()
            torch.matmul(
                buffer,
                weight_shard.T,
                out=out_chunks[(my_rank + i) % d],
            )
    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)
    return out

@pytest.mark.mpi
@pytest.mark.parametrize(
    "column_parallel_linear_fn",
    [
        column_parallel_linear_forward_cyclic_reference,
        column_parallel_linear_forward_uncyclic_reference,
    ],
    ids=["cyclic", "uncyclic"],
)
def test_column_parallel_linear_forward_reference(
    setup_default_process_group, column_parallel_linear_fn
):
    h, t = 6, 24
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    if 4 * h % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {4 * h} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )

    inp_ref = torch.testing.make_tensor(t, h, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(0)]).to_local()
    weight_ref = torch.testing.make_tensor(
        4 * h, h, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    weight_shard = distribute_tensor(weight_ref, mesh, placements=[Shard(0)]).to_local()
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_shard)
    stream_pool = StreamPool()
    out = column_parallel_linear_fn(inp_shard, weight_shard, stream_pool)
    torch.testing.assert_close(out, out_ref)

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize(
    "column_parallel_linear_fn",
    [
        column_parallel_linear_forward_cyclic_reference,
        column_parallel_linear_forward_uncyclic_reference,
    ],
    ids=["cyclic", "uncyclic"],
)
def test_column_parallel_linear_forward_reference_benchmark(
    benchmark, setup_default_process_group, column_parallel_linear_fn
):
    h, t = 8192, 8192
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_ref = torch.randn(t, h, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(4 * h, h, dtype=torch.bfloat16, device="cpu")
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(0)]).to_local()
    weight_shard = distribute_tensor(weight_ref, mesh, placements=[Shard(0)]).to_local()
    stream_pool = StreamPool()
    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: column_parallel_linear_fn(inp_shard, weight_shard, stream_pool)
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

def column_parallel_linear_forward(h: int, d: int, parallelism: str):
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, h), contiguity=True, dtype=DataType.BFloat16)
        weight_tv = fd.define_tensor(
            (4 * h, h), contiguity=True, dtype=DataType.BFloat16
        )
        ag_out = fd.ops.set(inp_tv)
        out_tv = fd.ops.linear(ag_out, weight_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

        for tv in [inp_tv, weight_tv]:
            tv.set_device_mesh(mesh)
            tv.outer_split(0, d)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        ag_out.set_device_mesh(mesh)
        ag_out.outer_split(0, d)
        if parallelism == "collective_permute":
            ag_out.swizzle1d(0, nvfuser.ParallelType.mesh_x)
        ag_out.axis(0).parallelize(nvfuser.ParallelType.stream)

        # Fusion IR before segmentation will look like this:
        #   [t, h]
        #   /\.
        #  d
        # (deviceIdx.x)
        #    |
        #    | set (lowered to Broadcast/CollectivePermute. This decomposition is done
        #    |    manually in the definition above. It will later be done
        #    |    by preseg.)
        #    |
        #   [t, h]                                  [4h,  h]
        #   /\                                      /\.
        #  d                                       d
        #  | swizzle1d (if parallelism == "collective_permute")
        #  s
        # (streamIdx)
        #                      |
        #                      | linear
        #                      |
        #                   [t, 4h, r{h}]
        #                   /\  /\.
        #                  s*   d

    return fd

@pytest.mark.mpi
@pytest.mark.parametrize("parallelism", ["collective_permute", "broadcast"])
def test_column_parallel_linear_forward(multidevice_test, parallelism: str):
    # This is a port of CollectiveBasedOverlapTest.ColumnAndSequenceParallelLinear_Forward.
    # The difference is we are using broadcast based overlapping instead of send/recv.
    h, t = 2, 24
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )

    fd = column_parallel_linear_forward(h, d, parallelism)

    inp_ref = torch.testing.make_tensor(t, h, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        4 * h, h, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight)

    with torch.profiler.profile(record_shapes=True) as prof:
        (out,) = fd.execute([inp, weight], _enable_options=["host_ir_lowering"])
    torch.testing.assert_close(out, out_ref)

    if parallelism == "collective_permute":
        collective_permute_events = [
            event for event in prof.events() if "ncclDevKernel_SendRecv" in event.name
        ]
        assert len(collective_permute_events) == (d - 1)
    else:
        broadcast_events = [
            event for event in prof.events() if "ncclDevKernel_Broadcast" in event.name
        ]
        assert len(broadcast_events) == (d if d > 1 else 0)

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize("parallelism", ["collective_permute", "broadcast"])
def test_column_parallel_linear_forward_benchmark(
    multidevice_test, benchmark, parallelism: str
):
    # This is a port of CollectiveBasedOverlapTest.ColumnParall
# ... truncated ...


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import os

import torch
import torch.distributed as dist
from torch.distributed.tensor import distribute_tensor, Shard

import nvfuser_direct as nvfuser
from .benchmark_utils import get_benchmark_fns
from nvfuser_direct import DataType, FusionDefinition, CommunicatorBackend, TensorView
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype

def row_parallel_linear_forward(
    h: int, num_devices: int, num_chunks: int
) -> FusionDefinition:
    with FusionDefinition() as fd:
        inp = fd.define_tensor(
            shape=[-1, h * 4], contiguity=True, dtype=DataType.BFloat16
        )
        weight = fd.define_tensor(
            shape=[h, h * 4], contiguity=True, dtype=DataType.BFloat16
        )
        out = fd.ops.linear(inp, weight)
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
        for tv in (inp, weight):
            tv.set_device_mesh(mesh)

        inp.outer_split(0, num_chunks)
        inp.axis(0).parallelize(nvfuser.ParallelType.stream)
        inp.outer_split(2, mesh.size)
        inp.axis(2).parallelize(nvfuser.ParallelType.mesh_x)
        weight.outer_split(1, mesh.size)
        weight.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    # Expected pre-segmentation IR:
    #
    #   [t, 4h]                                 [h, 4h]
    #   /\  /\                                      /\.
    #  s*  d                                       d
    #                      |
    #                      | linear
    #                      |
    #                          r{4h}
    #                          /  \.
    #                 [t, h, d, r{4h/d}]
    #                 /\.
    #                s
    #                     |
    #                     | sum
    #                     |
    #                  [t, h, r{d}]
    #                  /\.
    #                 s*

    # The host IR dumped with NVFUSER_DUMP=host_ir is similar to `row_parallel_linear_forward_reference`:
    #
    # %HostIrContainer { (T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1})) -> (T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1})) :
    #   T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}) = ALLOCATE(buffer=T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), mem_type=global, size=( i0 * 2 ), zero_init=false, resets_to_zero=false)
    #   Stream 0x174e5c80 = GetCurrentStream()
    #   FOR i535 from 0 to 3:
    #     SetCurrentStream(Stream i535)
    #     Synchronize(Stream 0x174e5c80)
    #     T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}) = ShardByStream(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), stream_index = i535)
    #     T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}) = ALLOCATE(buffer=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), mem_type=global, size=( ( ceilDiv(i0, 3) ) * 12 ), zero_init=false, resets_to_zero=false)
    #     T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1})
    #        = linear(T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}),
    #                 T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1})      )
    #     T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}) = ShardByStream(T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), stream_index = i535)
    #     Communication 272 (type=Allreduce, team=(0 1), input=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), output=T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}), backend=NCCL)
    #     Wait(Communication 272)
    #   SetCurrentStream(Stream 0x174e5c80)
    #   FOR i535 from 0 to 3:
    #     Synchronize(Stream i535)
    # } // %HostIrContainer

    return fd

@pytest.mark.mpi
def test_row_parallel_linear_forward(multidevice_test):
    # This is a port of CollectiveBasedOverlapTest.RowParallelLinear_Forward.
    h, s, t = 2, 3, 6
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    fd = row_parallel_linear_forward(h, d, s)

    inp_ref = torch.testing.make_tensor(t, h * 4, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"


    #   FOR i535 from 0 to 3:
    #     SetCurrentStream(Stream i535)
    #     Synchronize(Stream 0x174e5c80)
    #     T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}) = ShardByStream(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), stream_index = i535)
    #     T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}) = ALLOCATE(buffer=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), mem_type=global, size=( ( ceilDiv(i0, 3) ) * 12 ), zero_init=false, resets_to_zero=false)
    #     T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1})
    #        = linear(T4_l___bfloat[istreamIdx37{3}, iS38{( ceilDiv(i0, 3) )}, ideviceIdx.x35{2}, iS36{4}] (DeviceMesh{0 1}),
    #                 T1_g___bfloat[ideviceIdx.x11{2}, iS2{2}, iS12{4}] (DeviceMesh{0 1})      )
    #     T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}) = ShardByStream(T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{2}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0 1}), stream_index = i535)
    #     Communication 272 (type=Allreduce, team=(0 1), input=T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{2}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{4}rf] (DeviceMesh{0 1}), output=T5_l___bfloat[istreamIdx41{3}, iS42{( ceilDiv(i0, 3) )}, iS40{2}] (DeviceMesh{0 1}), backend=NCCL)
    #     Wait(Communication 272)
    #   SetCurrentStream(Stream 0x174e5c80)
    #   FOR i535 from 0 to 3:
    #     Synchronize(Stream i535)
    # } // %HostIrContainer

    return fd

@pytest.mark.mpi
def test_row_parallel_linear_forward(multidevice_test):
    # This is a port of CollectiveBasedOverlapTest.RowParallelLinear_Forward.
    h, s, t = 2, 3, 6
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    fd = row_parallel_linear_forward(h, d, s)

    inp_ref = torch.testing.make_tensor(t, h * 4, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])
    # nvfuser_direct.PythonProfiler failed with host IR lowering. The main
    # reason is that HostIrContainer doesn't keep segments while SegmentProfiler
    # is still expecting data.  It's unclear to me whether we should relax
    # SegmentProfiler's assumptions or stop creating them in the first place.
    with torch.profiler.profile(record_shapes=True) as prof:
        (out,) = fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    torch.testing.assert_close(out.cpu(), out_ref)

    matmul_events = [event for event in prof.events() if event.name == "aten::mm"]
    assert len(matmul_events) == s

    m = t // s
    n = h
    k = h * 4 // d
    for event in matmul_events:
        assert event.input_shapes == [[m, k], [k, n], [m, n]]

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize("s", [1, 2, 4])
def test_row_parallel_linear_forward_benchmark(multidevice_test, benchmark, s):
    # This is a port of CollectiveBasedOverlapTest.RowParallelLinear_Forward.
    h, t = 8192, 8192
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    fd = row_parallel_linear_forward(h, d, s)

    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16, device="cpu")


        (out,) = fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    torch.testing.assert_close(out.cpu(), out_ref)

    matmul_events = [event for event in prof.events() if event.name == "aten::mm"]
    assert len(matmul_events) == s

    m = t // s
    n = h
    k = h * 4 // d
    for event in matmul_events:
        assert event.input_shapes == [[m, k], [k, n], [m, n]]

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize("s", [1, 2, 4])
def test_row_parallel_linear_forward_benchmark(multidevice_test, benchmark, s):
    # This is a port of CollectiveBasedOverlapTest.RowParallelLinear_Forward.
    h, t = 8192, 8192
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    fd = row_parallel_linear_forward(h, d, s)

    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16, device="cpu")

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

# The caching allocator in PyTorch can't cache buffers across streams, so we
# have to reuse streams to avoid repeated cudaMalloc. torch.cuda.Stream() is
# backed by a stream pool as well but I failed to find a way to set its size.
class StreamPool:
    def __init__(self):
        self._streams = {}

    def get(self, sid: int) -> torch.cuda.Stream:
        s = self._streams.get(sid)
        if s is None:
            s = torch.cuda.Stream()
            self._streams[sid] = s
        return s

def row_parallel_linear_forward_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,


StreamPool

    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16, device="cpu")

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

# The caching allocator in PyTorch can't cache buffers across streams, so we
# have to reuse streams to avoid repeated cudaMalloc. torch.cuda.Stream() is
# backed by a stream pool as well but I failed to find a way to set its size.
class StreamPool:
    def __init__(self):
        self._streams = {}

    def get(self, sid: int) -> torch.cuda.Stream:
        s = self._streams.get(sid)
        if s is None:
            s = torch.cuda.Stream()
            self._streams[sid] = s
        return s

def row_parallel_linear_forward_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    num_chunks: int,
    stream_pool: StreamPool,
) -> torch.Tensor:
    out = torch.empty(
        inp_shard.size(0),
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    inp_chunks = inp_shard.chunk(num_chunks)
    out_chunks = out.chunk(num_chunks)

    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i, (inp_chunk, out_chunk) in enumerate(zip(inp_chunks, out_chunks)):


    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

# The caching allocator in PyTorch can't cache buffers across streams, so we
# have to reuse streams to avoid repeated cudaMalloc. torch.cuda.Stream() is
# backed by a stream pool as well but I failed to find a way to set its size.
class StreamPool:
    def __init__(self):
        self._streams = {}

    def get(self, sid: int) -> torch.cuda.Stream:
        s = self._streams.get(sid)
        if s is None:
            s = torch.cuda.Stream()
            self._streams[sid] = s
        return s

def row_parallel_linear_forward_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    num_chunks: int,
    stream_pool: StreamPool,
) -> torch.Tensor:
    out = torch.empty(
        inp_shard.size(0),
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    inp_chunks = inp_shard.chunk(num_chunks)
    out_chunks = out.chunk(num_chunks)

    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i, (inp_chunk, out_chunk) in enumerate(zip(inp_chunks, out_chunks)):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            torch.matmul(inp_chunk, weight_shard.T, out=out_chunk)
            work = dist.all_reduce(out_chunk, op=dist.ReduceOp.SUM, async_op=True)
            work.wait()

    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)

    return out

@pytest.mark.mpi
def test_row_parallel_linear_forward_reference(setup_default_process_group):
    h, s, t = 2, 3, 6
    d = dist.get_world_size()
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    torch.manual_seed(0)
    inp_ref = torch.testing.make_tensor(t, h * 4, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_ref.cuda()).cpu()


    out_chunks = out.chunk(num_chunks)

    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i, (inp_chunk, out_chunk) in enumerate(zip(inp_chunks, out_chunks)):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            torch.matmul(inp_chunk, weight_shard.T, out=out_chunk)
            work = dist.all_reduce(out_chunk, op=dist.ReduceOp.SUM, async_op=True)
            work.wait()

    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)

    return out

@pytest.mark.mpi
def test_row_parallel_linear_forward_reference(setup_default_process_group):
    h, s, t = 2, 3, 6
    d = dist.get_world_size()
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    torch.manual_seed(0)
    inp_ref = torch.testing.make_tensor(t, h * 4, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_ref.cuda()).cpu()

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(-1)]).to_local()
    weight_shard = distribute_tensor(
        weight_ref, mesh, placements=[Shard(-1)]
    ).to_local()
    stream_pool = StreamPool()
    out = row_parallel_linear_forward_reference(inp_shard, weight_shard, s, stream_pool)

    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
@pytest.mark.benchmark
def test_row_parallel_linear_forward_reference_benchmark(
    setup_default_process_group, benchmark
):
    h, s, t = 8192, 2, 8192
    d = dist.get_world_size()
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    torch.manual_seed(0)
    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16)
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16)

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])


        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        h, h * 4, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_ref.cuda()).cpu()

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(-1)]).to_local()
    weight_shard = distribute_tensor(
        weight_ref, mesh, placements=[Shard(-1)]
    ).to_local()
    stream_pool = StreamPool()
    out = row_parallel_linear_forward_reference(inp_shard, weight_shard, s, stream_pool)

    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
@pytest.mark.benchmark
def test_row_parallel_linear_forward_reference_benchmark(
    setup_default_process_group, benchmark
):
    h, s, t = 8192, 2, 8192
    d = dist.get_world_size()
    if (h * 4) % d != 0:
        pytest.skip(
            f"Row-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    assert t % s == 0

    torch.manual_seed(0)
    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16)
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16)

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(-1)]).to_local()
    weight_shard = distribute_tensor(
        weight_ref, mesh, placements=[Shard(-1)]
    ).to_local()

    stream_pool = StreamPool()
    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: row_parallel_linear_forward_reference(
            inp_shard, weight_shard, s, stream_pool
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

def column_parallel_linear_forward_cyclic_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    stream_pool: StreamPool,
) -> torch.Tensor:
    d = dist.get_world_size()
    my_rank = dist.get_rank()
    out = torch.empty(
        inp_shard.size(0) * d,
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    out_chunks = out.chunk(d)
    main_stream = torch.cuda.current_stream()
    worker_streams = []
    buffers = [inp_shard] + [torch.empty_like(inp_shard) for _ in range(d - 1)]
    for i in range(d):


    torch.manual_seed(0)
    inp_ref = torch.randn(t, h * 4, dtype=torch.bfloat16)
    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16)

    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(-1)]).to_local()
    weight_shard = distribute_tensor(
        weight_ref, mesh, placements=[Shard(-1)]
    ).to_local()

    stream_pool = StreamPool()
    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: row_parallel_linear_forward_reference(
            inp_shard, weight_shard, s, stream_pool
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

def column_parallel_linear_forward_cyclic_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    stream_pool: StreamPool,
) -> torch.Tensor:
    d = dist.get_world_size()
    my_rank = dist.get_rank()
    out = torch.empty(
        inp_shard.size(0) * d,
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    out_chunks = out.chunk(d)
    main_stream = torch.cuda.current_stream()
    worker_streams = []
    buffers = [inp_shard] + [torch.empty_like(inp_shard) for _ in range(d - 1)]
    for i in range(d):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            if i > 0:
                send_op = dist.P2POp(dist.isend, buffers[i - 1], (my_rank + 1) % d)
                recv_op = dist.P2POp(dist.irecv, buffers[i], (my_rank - 1 + d) % d)
                reqs = dist.batch_isend_irecv([send_op, recv_op])
                for req in reqs:
                    req.wait()
            torch.matmul(
                buffers[i],
                weight_shard.T,
                out=out_chunks[(my_rank - i + d) % d],
            )
    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)
    return out

def column_parallel_linear_forward_uncyclic_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    stream_pool: StreamPool,
) -> torch.Tensor:
    d = dist.get_world_size()
    my_rank = dist.get_rank()
    out = torch.empty(
        inp_shard.size(0) * d,
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    out_chunks = out.chunk(d)
    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i in range(d):
        worker_stream = stream_pool.get(i)


        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            if i > 0:
                send_op = dist.P2POp(dist.isend, buffers[i - 1], (my_rank + 1) % d)
                recv_op = dist.P2POp(dist.irecv, buffers[i], (my_rank - 1 + d) % d)
                reqs = dist.batch_isend_irecv([send_op, recv_op])
                for req in reqs:
                    req.wait()
            torch.matmul(
                buffers[i],
                weight_shard.T,
                out=out_chunks[(my_rank - i + d) % d],
            )
    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)
    return out

def column_parallel_linear_forward_uncyclic_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    stream_pool: StreamPool,
) -> torch.Tensor:
    d = dist.get_world_size()
    my_rank = dist.get_rank()
    out = torch.empty(
        inp_shard.size(0) * d,
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    out_chunks = out.chunk(d)
    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i in range(d):
        worker_stream = stream_pool.get(i)
        worker_streams.append(worker_stream)
        worker_stream.wait_stream(main_stream)
        with torch.cuda.stream(worker_stream):
            if i == 0:
                buffer = inp_shard
            else:
                buffer = torch.empty_like(inp_shard)
                send_op = dist.P2POp(dist.isend, inp_shard, (my_rank - i + d) % d)
                recv_op = dist.P2POp(dist.irecv, buffer, (my_rank + i) % d)
                reqs = dist.batch_isend_irecv([send_op, recv_op])
                for req in reqs:
                    req.wait()
            torch.matmul(
                buffer,
                weight_shard.T,
                out=out_chunks[(my_rank + i) % d],
            )
    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)
    return out

@pytest.mark.mpi
@pytest.mark.parametrize(
    "column_parallel_linear_fn",
    [
        column_parallel_linear_forward_cyclic_reference,
        column_parallel_linear_forward_uncyclic_reference,
    ],
    ids=["cyclic", "uncyclic"],
)
def test_column_parallel_linear_forward_reference(
    setup_default_process_group, column_parallel_linear_fn
):
    h, t = 6, 24
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    if 4 * h % d != 0:


                    req.wait()
            torch.matmul(
                buffer,
                weight_shard.T,
                out=out_chunks[(my_rank + i) % d],
            )
    for worker_stream in worker_streams:
        main_stream.wait_stream(worker_stream)
    return out

@pytest.mark.mpi
@pytest.mark.parametrize(
    "column_parallel_linear_fn",
    [
        column_parallel_linear_forward_cyclic_reference,
        column_parallel_linear_forward_uncyclic_reference,
    ],
    ids=["cyclic", "uncyclic"],
)
def test_column_parallel_linear_forward_reference(
    setup_default_process_group, column_parallel_linear_fn
):
    h, t = 6, 24
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])

    if 4 * h % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {4 * h} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )

    inp_ref = torch.testing.make_tensor(t, h, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(0)]).to_local()
    weight_ref = torch.testing.make_tensor(
        4 * h, h, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    weight_shard = distribute_tensor(weight_ref, mesh, placements=[Shard(0)]).to_local()
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_shard)
    stream_pool = StreamPool()
    out = column_parallel_linear_fn(inp_shard, weight_shard, stream_pool)
    torch.testing.assert_close(out, out_ref)

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize(
    "column_parallel_linear_fn",
    [
        column_parallel_linear_forward_cyclic_reference,
        column_parallel_linear_forward_uncyclic_reference,
    ],
    ids=["cyclic", "uncyclic"],
)
def test_column_parallel_linear_forward_reference_benchmark(
    benchmark, setup_default_process_group, column_parallel_linear_fn
):
    h, t = 8192, 8192
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_ref = torch.randn(t, h, dtype=torch.bfloat16, device="cpu")


    weight_ref = torch.testing.make_tensor(
        4 * h, h, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)
    weight_shard = distribute_tensor(weight_ref, mesh, placements=[Shard(0)]).to_local()
    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight_shard)
    stream_pool = StreamPool()
    out = column_parallel_linear_fn(inp_shard, weight_shard, stream_pool)
    torch.testing.assert_close(out, out_ref)

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize(
    "column_parallel_linear_fn",
    [
        column_parallel_linear_forward_cyclic_reference,
        column_parallel_linear_forward_uncyclic_reference,
    ],
    ids=["cyclic", "uncyclic"],
)
def test_column_parallel_linear_forward_reference_benchmark(
    benchmark, setup_default_process_group, column_parallel_linear_fn
):
    h, t = 8192, 8192
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_ref = torch.randn(t, h, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(4 * h, h, dtype=torch.bfloat16, device="cpu")
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(0)]).to_local()
    weight_shard = distribute_tensor(weight_ref, mesh, placements=[Shard(0)]).to_local()
    stream_pool = StreamPool()
    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: column_parallel_linear_fn(inp_shard, weight_shard, stream_pool)
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

def column_parallel_linear_forward(h: int, d: int, parallelism: str):
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, h), contiguity=True, dtype=DataType.BFloat16)
        weight_tv = fd.define_tensor(
            (4 * h, h), contiguity=True, dtype=DataType.BFloat16
        )
        ag_out = fd.ops.set(inp_tv)
        out_tv = fd.ops.linear(ag_out, weight_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

        for tv in [inp_tv, weight_tv]:
            tv.set_device_mesh(mesh)
            tv.outer_split(0, d)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        ag_out.set_device_mesh(mesh)


)
def test_column_parallel_linear_forward_reference_benchmark(
    benchmark, setup_default_process_group, column_parallel_linear_fn
):
    h, t = 8192, 8192
    d = dist.get_world_size()
    torch.manual_seed(0)
    mesh = dist.device_mesh.init_device_mesh("cuda", [d])
    inp_ref = torch.randn(t, h, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(4 * h, h, dtype=torch.bfloat16, device="cpu")
    inp_shard = distribute_tensor(inp_ref, mesh, placements=[Shard(0)]).to_local()
    weight_shard = distribute_tensor(weight_ref, mesh, placements=[Shard(0)]).to_local()
    stream_pool = StreamPool()
    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: column_parallel_linear_fn(inp_shard, weight_shard, stream_pool)
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

def column_parallel_linear_forward(h: int, d: int, parallelism: str):
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, h), contiguity=True, dtype=DataType.BFloat16)
        weight_tv = fd.define_tensor(
            (4 * h, h), contiguity=True, dtype=DataType.BFloat16
        )
        ag_out = fd.ops.set(inp_tv)
        out_tv = fd.ops.linear(ag_out, weight_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

        for tv in [inp_tv, weight_tv]:
            tv.set_device_mesh(mesh)
            tv.outer_split(0, d)
            tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        ag_out.set_device_mesh(mesh)
        ag_out.outer_split(0, d)
        if parallelism == "collective_permute":
            ag_out.swizzle1d(0, nvfuser.ParallelType.mesh_x)
        ag_out.axis(0).parallelize(nvfuser.ParallelType.stream)

        # Fusion IR before segmentation will look like this:
        #   [t, h]
        #   /\.
        #  d
        # (deviceIdx.x)
        #    |
        #    | set (lowered to Broadcast/CollectivePermute. This decomposition is done
        #    |    manually in the definition above. It will later be done
        #    |    by preseg.)
        #    |
        #   [t, h]                                  [4h,  h]
        #   /\                                      /\.
        #  d                                       d
        #  | swizzle1d (if parallelism == "collective_permute")
        #  s
        # (streamIdx)
        #                      |
        #                      | linear
        #                      |
        #                   [t, 4h, r{h}]
        #                   /\  /\.
        #                  s*   d

    return fd

@pytest.mark.mpi
@pytest.mark.parametrize("parallelism", ["collective_permute", "broadcast"])
def test_column_parallel_linear_forward(multidevice_test, parallelism: str):
    # This is a port of CollectiveBasedOverlapTest.ColumnAndSequenceParallelLinear_Forward.
    # The difference is we are using broadcast based overlapping instead of send/recv.
    h, t = 2, 24
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )

    fd = column_parallel_linear_forward(h, d, parallelism)


        #    |    by preseg.)
        #    |
        #   [t, h]                                  [4h,  h]
        #   /\                                      /\.
        #  d                                       d
        #  | swizzle1d (if parallelism == "collective_permute")
        #  s
        # (streamIdx)
        #                      |
        #                      | linear
        #                      |
        #                   [t, 4h, r{h}]
        #                   /\  /\.
        #                  s*   d

    return fd

@pytest.mark.mpi
@pytest.mark.parametrize("parallelism", ["collective_permute", "broadcast"])
def test_column_parallel_linear_forward(multidevice_test, parallelism: str):
    # This is a port of CollectiveBasedOverlapTest.ColumnAndSequenceParallelLinear_Forward.
    # The difference is we are using broadcast based overlapping instead of send/recv.
    h, t = 2, 24
    d = multidevice_test.size
    if (h * 4) % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {h * 4} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )

    fd = column_parallel_linear_forward(h, d, parallelism)

    inp_ref = torch.testing.make_tensor(t, h, dtype=torch.int32, device="cpu").to(
        torch.bfloat16
    )
    weight_ref = torch.testing.make_tensor(
        4 * h, h, dtype=torch.int32, device="cpu"
    ).to(torch.bfloat16)

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    out_ref = torch.nn.functional.linear(inp_ref.cuda(), weight)

    with torch.profiler.profile(record_shapes=True) as prof:
        (out,) = fd.execute([inp, weight], _enable_options=["host_ir_lowering"])
    torch.testing.assert_close(out, out_ref)

    if parallelism == "collective_permute":
        collective_permute_events = [
            event for event in prof.events() if "ncclDevKernel_SendRecv" in event.name
        ]
        assert len(collective_permute_events) == (d - 1)
    else:
        broadcast_events = [
            event for event in prof.events() if "ncclDevKernel_Broadcast" in event.name
        ]
        assert len(broadcast_events) == (d if d > 1 else 0)

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize("parallelism", ["collective_permute", "broadcast"])
def test_column_parallel_linear_forward_benchmark(
    multidevice_test, benchmark, parallelism: str
):
    # This is a port of CollectiveBasedOverlapTest.ColumnParallelLinear_Forward.
    h, t = 8192, 8192
    d = multidevice_test.size
    if (4 * h) % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {4 * h} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )


    with torch.profiler.profile(record_shapes=True) as prof:
        (out,) = fd.execute([inp, weight], _enable_options=["host_ir_lowering"])
    torch.testing.assert_close(out, out_ref)

    if parallelism == "collective_permute":
        collective_permute_events = [
            event for event in prof.events() if "ncclDevKernel_SendRecv" in event.name
        ]
        assert len(collective_permute_events) == (d - 1)
    else:
        broadcast_events = [
            event for event in prof.events() if "ncclDevKernel_Broadcast" in event.name
        ]
        assert len(broadcast_events) == (d if d > 1 else 0)

@pytest.mark.mpi
@pytest.mark.benchmark
@pytest.mark.parametrize("parallelism", ["collective_permute", "broadcast"])
def test_column_parallel_linear_forward_benchmark(
    multidevice_test, benchmark, parallelism: str
):
    # This is a port of CollectiveBasedOverlapTest.ColumnParallelLinear_Forward.
    h, t = 8192, 8192
    d = multidevice_test.size
    if (4 * h) % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {4 * h} to be divisible by world size {d}."
        )
    if t % d != 0:
        pytest.skip(
            f"Column-parallel linear requires {t} to be divisible by world size {d}."
        )

    fd = column_parallel_linear_forward(h, d, parallelism)

    inp_ref = torch.randn(t, h, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(4 * h, h, dtype=torch.bfloat16, device="cpu")

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

@pytest.mark.mpi
@pytest.mark.parametrize("backend_type", [CommunicatorBackend.nccl])
@pytest.mark.parametrize("s", [1, 8])
def test_overlap_allgather_matmul_stream_outermost(
    multidevice_test, benchmark, backend_type, s
):
    def fusion_definition(fd, m, k, n, s, d) -> list[TensorView]:
        x = fd.define_tensor(
            shape=[s, d, m // (s * d), k], contiguity=True, dtype=DataType.BFloat16
        )
        weight = fd.define_tensor(
            shape=[n, k], contiguity=True, dtype=DataType.BFloat16
        )
        bias = fd.define_tensor(shape=[n], contiguity=True, dtype=DataType.BFloat16)

        # [s, d, m//(s*d), n]
        out = fd.ops.linear(x, weight, bias)


    inp_ref = torch.randn(t, h, dtype=torch.bfloat16, device="cpu")
    weight_ref = torch.randn(4 * h, h, dtype=torch.bfloat16, device="cpu")

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

@pytest.mark.mpi
@pytest.mark.parametrize("backend_type", [CommunicatorBackend.nccl])
@pytest.mark.parametrize("s", [1, 8])
def test_overlap_allgather_matmul_stream_outermost(
    multidevice_test, benchmark, backend_type, s
):
    def fusion_definition(fd, m, k, n, s, d) -> list[TensorView]:
        x = fd.define_tensor(
            shape=[s, d, m // (s * d), k], contiguity=True, dtype=DataType.BFloat16
        )
        weight = fd.define_tensor(
            shape=[n, k], contiguity=True, dtype=DataType.BFloat16
        )
        bias = fd.define_tensor(shape=[n], contiguity=True, dtype=DataType.BFloat16)

        # [s, d, m//(s*d), n]
        out = fd.ops.linear(x, weight, bias)

        fd.add_output(out)
        return [x, weight, bias, out]

    def multidevice_schedule(fd, tensors, num_devices) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        x, weight, bias, out = tensors
        x.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
        out.axis(0).parallelize(nvfuser.ParallelType.stream)

    N_WARMUPS, N_ITERATIONS = 5, 25
    m, k, n, d = 2**10, 2**10, 2**10, multidevice_test.size
    assert m % (s * d) == 0

    os.environ["UCC_CL_BASIC_TLS"] = "nccl"
    x_unsharded = torch.testing.make_tensor(
        s, d, m // (s * d), k, dtype=torch.bfloat16, device="cpu"
    )
    x = multidevice_test.shard_tensor_1d(
        x_unsharded,
        1,
        nvfuser.multidevice.DeviceMesh(range(multidevice_test.size)),
    )
    weight = torch.testing.make_tensor(n, k, dtype=torch.bfloat16, device="cuda")
    bias = torch.testing.make_tensor(n, dtype=torch.bfloat16, device="cuda")
    ins = [x, weight, bias]
    out_ref = torch.nn.functional.linear(x_unsharded, weight.cpu(), bias.cpu())

    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, s, d)
        multidevice_schedule(fd, tensors, d)

    params = nvfuser.multidevice.MultiDeviceExecutorParams()
    params.backend_type = backend_type
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    # warmup
    for _ in range(N_WARMUPS):
        outputs = multidevice_executor.run(ins)
        out = outputs[0].cpu()
        assert out.dtype == torch.bfloat16
        assert out.shape == torch.Size([s, d, m // (s * d), n])
        torch.testing.assert_close(out, out_ref, rtol=1e-1, atol=1e-1)

    # benchmark
    benchmark.pedantic(lambda: multidevice_executor.run(ins), rounds=N_ITERATIONS)

@pytest.mark.mpi
@pytest.mark.parametrize(
    "backend_type", [CommunicatorBackend.nccl, CommunicatorBackend.cuda]
)
def test_overlap_allgather_matmul_shard_outermost(
    multidevice_test, benchmark, backend_type
):
    def fusion_definition(fd, m, k, n, d) -> list[TensorView]:
        x = fd.define_tensor(
            shape=[d, m // d, k], contiguity=True, dtype=DataType.BFloat16
        )
        weight = fd.define_tensor(
            shape=[n, k], contiguity=True, dtype=DataType.BFloat16
        )
        bias = fd.define_tensor(shape=[n], contiguity=True, dtype=DataType.BFloat16)

        # [d, m//d, n]
        out = fd.ops.linear(x, weight, bias)


    params = nvfuser.multidevice.MultiDeviceExecutorParams()
    params.backend_type = backend_type
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    # warmup
    for _ in range(N_WARMUPS):
        outputs = multidevice_executor.run(ins)
        out = outputs[0].cpu()
        assert out.dtype == torch.bfloat16
        assert out.shape == torch.Size([s, d, m // (s * d), n])
        torch.testing.assert_close(out, out_ref, rtol=1e-1, atol=1e-1)

    # benchmark
    benchmark.pedantic(lambda: multidevice_executor.run(ins), rounds=N_ITERATIONS)

@pytest.mark.mpi
@pytest.mark.parametrize(
    "backend_type", [CommunicatorBackend.nccl, CommunicatorBackend.cuda]
)
def test_overlap_allgather_matmul_shard_outermost(
    multidevice_test, benchmark, backend_type
):
    def fusion_definition(fd, m, k, n, d) -> list[TensorView]:
        x = fd.define_tensor(
            shape=[d, m // d, k], contiguity=True, dtype=DataType.BFloat16
        )
        weight = fd.define_tensor(
            shape=[n, k], contiguity=True, dtype=DataType.BFloat16
        )
        bias = fd.define_tensor(shape=[n], contiguity=True, dtype=DataType.BFloat16)

        # [d, m//d, n]
        out = fd.ops.linear(x, weight, bias)

        fd.add_output(out)
        return [x, weight, bias, out]

    def multidevice_schedule(fd, tensors, num_devices) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        x, weight, bias, out = tensors
        x.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
        out.axis(0).parallelize(nvfuser.ParallelType.stream)

    N_WARMUPS, N_ITERATIONS = 5, 25
    m, k, n, d = 2**10, 2**10, 2**10, multidevice_test.size
    assert m % d == 0

    os.environ["UCC_CL_BASIC_TLS"] = "nccl"
    x_unsharded = torch.testing.make_tensor(
        d, m // d, k, dtype=torch.bfloat16, device="cpu"
    )
    x = multidevice_test.shard_tensor_1d(
        x_unsharded,
        0,
        nvfuser.multidevice.DeviceMesh(range(multidevice_test.size)),
    )
    weight = torch.testing.make_tensor(n, k, dtype=torch.bfloat16, device="cuda")
    bias = torch.testing.make_tensor(n, dtype=torch.bfloat16, device="cuda")
    ins = [x, weight, bias]
    out_ref = torch.nn.functional.linear(x_unsharded, weight.cpu(), bias.cpu())

    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, d)
        multidevice_schedule(fd, tensors, d)

    params = nvfuser.multidevice.MultiDeviceExecutorParams()
    params.backend_type = backend_type
    params.use_allocation_cache = True
    params.offset_stream_indexing_by_rank = True
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    # warmup
    for _ in range(N_WARMUPS):
        outputs = multidevice_executor.run(ins)
        out = outputs[0].cpu()
        assert out.dtype == torch.bfloat16
        assert out.shape == torch.Size([d, m // d, n])
        torch.testing.assert_close(out, out_ref, rtol=1e-1, atol=1e-1)

    # benchmark
    benchmark.pedantic(lambda: multidevice_executor.run(ins), rounds=N_ITERATIONS)

@pytest.mark.mpi
@pytest.mark.parametrize(
    "backend_type", [CommunicatorBackend.nccl, CommunicatorBackend.cuda]
)
def test_allgather_matmul_no_overlap(multidevice_test, backend_type):
    def fusion_definition(
        fd: FusionDefinition, m: int, k: int, n: int, d: int
    ) -> list[TensorView]:
        a = fd.define_tensor(
            shape=[d, m // d, k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        b = fd.define_tensor(
            shape=[k, n],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )


    params.use_allocation_cache = True
    params.offset_stream_indexing_by_rank = True
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    # warmup
    for _ in range(N_WARMUPS):
        outputs = multidevice_executor.run(ins)
        out = outputs[0].cpu()
        assert out.dtype == torch.bfloat16
        assert out.shape == torch.Size([d, m // d, n])
        torch.testing.assert_close(out, out_ref, rtol=1e-1, atol=1e-1)

    # benchmark
    benchmark.pedantic(lambda: multidevice_executor.run(ins), rounds=N_ITERATIONS)

@pytest.mark.mpi
@pytest.mark.parametrize(
    "backend_type", [CommunicatorBackend.nccl, CommunicatorBackend.cuda]
)
def test_allgather_matmul_no_overlap(multidevice_test, backend_type):
    def fusion_definition(
        fd: FusionDefinition, m: int, k: int, n: int, d: int
    ) -> list[TensorView]:
        a = fd.define_tensor(
            shape=[d, m // d, k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        b = fd.define_tensor(
            shape=[k, n],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )

        c = fd.ops.matmul(a, b)
        if backend_type == CommunicatorBackend.cuda:
            c.set_memory_type(nvfuser.MemoryType.symmetric)
        fd.add_output(c)
        return [a, b, c]

    def multidevice_schedule(tensors: list[TensorView], num_devices: int) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        a, _, _ = tensors
        a.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    d = multidevice_test.size
    m, k, n = 128, 64, 96
    if m % d != 0:
        pytest.skip(f"m ({m}) must be divisible by world size ({d}).")

    dtype = torch.bfloat16
    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, d)
        multidevice_schedule(tensors, d)

    a_ref = torch.testing.make_tensor(d, m // d, k, dtype=torch.int32, device="cpu").to(
        dtype
    )
    b_ref = torch.testing.make_tensor(k, n, dtype=torch.int32, device="cpu").to(dtype)

    a = multidevice_test.shard_tensor(a_ref, fd.fusion.inputs()[0])
    b = b_ref.cuda(multidevice_test.local_rank)

    params = nvfuser.multidevice.MultiDeviceExecutorParams()
    params.backend_type = backend_type
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    # Set a new CUDA stream because cudaMemcpyBatchAsync does not support default stream
    with torch.cuda.Stream(device=multidevice_test.local_rank):
        (out,) = multidevice_executor.run([a, b])
        out_ref = torch.matmul(a_ref, b_ref)
        torch.testing.assert_close(out.cpu(), out_ref)


    weight_ref = torch.randn(h, h * 4, dtype=torch.bfloat16, device="cpu")

    inp = multidevice_test.shard_tensor(inp_ref, fd.fusion.inputs()[0])
    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

# The caching allocator in PyTorch can't cache buffers across streams, so we
# have to reuse streams to avoid repeated cudaMalloc. torch.cuda.Stream() is
# backed by a stream pool as well but I failed to find a way to set its size.
class StreamPool:
    def __init__(self):
        self._streams = {}

    def get(self, sid: int) -> torch.cuda.Stream:
        s = self._streams.get(sid)
        if s is None:
            s = torch.cuda.Stream()
            self._streams[sid] = s
        return s

def row_parallel_linear_forward_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    num_chunks: int,
    stream_pool: StreamPool,
) -> torch.Tensor:
    out = torch.empty(
        inp_shard.size(0),
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,


get

    weight = multidevice_test.shard_tensor(weight_ref, fd.fusion.inputs()[1])

    warmup_fn, benchmark_fn = get_benchmark_fns(
        lambda: fd.execute(
            [inp, weight],
            _enable_options=["host_ir_lowering"],
            _disable_options=["infer_contiguity"],
        )
    )
    warmup_fn()
    benchmark.pedantic(benchmark_fn, rounds=5)

# The caching allocator in PyTorch can't cache buffers across streams, so we
# have to reuse streams to avoid repeated cudaMalloc. torch.cuda.Stream() is
# backed by a stream pool as well but I failed to find a way to set its size.
class StreamPool:
    def __init__(self):
        self._streams = {}

    def get(self, sid: int) -> torch.cuda.Stream:
        s = self._streams.get(sid)
        if s is None:
            s = torch.cuda.Stream()
            self._streams[sid] = s
        return s

def row_parallel_linear_forward_reference(
    inp_shard: torch.Tensor,
    weight_shard: torch.Tensor,
    num_chunks: int,
    stream_pool: StreamPool,
) -> torch.Tensor:
    out = torch.empty(
        inp_shard.size(0),
        weight_shard.size(0),
        device="cuda",
        dtype=inp_shard.dtype,
    )
    inp_chunks = inp_shard.chunk(num_chunks)
    out_chunks = out.chunk(num_chunks)

    main_stream = torch.cuda.current_stream()
    worker_streams = []
    for i, (inp_chunk, out_chunk) in enumerate(zip(inp_chunks, out_chunks)):


    params.offset_stream_indexing_by_rank = True
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    # warmup
    for _ in range(N_WARMUPS):
        outputs = multidevice_executor.run(ins)
        out = outputs[0].cpu()
        assert out.dtype == torch.bfloat16
        assert out.shape == torch.Size([d, m // d, n])
        torch.testing.assert_close(out, out_ref, rtol=1e-1, atol=1e-1)

    # benchmark
    benchmark.pedantic(lambda: multidevice_executor.run(ins), rounds=N_ITERATIONS)

@pytest.mark.mpi
@pytest.mark.parametrize(
    "backend_type", [CommunicatorBackend.nccl, CommunicatorBackend.cuda]
)
def test_allgather_matmul_no_overlap(multidevice_test, backend_type):
    def fusion_definition(
        fd: FusionDefinition, m: int, k: int, n: int, d: int
    ) -> list[TensorView]:
        a = fd.define_tensor(
            shape=[d, m // d, k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        b = fd.define_tensor(
            shape=[k, n],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )

        c = fd.ops.matmul(a, b)
        if backend_type == CommunicatorBackend.cuda:
            c.set_memory_type(nvfuser.MemoryType.symmetric)
        fd.add_output(c)
        return [a, b, c]

    def multidevice_schedule(tensors: list[TensorView], num_devices: int) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        a, _, _ = tensors
        a.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    d = multidevice_test.size
    m, k, n = 128, 64, 96
    if m % d != 0:
        pytest.skip(f"m ({m}) must be divisible by world size ({d}).")

    dtype = torch.bfloat16
    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, d)
        multidevice_schedule(tensors, d)

    a_ref = torch.testing.make_tensor(d, m // d, k, dtype=torch.int32, device="cpu").to(


    def fusion_definition(
        fd: FusionDefinition, m: int, k: int, n: int, d: int
    ) -> list[TensorView]:
        a = fd.define_tensor(
            shape=[d, m // d, k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        b = fd.define_tensor(
            shape=[k, n],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )

        c = fd.ops.matmul(a, b)
        if backend_type == CommunicatorBackend.cuda:
            c.set_memory_type(nvfuser.MemoryType.symmetric)
        fd.add_output(c)
        return [a, b, c]

    def multidevice_schedule(tensors: list[TensorView], num_devices: int) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        a, _, _ = tensors
        a.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    d = multidevice_test.size
    m, k, n = 128, 64, 96
    if m % d != 0:
        pytest.skip(f"m ({m}) must be divisible by world size ({d}).")

    dtype = torch.bfloat16
    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, d)
        multidevice_schedule(tensors, d)

    a_ref = torch.testing.make_tensor(d, m // d, k, dtype=torch.int32, device="cpu").to(
        dtype
    )
    b_ref = torch.testing.make_tensor(k, n, dtype=torch.int32, device="cpu").to(dtype)

    a = multidevice_test.shard_tensor(a_ref, fd.fusion.inputs()[0])
    b = b_ref.cuda(multidevice_test.local_rank)

    params = nvfuser.multidevice.MultiDeviceExecutorParams()


tests.python.multidevice.linear

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

from dataclasses import dataclass
from enum import auto, Enum
from functools import lru_cache
from typing import Iterable

import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement

from nvfuser_direct import DataType, FusionDefinition
from .fusion_definition_wrapper import FusionDefinitionWrapper

@dataclass(frozen=True)
class LinearConfig:
    in_features: int
    out_features: int
    # Whether the input/output tensors have a leading batch dimension. This is
    # typically true for MLA and false for MoE.
    has_batch: bool

# I omitted biases because DeepSeek V3 uses non-biased linear layers in MLA and
# MoE.
def define_linear_forward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    inp_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    inp = fd.define_tensor(inp_shape, contiguity=True, dtype=DataType.BFloat16)
    weight = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    out = fd.ops.linear(inp, weight)
    fd.add_output(out)

def define_linear_backward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    x_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    x = fd.define_tensor(x_shape, contiguity=True, dtype=DataType.BFloat16)
    w = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    grad_shape = [-1, -1, e_out] if config.has_batch else [-1, e_out]
    grad = fd.define_tensor(grad_shape, contiguity=True, dtype=DataType.BFloat16)

    grad_x = fd.ops.matmul(grad, w)

    grad_flat = fd.ops.reshape(grad, [-1, e_out]) if config.has_batch else grad
    grad_flat_t = fd.ops.permute(grad_flat, [1, 0])

    x_flat = fd.ops.reshape(x, [-1, e_in]) if config.has_batch else x

    grad_w = fd.ops.matmul(grad_flat_t, x_flat)

    fd.add_output(grad_x)
    fd.add_output(grad_w)

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()

# Cache based on forward and backward linear definition and its configuration.
# FusionDefinitionWrapper caches based on input dtensors.
@lru_cache
def get_fusion_definition_wrapper(
    compute_type: ComputeType,
    linear_config: LinearConfig,
) -> FusionDefinition:
    match compute_type:
        case ComputeType.FORWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_forward(linear_config, fd)
            )
        case ComputeType.BACKWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_backward(linear_config, fd)
            )

class LinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        input: DTensor,
        weight: DTensor,
    ):
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.FORWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        (output,) = op([input, weight])
        ctx.save_for_backward(input, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output: DTensor):
        input, weight = ctx.saved_tensors
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.BACKWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        grad_x, grad_w = op([input, weight, grad_output])
        return (grad_x, grad_w)

class TensorParallelLinear(torch.nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        in_placements: Iterable[Placement] = [],
    ):
        # Unlike normal layers whose `__init__` allocates the parameters,
        # `TensorParallelLinear` is expected to be created from a
        # non-distributed Linear layer via the `distribute` method. Therefore,
        # here, we construct super() with no memory allocated for weights. The
        # weights will be derived by the `distribute` method.
        super().__init__(in_features, out_features, bias=False, device="meta")
        self.in_placements = in_placements

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, in_placements={self.in_placements}, weight_placements={self.weight.data.placements})"

    @classmethod
    def distribute(
        cls,
        linear: torch.nn.Linear,
        mesh: dist.device_mesh.DeviceMesh,
        in_placements: Iterable[Placement] = [],
        weight_placements: Iterable[Placement] = [],
    ):
        tp_linear = cls(linear.in_features, linear.out_features, in_placements)
        tp_linear.weight = torch.nn.Parameter(
            dist.tensor.distribute_tensor(linear.weight, mesh, weight_placements)
        )
        return tp_linear

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        in_dtensor = DTensor.from_local(
            input, self.weight.device_mesh, self.in_placements
        )
        out_dtensor = LinearFunction.apply(in_dtensor, self.weight)
        return out_dtensor.to_local()


LinearConfig

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

from dataclasses import dataclass
from enum import auto, Enum
from functools import lru_cache
from typing import Iterable

import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement

from nvfuser_direct import DataType, FusionDefinition
from .fusion_definition_wrapper import FusionDefinitionWrapper

@dataclass(frozen=True)
class LinearConfig:
    in_features: int
    out_features: int
    # Whether the input/output tensors have a leading batch dimension. This is
    # typically true for MLA and false for MoE.
    has_batch: bool

# I omitted biases because DeepSeek V3 uses non-biased linear layers in MLA and
# MoE.
def define_linear_forward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    inp_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    inp = fd.define_tensor(inp_shape, contiguity=True, dtype=DataType.BFloat16)
    weight = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    out = fd.ops.linear(inp, weight)
    fd.add_output(out)

def define_linear_backward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    x_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    x = fd.define_tensor(x_shape, contiguity=True, dtype=DataType.BFloat16)
    w = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)


import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement

from nvfuser_direct import DataType, FusionDefinition
from .fusion_definition_wrapper import FusionDefinitionWrapper

@dataclass(frozen=True)
class LinearConfig:
    in_features: int
    out_features: int
    # Whether the input/output tensors have a leading batch dimension. This is
    # typically true for MLA and false for MoE.
    has_batch: bool

# I omitted biases because DeepSeek V3 uses non-biased linear layers in MLA and
# MoE.
def define_linear_forward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    inp_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    inp = fd.define_tensor(inp_shape, contiguity=True, dtype=DataType.BFloat16)
    weight = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    out = fd.ops.linear(inp, weight)
    fd.add_output(out)

def define_linear_backward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    x_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    x = fd.define_tensor(x_shape, contiguity=True, dtype=DataType.BFloat16)
    w = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    grad_shape = [-1, -1, e_out] if config.has_batch else [-1, e_out]
    grad = fd.define_tensor(grad_shape, contiguity=True, dtype=DataType.BFloat16)

    grad_x = fd.ops.matmul(grad, w)

    grad_flat = fd.ops.reshape(grad, [-1, e_out]) if config.has_batch else grad
    grad_flat_t = fd.ops.permute(grad_flat, [1, 0])

    x_flat = fd.ops.reshape(x, [-1, e_in]) if config.has_batch else x

    grad_w = fd.ops.matmul(grad_flat_t, x_flat)


class LinearConfig:
    in_features: int
    out_features: int
    # Whether the input/output tensors have a leading batch dimension. This is
    # typically true for MLA and false for MoE.
    has_batch: bool

# I omitted biases because DeepSeek V3 uses non-biased linear layers in MLA and
# MoE.
def define_linear_forward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    inp_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    inp = fd.define_tensor(inp_shape, contiguity=True, dtype=DataType.BFloat16)
    weight = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    out = fd.ops.linear(inp, weight)
    fd.add_output(out)

def define_linear_backward(config: LinearConfig, fd: FusionDefinition) -> None:
    e_in, e_out = config.in_features, config.out_features

    x_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    x = fd.define_tensor(x_shape, contiguity=True, dtype=DataType.BFloat16)
    w = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    grad_shape = [-1, -1, e_out] if config.has_batch else [-1, e_out]
    grad = fd.define_tensor(grad_shape, contiguity=True, dtype=DataType.BFloat16)

    grad_x = fd.ops.matmul(grad, w)

    grad_flat = fd.ops.reshape(grad, [-1, e_out]) if config.has_batch else grad
    grad_flat_t = fd.ops.permute(grad_flat, [1, 0])

    x_flat = fd.ops.reshape(x, [-1, e_in]) if config.has_batch else x

    grad_w = fd.ops.matmul(grad_flat_t, x_flat)

    fd.add_output(grad_x)
    fd.add_output(grad_w)

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()

# Cache based on forward and backward linear definition and its configuration.
# FusionDefinitionWrapper caches based on input dtensors.
@lru_cache
def get_fusion_definition_wrapper(
    compute_type: ComputeType,
    linear_config: LinearConfig,
) -> FusionDefinition:
    match compute_type:
        case ComputeType.FORWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_forward(linear_config, fd)
            )
        case ComputeType.BACKWARD:


ComputeType

    x_shape = [-1, -1, e_in] if config.has_batch else [-1, e_in]
    x = fd.define_tensor(x_shape, contiguity=True, dtype=DataType.BFloat16)
    w = fd.define_tensor([e_out, e_in], contiguity=True, dtype=DataType.BFloat16)
    grad_shape = [-1, -1, e_out] if config.has_batch else [-1, e_out]
    grad = fd.define_tensor(grad_shape, contiguity=True, dtype=DataType.BFloat16)

    grad_x = fd.ops.matmul(grad, w)

    grad_flat = fd.ops.reshape(grad, [-1, e_out]) if config.has_batch else grad
    grad_flat_t = fd.ops.permute(grad_flat, [1, 0])

    x_flat = fd.ops.reshape(x, [-1, e_in]) if config.has_batch else x

    grad_w = fd.ops.matmul(grad_flat_t, x_flat)

    fd.add_output(grad_x)
    fd.add_output(grad_w)

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()

# Cache based on forward and backward linear definition and its configuration.
# FusionDefinitionWrapper caches based on input dtensors.
@lru_cache
def get_fusion_definition_wrapper(
    compute_type: ComputeType,
    linear_config: LinearConfig,
) -> FusionDefinition:
    match compute_type:
        case ComputeType.FORWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_forward(linear_config, fd)
            )
        case ComputeType.BACKWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_backward(linear_config, fd)
            )


    grad_flat = fd.ops.reshape(grad, [-1, e_out]) if config.has_batch else grad
    grad_flat_t = fd.ops.permute(grad_flat, [1, 0])

    x_flat = fd.ops.reshape(x, [-1, e_in]) if config.has_batch else x

    grad_w = fd.ops.matmul(grad_flat_t, x_flat)

    fd.add_output(grad_x)
    fd.add_output(grad_w)

class ComputeType(Enum):
    FORWARD = auto()
    BACKWARD = auto()

# Cache based on forward and backward linear definition and its configuration.
# FusionDefinitionWrapper caches based on input dtensors.
@lru_cache
def get_fusion_definition_wrapper(
    compute_type: ComputeType,
    linear_config: LinearConfig,
) -> FusionDefinition:
    match compute_type:
        case ComputeType.FORWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_forward(linear_config, fd)
            )
        case ComputeType.BACKWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_backward(linear_config, fd)
            )

class LinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        input: DTensor,
        weight: DTensor,
    ):
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.FORWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        (output,) = op([input, weight])
        ctx.save_for_backward(input, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output: DTensor):


LinearFunction

# Cache based on forward and backward linear definition and its configuration.
# FusionDefinitionWrapper caches based on input dtensors.
@lru_cache
def get_fusion_definition_wrapper(
    compute_type: ComputeType,
    linear_config: LinearConfig,
) -> FusionDefinition:
    match compute_type:
        case ComputeType.FORWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_forward(linear_config, fd)
            )
        case ComputeType.BACKWARD:
            return FusionDefinitionWrapper(
                lambda fd: define_linear_backward(linear_config, fd)
            )

class LinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        input: DTensor,
        weight: DTensor,
    ):
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.FORWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        (output,) = op([input, weight])
        ctx.save_for_backward(input, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output: DTensor):
        input, weight = ctx.saved_tensors
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.BACKWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        grad_x, grad_w = op([input, weight, grad_output])
        return (grad_x, grad_w)

class TensorParallelLinear(torch.nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        in_placements: Iterable[Placement] = [],
    ):
        # Unlike normal layers whose `__init__` allocates the parameters,
        # `TensorParallelLinear` is expected to be created from a
        # non-distributed Linear layer via the `distribute` method. Therefore,
        # here, we construct super() with no memory allocated for weights. The
        # weights will be derived by the `distribute` method.
        super().__init__(in_features, out_features, bias=False, device="meta")
        self.in_placements = in_placements

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, in_placements={self.in_placements}, weight_placements={self.weight.data.placements})"


TensorParallelLinear

        op = get_fusion_definition_wrapper(
            ComputeType.FORWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        (output,) = op([input, weight])
        ctx.save_for_backward(input, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output: DTensor):
        input, weight = ctx.saved_tensors
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.BACKWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        grad_x, grad_w = op([input, weight, grad_output])
        return (grad_x, grad_w)

class TensorParallelLinear(torch.nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        in_placements: Iterable[Placement] = [],
    ):
        # Unlike normal layers whose `__init__` allocates the parameters,
        # `TensorParallelLinear` is expected to be created from a
        # non-distributed Linear layer via the `distribute` method. Therefore,
        # here, we construct super() with no memory allocated for weights. The
        # weights will be derived by the `distribute` method.
        super().__init__(in_features, out_features, bias=False, device="meta")
        self.in_placements = in_placements

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, in_placements={self.in_placements}, weight_placements={self.weight.data.placements})"

    @classmethod
    def distribute(
        cls,
        linear: torch.nn.Linear,
        mesh: dist.device_mesh.DeviceMesh,
        in_placements: Iterable[Placement] = [],
        weight_placements: Iterable[Placement] = [],
    ):
        tp_linear = cls(linear.in_features, linear.out_features, in_placements)
        tp_linear.weight = torch.nn.Parameter(
            dist.tensor.distribute_tensor(linear.weight, mesh, weight_placements)
        )
        return tp_linear

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        in_dtensor = DTensor.from_local(
            input, self.weight.device_mesh, self.in_placements
        )
        out_dtensor = LinearFunction.apply(in_dtensor, self.weight)
        return out_dtensor.to_local()


forward

        self.in_placements = in_placements

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, in_placements={self.in_placements}, weight_placements={self.weight.data.placements})"

    @classmethod
    def distribute(
        cls,
        linear: torch.nn.Linear,
        mesh: dist.device_mesh.DeviceMesh,
        in_placements: Iterable[Placement] = [],
        weight_placements: Iterable[Placement] = [],
    ):
        tp_linear = cls(linear.in_features, linear.out_features, in_placements)
        tp_linear.weight = torch.nn.Parameter(
            dist.tensor.distribute_tensor(linear.weight, mesh, weight_placements)
        )
        return tp_linear

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        in_dtensor = DTensor.from_local(
            input, self.weight.device_mesh, self.in_placements
        )
        out_dtensor = LinearFunction.apply(in_dtensor, self.weight)
        return out_dtensor.to_local()


backward

            )

class LinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        input: DTensor,
        weight: DTensor,
    ):
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.FORWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        (output,) = op([input, weight])
        ctx.save_for_backward(input, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output: DTensor):
        input, weight = ctx.saved_tensors
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.BACKWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        grad_x, grad_w = op([input, weight, grad_output])
        return (grad_x, grad_w)

class TensorParallelLinear(torch.nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        in_placements: Iterable[Placement] = [],
    ):
        # Unlike normal layers whose `__init__` allocates the parameters,
        # `TensorParallelLinear` is expected to be created from a
        # non-distributed Linear layer via the `distribute` method. Therefore,
        # here, we construct super() with no memory allocated for weights. The
        # weights will be derived by the `distribute` method.
        super().__init__(in_features, out_features, bias=False, device="meta")
        self.in_placements = in_placements

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, in_placements={self.in_placements}, weight_placements={self.weight.data.placements})"


            ComputeType.FORWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        (output,) = op([input, weight])
        ctx.save_for_backward(input, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output: DTensor):
        input, weight = ctx.saved_tensors
        assert input.dim() in (2, 3)
        op = get_fusion_definition_wrapper(
            ComputeType.BACKWARD,
            LinearConfig(weight.size(1), weight.size(0), has_batch=(input.dim() == 3)),
        )
        grad_x, grad_w = op([input, weight, grad_output])
        return (grad_x, grad_w)

class TensorParallelLinear(torch.nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        in_placements: Iterable[Placement] = [],
    ):
        # Unlike normal layers whose `__init__` allocates the parameters,
        # `TensorParallelLinear` is expected to be created from a
        # non-distributed Linear layer via the `distribute` method. Therefore,
        # here, we construct super() with no memory allocated for weights. The
        # weights will be derived by the `distribute` method.
        super().__init__(in_features, out_features, bias=False, device="meta")
        self.in_placements = in_placements

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, in_placements={self.in_placements}, weight_placements={self.weight.data.placements})"

    @classmethod
    def distribute(
        cls,
        linear: torch.nn.Linear,
        mesh: dist.device_mesh.DeviceMesh,
        in_placements: Iterable[Placement] = [],
        weight_placements: Iterable[Placement] = [],
    ):
        tp_linear = cls(linear.in_features, linear.out_features, in_placements)
        tp_linear.weight = torch.nn.Parameter(
            dist.tensor.distribute_tensor(linear.weight, mesh, weight_placements)
        )
        return tp_linear

    def forward(self, input: torch.Tensor) -> torch.Tensor: