from __future__ import annotations

from tilefoundry.dsl import *
from tilefoundry.target import CudaTarget

@module(entry="composed_mesh_pipeline", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 4), Topology("thread", 8),))
class TypePrinterSugar:
    @func(mesh=Mesh((Topology("thread", 8),), layout=(8,), names=("lane",)))
    def nested_loop_tuple(
        x: Tensor[(8, 16), "f32"],
        weight: Tensor[(8, 16), "f32",
            ShardLayout(
                layout=Layout((8, 16), None),
                attrs=(B(), P("max")),
                mesh=Mesh((Topology("thread", 8),), layout=(2, 4), names=("warp", "lane")),
            )]
    ):
        split = tf.reshard(x, layout=(8 @ mesh.lane, 16), storage="rmem")
        whole = tf.reshard(x, layout=(8, 16), storage="rmem")
        for _ in range(3):
            whole_2 = tf.add(whole, whole)
            split_2 = tf.unary(split, kind=UnaryKind.SQUARE)
            split = split_2
            whole = whole_2
        v7 = tf.reshard(split, layout=(8, 16), storage="gmem")
        v10 = tf.reshard(whole, layout=(8, 16), storage="gmem")
        with Mesh(("thread",), layout=(2, 4), names=("warp", "lane")) as thread:
            per_warp = tf.reshard(weight, layout=(2 @ thread.warp, 4, 16), storage="rmem")
            unfolded = tf.reshard(per_warp, layout=(8, 16), storage="gmem")
        return (v7, v10, unfolded)

    @func(mesh=Mesh((Topology("cta", 4),), layout=(4,), names=("tile",)))
    def named_and_out_of_scope(
        x: Tensor[(8, 16), "f32"],
        held: Tensor[(8, 16), "f32", (4 @ mesh.tile, 2, 16), "rmem"],
        frag: Tensor[(16,), "f32",
            ShardLayout(
                layout=Layout((2, 4, 2), (8, 2, 1)),
                attrs=(S(0), S(1)),
                mesh=Mesh((Topology("thread", 8),), layout=(2, 4), names=("warp", "lane")),
            ), "rmem"]
    ):
        mine = tf.reshard(x, layout=(4 @ mesh.tile, 2, 16), storage="rmem")
        v1 = tf.reshard(mine, layout=(8, 16), storage="gmem")
        escaped = tf.reshard(x, layout=ShardLayout(
            layout=Layout((2, 4, 16), None),
            attrs=(S(0), B()),
            mesh=Mesh((Topology("thread", 8),), layout=(2, 4), names=("warp", "lane")),
        ), storage="rmem")
        return (v1, held, frag, escaped)

    @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), None),
                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(("cta",), layout=(4,), names=("tile",)) as cta:
            composed = tf.reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage="rmem")
            v1 = tf.unary(composed, kind=UnaryKind.SQUARE)
            staged = tf.reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage="smem")
            narrowed = tf.cast(staged, dtype="bf16")
            swapped = tf.transpose(narrowed, perm=(0, 2, 1))
            gathered = tf.reshard(swapped, layout=(8, 16, 4), storage="gmem")
            seeded = tf.reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage="rmem")
        folded = tf.reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage="rmem")
        summed = tf.reshard(mixed, layout=(8, 16), storage="rmem")
        for _ in range(3):
            summed_2 = tf.add(summed, summed)
            folded_2 = tf.unary(folded, kind=UnaryKind.SQUARE)
            folded = folded_2
            summed = summed_2
        v18 = tf.reshard(folded, layout=(8, 16), storage="gmem")
        v21 = tf.reshard(summed, layout=(8, 16), storage="gmem")
        v24 = tf.reshard(seeded, layout=(16,), storage="gmem")
        return (gathered, v18, v21, v24)
