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

Code before:
            ThreadsPerGemm = ptrdiff_t(M);
        }

        WorkBlock.ThreadCountM = ThreadsPerGemm;
        WorkBlock.ThreadCountN = 1;
    }
    TargetThreadCount = ThreadsPerGemm * BatchN;

    MlasTrySimpleParallel(ThreadPool, TargetThreadCount, [&](ptrdiff_t tid) {
        const auto gemm_i = tid / ThreadsPerGemm;
        const auto blk_i = tid % ThreadsPerGemm;
        MlasGemmQuantThreaded(&WorkBlock, &Shape, &DataParams[gemm_i], blk_i);
    });
}

void
MLASCALL
MlasDynamicQGemmBatch (
    const MLAS_GEMM_DYN_QUANT_SHAPE_PARAMS& Shape,
    const MLAS_GEMM_DYN_QUANT_DATA_PARAMS* DataParams,
    const size_t BatchN,
    MLAS_THREADPOOL* ThreadPool
) {
#if defined(USE_KLEIDIAI) && !defined(_MSC_VER)
    //No fallback and putting in guards
    if(ArmKleidiAI::SMEInfo::CanUseSME2){
        ArmKleidiAI::MlasDynamicQGemmBatch(Shape, DataParams, BatchN, ThreadPool);
    }
#endif

    MLAS_UNREFERENCED_PARAMETER(Shape);
    MLAS_UNREFERENCED_PARAMETER(DataParams);
    MLAS_UNREFERENCED_PARAMETER(BatchN);
    MLAS_UNREFERENCED_PARAMETER(ThreadPool);
}

int32_t
MlasSymmQgemmGetKernelOutputCnt()
{
    const MLAS_SYMM_QGEMM_DISPATCH* dispatch = GetMlasPlatform().SymmQgemmDispatch;
    return int32_t(dispatch->StrideM);
}


void
MLASCALL
MlasSymmQgemmBatch(
    const MLAS_GEMM_QUANT_SHAPE_PARAMS& Shape,
    const MLAS_SYMM_QGEMM_DATA_PARAMS* DataParams,
    const size_t BatchN,
    MLAS_THREADPOOL* ThreadPool

Review comment (refactor):
there are other places that need to be updated, like:

https://github.com/microsoft/onnxruntime/blob/b3ba580346293ed1cab35b59cf93ae1da56e65f7/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc#L218

https://github.com/microsoft/onnxruntime/blob/b3ba580346293ed1cab35b59cf93ae1da56e65f7/onnxruntime/test/mlas/unittest/test_dynamic_qgemm.cpp#L24

I might be missing some.

I think it would be worth making a helper function like `MlasIsDynamicQGemmAvailable` that has the appropriate checks and using that instead.

Code after:
            ThreadsPerGemm = ptrdiff_t(M);
        }

        WorkBlock.ThreadCountM = ThreadsPerGemm;
        WorkBlock.ThreadCountN = 1;
    }
    TargetThreadCount = ThreadsPerGemm * BatchN;

    MlasTrySimpleParallel(ThreadPool, TargetThreadCount, [&](ptrdiff_t tid) {
        const auto gemm_i = tid / ThreadsPerGemm;
        const auto blk_i = tid % ThreadsPerGemm;
        MlasGemmQuantThreaded(&WorkBlock, &Shape, &DataParams[gemm_i], blk_i);
    });
}

bool
MLASCALL
MlasIsDynamicQGemmAvailable()
{
#if defined(USE_KLEIDIAI) && !defined(_MSC_VER)
  return ArmKleidiAI::UseSME2;
#else
  return false;
#endif
}

void
MLASCALL
MlasDynamicQGemmBatch (
    const MLAS_GEMM_DYN_QUANT_SHAPE_PARAMS& Shape,
    const MLAS_GEMM_DYN_QUANT_DATA_PARAMS* DataParams,
    const size_t BatchN,
    MLAS_THREADPOOL* ThreadPool
) {
    assert(MlasIsDynamicQGemmAvailable());

#if defined(USE_KLEIDIAI) && !defined(_MSC_VER)
    //No fallback
    ArmKleidiAI::MlasDynamicQGemmBatch(Shape, DataParams, BatchN, ThreadPool);
#endif

    MLAS_UNREFERENCED_PARAMETER(Shape);
    MLAS_UNREFERENCED_PARAMETER(DataParams);
    MLAS_UNREFERENCED_PARAMETER(BatchN);
    MLAS_UNREFERENCED_PARAMETER(ThreadPool);
}

int32_t
MlasSymmQgemmGetKernelOutputCnt()
{
    const MLAS_SYMM_QGEMM_DISPATCH* dispatch = GetMlasPlatform().SymmQgemmDispatch;