"""
Stage B: Two MMAs + Identity Softmax with Layout Transform

Following NVIDIA's fmha.py softmax_step pattern exactly.

Architecture:
  MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False)
  Identity softmax: tcgen05.ld from C-layout → convert F32→BF16 → tcgen05.st to A-layout
  MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True)

Reference: output = (Q @ K^T) @ V (no softmax, P = raw scores)

TMEM Layout (following fmha.py for 128x128 MMA tile):
  tmem_s0 = 0   (scores, QK accumulator, C-layout)
  tmem_p0 = 32  (P, PV A-operand, A-layout — written by identity softmax)
  tmem_o0 = 256 (output, PV accumulator, C-layout)

The identity softmax performs the C-layout → A-layout transform via tcgen05.ld + tcgen05.st.
This is the critical bridge that makes the two-MMA pipeline work on Blackwell.
"""

import torch
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
from cutlass.utils import LayoutEnum
import cuda.bindings.driver as cuda


class StageBIdentitySoftmaxKernel:
    def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True):
        self.acc_dtype = Float32
        self.qk_acc_dtype = Float32
        self.q_dtype = BFloat16
        self.o_dtype = BFloat16
        self.use_2cta_instrs = use_2cta_instrs
        self.mma_tiler_mn = mma_tiler_mn
        self.mma_tiler = (*mma_tiler_mn, 1)
        self.use_tma_store = use_tma_store
        self.cluster_shape_mn = (1, 1)
        self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE

        self.softmax_warp_ids = (0, 1, 2, 3)
        self.epilogue_warp_id = self.softmax_warp_ids  # same warps do softmax + epilogue
        self.mma_warp_id = 4
        self.tma_warp_id = 5
        self.threads_per_cta = 32 * 6

        # TMEM offsets (fmha.py pattern)
        self.tmem_s0_offset = 0
        self.tmem_o0_offset = 256
        self.tmem_p0_offset = 32

        self.tmem_alloc_cols = 512

        self.epilog_sync_bar_id = 1
        self.tmem_alloc_sync_bar_id = 2
        self.tmem_dealloc_sync_bar_id = 3
        self.scores_full_bar_id = 4
        self.softmax_done_bar_id = 5

        self.num_c_stage = 2

    def _setup_attributes(self, qk_mma, pv_mma):
        qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2])
        self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4)
        pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2])
        self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4)
        self.mma_tiler = self.qk_mma_tiler

        self.cta_tile_shape_mnk = (
            self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[2])
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((1, 1, 1)), (qk_mma.thr_id.shape,))

        self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk, self.use_2cta_instrs,
            self.c_layout, self.o_dtype)

        self.num_ab_stage = 1
        self.num_acc_stage = 1

        self.q_smem_layout_staged = utils.sm100.make_smem_layout_a(
            qk_mma, self.qk_mma_tiler, self.a_dtype, self.num_ab_stage)
        self.k_smem_layout_staged = utils.sm100.make_smem_layout_b(
            qk_mma, self.qk_mma_tiler, self.b_dtype, self.num_ab_stage)
        self.v_smem_layout_staged = utils.sm100.make_smem_layout_b(
            pv_mma, self.pv_mma_tiler, self.b_dtype, self.num_ab_stage)
        self.p_tmem_layout_staged = utils.sm100.make_smem_layout_a(
            pv_mma, self.pv_mma_tiler, self.q_dtype, self.num_ab_stage)
        self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi(
            self.o_dtype, self.c_layout, self.epi_tile, self.num_c_stage)

        # For TMEM allocation
        acc_shape_qk = qk_mma.partition_shape_C(self.mma_tiler_mn)
        tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape_qk, self.num_acc_stage))
        self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtS_fake, arch="sm_100")

        q_smem = cute.slice_(self.q_smem_layout_staged, (None, None, None, 0))
        k_smem = cute.slice_(self.k_smem_layout_staged, (None, None, None, 0))
        self.num_tma_load_bytes = (
            cute.size_in_bytes(self.a_dtype, q_smem) +
            cute.size_in_bytes(self.b_dtype, k_smem)
        ) * cute.size(qk_mma.thr_id.shape)

    @cute.jit
    def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor,
                 stream: cuda.CUstream):
        self.a_dtype = a.element_type
        self.b_dtype = b.element_type
        self.c_dtype = c.element_type
        self.a_major_mode = LayoutEnum.from_tensor(a).mma_major_mode()
        self.b_major_mode = LayoutEnum.from_tensor(b).mma_major_mode()
        self.c_layout = LayoutEnum.from_tensor(c)

        qk_mma = utils.sm100.make_trivial_tiled_mma(
            self.a_dtype, self.b_dtype, self.a_major_mode, self.b_major_mode,
            self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn,
            tcgen05.OperandSource.SMEM)

        pv_mma = utils.sm100.make_trivial_tiled_mma(
            self.a_dtype, self.b_dtype,
            cute.nvgpu.OperandMajorMode.K, self.b_major_mode,
            self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn,
            tcgen05.OperandSource.TMEM)

        self._setup_attributes(qk_mma, pv_mma)

        q_smem = cute.slice_(self.q_smem_layout_staged, (None, None, None, 0))
        k_smem = cute.slice_(self.k_smem_layout_staged, (None, None, None, 0))

        tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A(
            utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id),
            a, q_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape)
        tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B(
            utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id),
            b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape)

        epi_smem = cute.select(self.c_smem_layout_staged, mode=[0, 1])
        tma_c, tma_tc = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile)

        self._kernel(
            qk_mma, pv_mma,
            tma_q, tma_tq, tma_k, tma_tk,
            tma_c, tma_tc,
            self.cluster_layout_vmnk,
            self.q_smem_layout_staged, self.k_smem_layout_staged,
            self.v_smem_layout_staged, self.p_tmem_layout_staged,
            self.c_smem_layout_staged, self.epi_tile,
        ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream)

    @cute.kernel
    def _kernel(self, qk_mma, pv_mma,
                tma_q, mQ, tma_k, mK,
                tma_c, mC, cl_vmnk,
                q_smem_staged, k_smem_staged,
                v_smem_staged, p_tmem_staged,
                c_smem_staged, epi_tile):
        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
        tidx, _, _ = cute.arch.thread_idx()
        use_2cta_instrs = cute.size(qk_mma.thr_id.shape) == 2
        is_leader_cta = True

        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_q)
            cpasync.prefetch_descriptor(tma_k)
            cpasync.prefetch_descriptor(tma_c)

        @cute.struct
        class SharedStorage:
            ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
            tmem_dealloc_mbar: cutlass.Int64
            tmem_holding_buf: cutlass.Int32
            scores_full_mbar: cutlass.Int64
            softmax_done_mbar: cutlass.Int64

        smem = utils.SmemAllocator()
        storage = smem.allocate(SharedStorage)

        ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
            num_stages=self.num_ab_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
            tx_count=self.num_tma_load_bytes,
            cta_layout_vmnk=cl_vmnk,
            defer_sync=True,
        ).make_participants()

        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                32 * len(self.softmax_warp_ids) * (2 if use_2cta_instrs else 1)),
            cta_layout_vmnk=cl_vmnk,
            defer_sync=True,
        )

        tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=self.tmem_alloc_sync_bar_id,
            num_threads=32 * len((self.mma_warp_id, *self.softmax_warp_ids)),
        )
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=tmem_alloc_barrier,
            allocator_warp_id=self.softmax_warp_ids[0],
            is_two_cta=use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
        )

        scores_full_mbar = pipeline.NamedBarrier(
            barrier_id=self.scores_full_bar_id,
            num_threads=32 * (1 + len(self.softmax_warp_ids)),
        )
        softmax_done_mbar = pipeline.NamedBarrier(
            barrier_id=self.softmax_done_bar_id,
            num_threads=32 * (1 + len(self.softmax_warp_ids)),
        )

        pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True)

        sQ = smem.allocate_tensor(
            element_type=self.a_dtype, layout=q_smem_staged.outer,
            byte_alignment=128, swizzle=q_smem_staged.inner)
        sK = smem.allocate_tensor(
            element_type=self.b_dtype, layout=k_smem_staged.outer,
            byte_alignment=128, swizzle=k_smem_staged.inner)
        sC = smem.allocate_tensor(
            element_type=self.o_dtype, layout=c_smem_staged.outer,
            byte_alignment=128, swizzle=c_smem_staged.inner)

        gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None, 0, None)), (None, None, None))
        gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0, None, None)), (None, None, None))
        gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None, None, 0)), (None, None, None))
        k_tile_cnt = cute.size(gQ, mode=[3])

        qk_thr = qk_mma.get_slice(0)
        tCgQ = qk_thr.partition_A(gQ)
        tCgK = qk_thr.partition_B(gK)
        tCgC = qk_thr.partition_C(gC)

        a_cta_layout = cute.make_layout(cute.slice_(cl_vmnk, (0, 0, None, 0)).shape)
        tAsQ, tAgQ = cpasync.tma_partition(
            tma_q, 0, a_cta_layout,
            cute.group_modes(sQ, 0, 3), cute.group_modes(tCgQ, 0, 3))
        b_cta_layout = cute.make_layout(cute.slice_(cl_vmnk, (0, None, 0, 0)).shape)
        tAsK, tAgK = cpasync.tma_partition(
            tma_k, 0, b_cta_layout,
            cute.group_modes(sK, 0, 3), cute.group_modes(tCgK, 0, 3))
        tAgQ = tAgQ[(None, 0, None, 0)]
        tAgK = tAgK[(None, 0, None, 0)]

        tCrQ = qk_mma.make_fragment_A(sQ)
        tCrK = qk_mma.make_fragment_B(sK)
        tCrV = pv_mma.make_fragment_B(sK)

        # ── TMEM tensor setup (following fmha.py) ──
        # QK accumulator (scores) — 2D C-layout (fmha.py pattern)
        qk_acc_shape = qk_thr.partition_shape_C(self.mma_tiler_mn)
        tStS = qk_thr.make_fragment_C(qk_acc_shape)
        tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout)

        # PV accumulator (output) — 2D C-layout
        pv_thr = pv_mma.get_slice(0)
        pv_acc_shape = pv_mma.partition_shape_C(self.mma_tiler_mn)
        tOtO = pv_thr.make_fragment_C(pv_acc_shape)
        tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout)

        # P fragment for PV MMA (a_source=TMEM, A-layout)
        tP = cute.make_tensor(tStS.iterator, p_tmem_staged.outer)
        tOrP_base = pv_mma.make_fragment_A(tP)
        tOrP = tOrP_base[(None, None, None, 0)]
        tOrP0 = cute.make_tensor(
            tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset,
            tOrP.layout,
        )

        # Fake accumulators with stage dim (for epilogue_tma_store + TMEM allocation)
        tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage))
        tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage))

        pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)

        # ══════════════════════════════════════════════════════════
        # TMA LOAD WARP (warp 5)
        # ══════════════════════════════════════════════════════════
        if warp_idx == self.tma_warp_id:
            ab_producer.reset()
            peek_ab_empty_status = ab_producer.try_acquire()

            for k_tile in cutlass.range(k_tile_cnt, unroll=1):
                handle = ab_producer.acquire_and_advance(peek_ab_empty_status)
                cute.copy(tma_q, tAgQ[(None, handle.count)], tAsQ[(None, handle.index)],
                          tma_bar_ptr=handle.barrier)
                cute.copy(tma_k, tAgK[(None, handle.count)], tAsK[(None, handle.index)],
                          tma_bar_ptr=handle.barrier)
                peek_ab_empty_status = cutlass.Boolean(1)
                if handle.count + 1 < k_tile_cnt:
                    peek_ab_empty_status = ab_producer.try_acquire()

            ab_producer.tail()

        # ══════════════════════════════════════════════════════════
        # MMA WARP (warp 4)
        # ══════════════════════════════════════════════════════════
        if warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()

            ab_consumer.reset()
            peek_ab_full_status = ab_consumer.try_wait()

            # QK MMA: Q @ K^T → tmem_scores
            qk_mma.set(tcgen05.Field.ACCUMULATE, False)
            for k_tile in range(k_tile_cnt):
                if is_leader_cta:
                    handle = ab_consumer.wait_and_advance(peek_ab_full_status)
                    num_kblocks = cute.size(tCrQ, mode=[2])
                    for kblk_idx in cutlass.range(num_kblocks, unroll_full=True):
                        kblk_crd = (None, None, kblk_idx, handle.index)
                        cute.gemm(qk_mma, tStS0, tCrQ[kblk_crd], tCrK[kblk_crd], tStS0)
                    handle.release()
                    peek_ab_full_status = cutlass.Boolean(1)
                    if handle.count + 1 < k_tile_cnt:
                        peek_ab_full_status = ab_consumer.try_wait()

            cute.arch.fence_view_async_tmem_store()
            scores_full_mbar.arrive()
            softmax_done_mbar.wait()

            # PV MMA: P @ V → tmem_output
            pv_mma.set(tcgen05.Field.ACCUMULATE, True)
            tCrV_s = tCrV[(None, None, None, 0)]
            num_pv_kblocks = cute.size(tOrP0, mode=[2])
            for kblk_idx in cutlass.range(num_pv_kblocks, unroll_full=True):
                cute.gemm(pv_mma, tOtO0, tOrP0[(None, None, kblk_idx)],
                          tCrV_s[(None, None, kblk_idx)], tOtO0)

            acc_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_acc_stage)
            acc_pipeline.producer_acquire(acc_producer_state)
            acc_pipeline.producer_commit(acc_producer_state)
            acc_producer_state.advance()
            acc_pipeline.producer_tail(acc_producer_state)

        # ══════════════════════════════════════════════════════════
        # SOFTMAX / EPILOGUE WARPS (0..3)
        # ══════════════════════════════════════════════════════════
        if warp_idx < self.mma_warp_id:
            tmem.allocate(self.tmem_alloc_cols)
            tmem.wait_for_alloc()
            tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)

            # ── Identity softmax: C-layout → A-layout transform ──

            # 1. LOAD pipeline (reads from QK C-layout)
            tmem_load_atom = cute.make_copy_atom(
                tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
                self.qk_acc_dtype,
            )
            tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0)
            softmax_thread_idx = tidx % (32 * len(self.softmax_warp_ids))
            thr_tmem_load = tiled_tmem_load.get_slice(softmax_thread_idx)
            tTMEM_LOADtS = thr_tmem_load.partition_S(tStS0)
            cS = cute.make_identity_tensor(
                (self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
            tScS = qk_thr.partition_C(cS)
            tTMEM_LOADcS = thr_tmem_load.partition_D(tScS)

            # 2. STORE pipeline (writes P in A-layout at tmem_p0_offset)
            tilePlikeFP32 = self.qk_mma_tiler[1] // 32 * self.o_dtype.width
            tStS_P_layout = cute.composition(
                tStS.layout, cute.make_layout((128, tilePlikeFP32)))
            tStS_P = cute.make_tensor(
                tStS.iterator + self.tmem_p0_offset, tStS_P_layout)
            tmem_store_atom = cute.make_copy_atom(
                tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)),
                self.qk_acc_dtype,
            )
            tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P)
            thr_tmem_store = tiled_tmem_store.get_slice(softmax_thread_idx)
            tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P)
            tScS_P_layout = cute.composition(
                tScS.layout, cute.make_layout((128, tilePlikeFP32)))
            tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout)
            tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P)

            # 3. Wait for scores
            scores_full_mbar.wait()

            # 4. Load scores from C-layout → registers
            tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype)
            cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
            cute.arch.fence_view_async_tmem_load()

            # 5. IDENTITY: convert F32 → Q dtype, no softmax math
            tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype)
            tTMEM_STORErS_x4_e = cute.make_tensor(
                cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype),
                tTMEM_LOADrS.layout,
            )
            s_vec = tTMEM_LOADrS.load()
            tTMEM_STORErS_x4_e.store(s_vec.to(self.q_dtype))

            # 6. Store into A-layout (P region)
            cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4)
            cute.arch.fence_view_async_tmem_store()

            # 7. Signal MMA warp
            softmax_done_mbar.arrive()

            # ── Epilogue: write output to GMEM ──
            tCtO_base = cute.make_tensor(
                tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)

            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage)
            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread, 32 * len(self.softmax_warp_ids))
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_c_stage, producer_group=c_producer_group)

            epilogue_op = const_expr(lambda x: x)
            acc_consumer_state = utils.gemm.sm100.epilogue_tma_store(
                self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC,
                epi_tile, 0, epilogue_op, (0, 0, 0),
                acc_consumer_state, acc_pipeline, c_pipeline)

            c_pipeline.producer_tail()
            tmem.relinquish_alloc_permit()
            tmem.free(tmem_ptr)


