Repository: microsoft/onnxruntime
File: onnxruntime/core/mlas/lib/kleidiai/sgemm_kleidiai.cpp (C++)

Code before:
        auto BTile = reinterpret_cast<const void*>(B_base + rhs_packed_offset);

        // Get lhs tile, A
        const size_t lhs_packed_offset =
            UseSME2 ? kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa(MIdx * m_step, K)
                    : kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa(MIdx * m_step, K);

        const std::byte* A_base = LhsPackedData + LhsPackedStride * BIdx;
        auto ATile = reinterpret_cast<const float*>(A_base + lhs_packed_offset);

        auto TileSizeM = (MIdx + 1) * m_step > M ? (M - MIdx * m_step) : m_step;
        auto TileSizeN = (NIdx + 1) * n_step > N ? (N - NIdx * n_step) : n_step;

        // Get result tile, C
        auto CTile = reinterpret_cast<void*>(
            reinterpret_cast<std::byte*>(Data[BIdx].C) +
            MIdx * m_step * Data[BIdx].ldc * sizeof(float) +
            NIdx * n_step * sizeof(float)
        );
        // Allocate temporary buffer for raw A*B result (TLS reusable buffer)

        size_t tile_elems = 0;
        if (mul_overflow_size_t_builtin(TileSizeM, TileSizeN, &tile_elems))
        {
            // size_t wraparound detected, exit
            return
        }
        if (g_kai_tls.output_tile.capacity() < tile_elems) {

            g_kai_tls.output_tile.reserve(tile_elems);
        }
        // resize the tile to the required size
        g_kai_tls.output_tile.resize(tile_elems);

        float* temp_tile = g_kai_tls.output_tile.data();
        std::fill_n(temp_tile, tile_elems, 0.0f);

        if (UseSME2) {
            kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa(
                TileSizeM,
                TileSizeN,
                K,
                ATile, BTile, temp_tile,
                TileSizeN * sizeof(float), sizeof(float),
                -std::numeric_limits<float>::max(), std::numeric_limits<float>::max()
            );
        } else {
            kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa(
                TileSizeM,
                TileSizeN,
                K,

Review comment (suggestion):
missing semi-colon?

also, this will just return early from the worker thread and not actually exit the function. perhaps this check should be moved before the `MlasTrySimpleParallel` call using the maximum possible values of `TileSizeM` and `TileSizeN`.

Code after:
        const size_t lhs_packed_offset =
            UseSME2 ? kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa(MIdx * m_step, K)
                    : kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa(MIdx * m_step, K);

        const std::byte* A_base = LhsPackedData + LhsPackedStride * BIdx;
        auto ATile = reinterpret_cast<const float*>(A_base + lhs_packed_offset);

        auto TileSizeM = (MIdx + 1) * m_step > M ? (M - MIdx * m_step) : m_step;
        auto TileSizeN = (NIdx + 1) * n_step > N ? (N - NIdx * n_step) : n_step;

        // Get result tile, C
        auto CTile = reinterpret_cast<void*>(
            reinterpret_cast<std::byte*>(Data[BIdx].C) +
            MIdx * m_step * Data[BIdx].ldc * sizeof(float) +
            NIdx * n_step * sizeof(float)
        );
        // Allocate temporary buffer for raw A*B result (TLS reusable buffer)
        size_t tile_elems = TileSizeM * TileSizeN;

        // resize the tile to the required size
        g_kai_tls.output_tile.resize(tile_elems);

        float* temp_tile = g_kai_tls.output_tile.data();
        std::fill_n(temp_tile, tile_elems, 0.0f);

        if (UseSME2) {
            kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa(
                TileSizeM,
                TileSizeN,
                K,
                ATile, BTile, temp_tile,
                TileSizeN * sizeof(float), sizeof(float),
                -std::numeric_limits<float>::max(), std::numeric_limits<float>::max()
            );
        } else {
            kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa(
                TileSizeM,
                TileSizeN,
                K,
                ATile, BTile, temp_tile,
                TileSizeN * sizeof(float), sizeof(float),
                -std::numeric_limits<float>::max(), std::numeric_limits<float>::max()
            );
        }

        // Final output tile pointer
        float* dst_tile = reinterpret_cast<float*>(CTile);

            // quick copy of data in cases where we are not scaling or accumulating anything
            // with bounds checking on tile sizing to ensure the data fits in the memory block
            bool can_memcpy = (