# analysis target=nvidia.h200_sxm module=GEMM_8192X17408X5120_TMA_STORE function=gemm topology=cta wave=1/1
# selection requested=memory executed=memory
# memory traffic=gmem:r522.00MB/w816.00MB@logical,r16.20GB/w816.00MB@total,r16.20GB/w816.00MB@cta,r16.20GB/w816.00MB@thread;rmem:r43.30GB/w42.77GB@logical,r43.31GB/w43.30GB@total,r43.31GB/w43.30GB@cta,r183.88MB/w173.19MB@thread;smem:r16.20GB/w522.00MB@logical,r16.20GB/w16.20GB@total,r16.20GB/w16.20GB@cta,r3.00GB/w15.94GB@thread footprint=a:16.00KB;b:32.00KB;tile_dst:64.00KB;out_3:64.00KB footprint-precision=exact peak=gmem:522.00MB;rmem:128.00KB;smem:208.00KB persistent=gmem:250.00MB
#   buffer=b holds=175.56MB time=m space=none reuse=10.46GB fits=no precision=exact
#   buffer=a holds=3.88MB time=n space=none reuse=83.75MB fits=yes precision=exact
#   error="l2 reuse window m holds 175.56MB at a 1-unit wave, exceeding capacity 47.68MB"

from __future__ import annotations

from tilefoundry.dsl import *