def test_stage_b_identity_softmax():
    """Test Stage B: (Q @ K^T) @ V with identity softmax layout transform"""
    torch.manual_seed(42)

    m, n, k = 128, 128, 128

    q = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda")
    kv = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda")
    c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda")

    qf = q[:, :, 0].float()
    kvf = kv[:, :, 0].float()
    scores = qf @ kvf.T
    ref = scores @ kvf

    import cutlass.torch as cutlass_torch
    mQ = cutlass_torch.from_dlpack(q).mark_layout_dynamic(
        leading_dim=cutlass_torch.get_leading_dim(q))
    mK = cutlass_torch.from_dlpack(kv).mark_layout_dynamic(
        leading_dim=cutlass_torch.get_leading_dim(kv))
    mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic(
        leading_dim=cutlass_torch.get_leading_dim(c))

    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

    kernel = StageBIdentitySoftmaxKernel(
        mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True)
    print("Compiling Stage B (identity softmax with layout transform)...", flush=True)
    compiled = cute.compile(kernel, mQ, mK, mC, stream)

    print("Running...", flush=True)
    compiled(mQ, mK, mC, stream)
    torch.cuda.synchronize()

    output = c[:, :, 0].float()
    cos = torch.nn.functional.cosine_similarity(
        output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
    max_err = (output - ref).abs().max().item()

    print("Stage B: (Q @ K^T) @ V with identity softmax layout transform")
    print("  Shape: Q({},{}), K/V({},{}), output({},{})".format(m, k, n, k, m, n))
    print("  Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err))
    print("  {}".format("PASS" if cos >= 0.99 else "FAIL"))
    return cos


if __name__ == "__main__":
    test_stage_b_identity_softmax()
