from __future__ import annotations

from tilefoundry.dsl import *

@func(mesh=Mesh((Topology("thread", 8),), layout=(2, 4), names=("warp", "lane")))
def composed_mesh_pipeline(
    x: Tensor[(8, 4, 16), "f32"],
    seed: Tensor[(16,), "f32"],
    acc: Tensor[(8, 16), "f32", (2 @ mesh.warp, 4, 16) @ {mesh.lane @ P("sum")}, "rmem"],
    mixed: Tensor[(8, 16), "f32",
        ShardLayout(
            layout=Layout((4, 2, 16), (0, 16, 1)),
            attrs=(S(0), B(), P("sum")),
            mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), layout=(4, 2, 4), names=("tile", "warp", "lane")),
        ), "rmem"]
):
    with Mesh((Topology("cta", 4),), layout=(4,), names=("tile",)) as cta:  # memory peak=rmem:640B;smem:512B persistent=rmem:128B
        composed = tf.reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage="rmem")  # Tensor[(8, 4, 16), "f32", (4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), "rmem"]; compute-cost precision=exact; memory traffic=gmem:r2.00KiB/w0@logical,r2.00KiB/w0@total,r512B/w0@cta,r64B/w0@thread;rmem:r0/w2.00KiB@logical,r0/w2.00KiB@total,r0/w512B@cta,r0/w64B@thread footprint=x:2.00KiB footprint-precision=exact
        v1 = tf.unary(composed, kind=UnaryKind.SQUARE)  # Tensor[(8, 4, 16), "f32", (4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), "rmem"]; compute-cost flops=f32:512@logical,512@total,128@cta,16@thread precision=exact; memory traffic=rmem:r2.00KiB/w2.00KiB@logical,r2.00KiB/w2.00KiB@total,r512B/w512B@cta,r64B/w64B@thread footprint-precision=exact
        staged = tf.reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage="smem")  # Tensor[(8, 4, 16), "f32", (4 @ cta.tile, 2, 4, 16), "smem"]; compute-cost precision=exact; memory traffic=rmem:r2.00KiB/w0@logical,r2.00KiB/w0@total,r512B/w0@cta,r64B/w0@thread;smem:r0/w2.00KiB@logical,r0/w2.00KiB@total,r0/w512B@cta,r0/w64B@thread footprint-precision=exact
        narrowed = tf.cast(staged, dtype="bf16")  # Tensor[(8, 4, 16), "bf16", (4 @ cta.tile, 2, 4, 16), "smem"]; compute-cost flops=bf16:512@logical,512@total,128@cta,128@thread precision=exact; memory traffic=smem:r2.00KiB/w1.00KiB@logical,r2.00KiB/w1.00KiB@total,r512B/w256B@cta,r512B/w256B@thread footprint-precision=exact
        swapped = tf.transpose(narrowed, perm=(0, 2, 1))  # Tensor[(8, 16, 4), "bf16", (4 @ cta.tile, 2, 16, 4), "smem"]; compute-cost precision=exact; memory traffic=smem:r1.00KiB/w1.00KiB@logical,r1.00KiB/w1.00KiB@total,r256B/w256B@cta,r256B/w256B@thread footprint-precision=exact
        gathered = tf.reshard(swapped, layout=(8, 16, 4), storage="gmem")  # Tensor[(8, 16, 4), "bf16"]; compute-cost precision=exact; memory traffic=gmem:r0/w1.00KiB@logical,r0/w1.00KiB@total,r0/w256B@cta,r0/w256B@thread;smem:r1.00KiB/w0@logical,r1.00KiB/w0@total,r256B/w0@cta,r256B/w0@thread footprint=gathered:1.00KiB footprint-precision=exact
        seeded = tf.reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage="rmem")  # Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]; compute-cost precision=exact; memory traffic=gmem:r64B/w0@logical,r256B/w0@total,r64B/w0@cta,r8B/w0@thread;rmem:r0/w64B@logical,r0/w256B@total,r0/w64B@cta,r0/w8B@thread footprint=seed:64B footprint-precision=exact
    folded = tf.reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage="rmem")  # Tensor[(8, 16), "f32", (2 @ mesh.warp, 4, 4 @ mesh.lane, 4), "rmem"]; compute-cost precision=exact; memory traffic=rmem:r512B/w512B@logical,r512B/w512B@total,r512B/w512B@cta,r64B/w64B@thread footprint-precision=exact
    summed = tf.reshard(mixed, layout=(8, 16), storage="rmem")  # Tensor[(8, 16), "f32", "rmem"]; compute-cost precision=exact; memory traffic=rmem:r512B/w512B@logical,r128B/w128B@total,r128B/w128B@cta,r128B/w128B@thread footprint-precision=exact
    for _ in range(3):
        summed_2 = tf.add(summed, summed)  # Tensor[(8, 16), "f32", "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,128@thread precision=exact; memory traffic=rmem:r1.00KiB/w512B@logical,r1.00KiB/w512B@total,r1.00KiB/w512B@cta,r1.00KiB/w512B@thread footprint-precision=exact
        folded_2 = tf.unary(folded, kind=UnaryKind.SQUARE)  # Tensor[(8, 16), "f32", (2 @ mesh.warp, 4, 4 @ mesh.lane, 4), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,16@thread precision=exact; memory traffic=rmem:r512B/w512B@logical,r512B/w512B@total,r512B/w512B@cta,r64B/w64B@thread footprint-precision=exact
        folded = folded_2
        summed = summed_2
    v18 = tf.reshard(folded, layout=(8, 16), storage="gmem")  # Tensor[(8, 16), "f32"]; compute-cost precision=exact; memory traffic=gmem:r0/w512B@logical,r0/w512B@total,r0/w512B@cta,r0/w64B@thread;rmem:r512B/w0@logical,r512B/w0@total,r512B/w0@cta,r64B/w0@thread footprint=v18:512B footprint-precision=exact
    v21 = tf.reshard(summed, layout=(8, 16), storage="gmem")  # Tensor[(8, 16), "f32"]; compute-cost precision=exact; memory traffic=gmem:r0/w512B@logical,r0/w512B@total,r0/w512B@cta,r0/w512B@thread;rmem:r512B/w0@logical,r512B/w0@total,r512B/w0@cta,r512B/w0@thread footprint=v21:512B footprint-precision=exact
    v24 = tf.reshard(seeded, layout=(16,), storage="gmem")  # Tensor[(16,), "f32"]; compute-cost precision=exact; memory traffic=gmem:r0/w64B@logical,r0/w64B@total,r0/w64B@cta,r0/w8B@thread;rmem:r64B/w0@logical,r64B/w0@total,r64B/w0@cta,r8B/w0@thread footprint=v24:64B footprint-precision=exact
    return (gathered, v18, v21, v24)