@func(mesh=Mesh((Topology("cta", 1),), layout=(1,), names=("block",)))
def gemm(
    a: Tensor[(8192, 5120), "bf16"],
    b: Tensor[(5120, 17408), "bf16"]
) -> Tensor[(8192, 17408), "bf16"]:
    with Mesh((Topology("thread", 384),), layout=(3, 128), names=("warpgroup", "participant")) as thread:  # memory peak=rmem:128.00KB;smem:208.00KB
        out = tf.zeros(type=Tensor[(8192, 17408), "bf16"])  # Tensor[(8192, 17408), "bf16"]; memory traffic=gmem:r0/w272.00MB@logical,r0/w272.00MB@total,r0/w272.00MB@cta,r0/w272.00MB@thread footprint-precision=exact
        for m in tf.tile(8192, 128):
            for n in tf.tile(17408, 256):
                with Mesh(thread[1:], layout=(2, 4, 8, 4), names=("group", "warp", "lane8", "lane4")) as thread_2:  # memory peak=rmem:128.00KB
                    acc = tf.zeros(type=Tensor[(128, 256), "f32", ((2 @ thread_2.group, 4 @ thread_2.warp, 2, 8 @ thread_2.lane8, 32, 4 @ thread_2.lane4, 2), (16384, 16, 8, 1, 512, 128, 64)), "rmem"])  # Tensor[(128, 256), "f32", ((2 @ thread_2.group, 4 @ thread_2.warp, 2, 8 @ thread_2.lane8, 32, 4 @ thread_2.lane4, 2), (16384, 16, 8, 1, 512, 128, 64)), "rmem"]; memory traffic=rmem:r0/w128.00KB@logical,r0/w128.00KB@total,r0/w128.00KB@cta,r0/w512B@thread footprint-precision=exact
                for k in tf.tile(5120, 64):
                    with thread[:1, :32] as thread_3:  # memory peak=smem:144.00KB
                        v4 = a[m, k]  # Tensor[(128, 64), "bf16", ((128, 64), (5120, 1))]; memory traffic=rmem:r16B/w0@logical,r16B/w0@total,r16B/w0@cta,r16B/w0@thread footprint-precision=exact
                        lhs_ring = tf.empty(type=Tensor[(128, 64), "bf16", Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))) | Swizzle(3, 4, 3), "smem"], buffers=3)  # Tensor[(128, 64), "bf16", Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))) | Swizzle(3, 4, 3), "smem"]; memory footprint-precision=exact
                        lhs = tf.schedule((v4, lhs_ring), op=T.copy_async_tensor())  # Tensor[(128, 64), "bf16", Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))) | Swizzle(3, 4, 3), "smem"]; memory traffic=gmem:r16.00KB/w0@logical,r16.00KB/w0@total,r16.00KB/w0@cta,r16.00KB/w0@thread;smem:r0/w16.00KB@logical,r0/w16.00KB@total,r0/w16.00KB@cta,r0/w16.00KB@thread footprint=a:16.00KB footprint-precision=exact
                        v8 = b[k, n]  # Tensor[(64, 256), "bf16", ((64, 256), (17408, 1))]; memory traffic=rmem:r16B/w0@logical,r16B/w0@total,r16B/w0@cta,r16B/w0@thread footprint-precision=exact
                        rhs_ring = tf.empty(type=Tensor[(64, 256), "bf16", Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))) | Swizzle(3, 4, 3), "smem"], buffers=3)  # Tensor[(64, 256), "bf16", Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))) | Swizzle(3, 4, 3), "smem"]; memory footprint-precision=exact
                        rhs = tf.schedule((v8, rhs_ring), op=T.copy_async_tensor())  # Tensor[(64, 256), "bf16", Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))) | Swizzle(3, 4, 3), "smem"]; memory traffic=gmem:r32.00KB/w0@logical,r32.00KB/w0@total,r32.00KB/w0@cta,r32.00KB/w0@thread;smem:r0/w32.00KB@logical,r0/w32.00KB@total,r0/w32.00KB@cta,r0/w32.00KB@thread footprint=b:32.00KB footprint-precision=exact
                    with Mesh(thread[1:], layout=(2, 4, 8, 4), names=("group", "warp", "lane8", "lane4")) as thread_4:  # memory peak=rmem:128.00KB;smem:144.00KB
                        acc_3 = tf.schedule((acc, lhs, rhs), op=T.tiled_mma(atom=T.cuda.sm90.Wgmma(n=256, dtype="bf16", form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, b_major=T.cuda.sm90.Major.MN)), repeat=(2, 1, 4))  # Tensor[(128, 256), "f32", ((2 @ thread_4.group, 4 @ thread_4.warp, 2, 8 @ thread_4.lane8, 32, 4 @ thread_4.lane4, 2), (16384, 16, 8, 1, 512, 128, 64)), "rmem"]; memory traffic=rmem:r128.00KB/w128.00KB@logical,r128.00KB/w128.00KB@total,r128.00KB/w128.00KB@cta,r512B/w512B@thread;smem:r48.00KB/w0@logical,r48.00KB/w0@total,r48.00KB/w0@cta,r8.25KB/w0@thread footprint-precision=exact
                    acc = acc_3
                with Mesh(thread[1:], layout=(2, 4, 8, 4), names=("group", "warp", "lane8", "lane4")) as thread_5:  # memory peak=rmem:128.00KB;smem:64.00KB
                    tile_out = tf.cast(acc, dtype="bf16")  # Tensor[(128, 256), "bf16", ((2 @ thread_5.group, 4 @ thread_5.warp, 2, 8 @ thread_5.lane8, 32, 4 @ thread_5.lane4, 2), (16384, 16, 8, 1, 512, 128, 64)), "rmem"]; memory traffic=rmem:r128.00KB/w64.00KB@logical,r128.00KB/w64.00KB@total,r128.00KB/w64.00KB@cta,r512B/w256B@thread footprint-precision=exact
                    staged_dst = tf.empty(type=Tensor[(128, 256), "bf16", Layout((128, (4, 64)), (64, (8192, 1))) | Swizzle(3, 4, 3), "smem"])  # Tensor[(128, 256), "bf16", Layout((128, (4, 64)), (64, (8192, 1))) | Swizzle(3, 4, 3), "smem"]; memory footprint-precision=exact
                    staged = tf.schedule((tile_out, staged_dst), op=T.copy())  # Tensor[(128, 256), "bf16", Layout((128, (4, 64)), (64, (8192, 1))) | Swizzle(3, 4, 3), "smem"]; memory traffic=rmem:r64.00KB/w0@logical,r64.00KB/w0@total,r64.00KB/w0@cta,r256B/w0@thread;smem:r0/w64.00KB@logical,r0/w64.00KB@total,r0/w64.00KB@cta,r0/w256B@thread footprint-precision=exact
                with thread[:1, :32] as thread_6:  # memory peak=smem:64.00KB
                    tile_dst = tf.empty(type=Tensor[(128, 256), "bf16", (128, 256)])  # Tensor[(128, 256), "bf16", (128, 256)]; memory footprint-precision=exact
                    tile = tf.schedule((staged, tile_dst), op=T.copy_async_tensor())  # Tensor[(128, 256), "bf16", (128, 256)]; memory traffic=gmem:r0/w64.00KB@logical,r0/w64.00KB@total,r0/w64.00KB@cta,r0/w64.00KB@thread;smem:r64.00KB/w0@logical,r64.00KB/w0@total,r64.00KB/w0@cta,r64.00KB/w0@thread footprint=tile_dst:64.00KB footprint-precision=exact
                    out_3 = tf.insert_slice(out, tile, (m, n))  # Tensor[(8192, 17408), "bf16"]; memory traffic=gmem:r64.00KB/w64.00KB@logical,r64.00KB/w64.00KB@total,r64.00KB/w64.00KB@cta,r64.00KB/w64.00KB@thread;rmem:r16B/w0@logical,r16B/w0@total,r16B/w0@cta,r16B/w0@thread footprint=tile_dst:64.00KB;out_3:64.00KB footprint-precision=exact
                out = out_3
            out = out
    return out
