mlas-sys 0.1.0-dev.5

// Shim: expose plain C entry points over MLAS's C++ SGEMM API so Rust can call
// it via FFI without bindgen.
//
// Threading (multi-thread MLAS): MLAS's high-level GEMM (`MlasGemmBatch`)
// computes its own cache-aware M/N thread partitioning and dispatches the tiles
// through `MlasTrySimpleParallel` / `MlasGetMaximumThreadCount`. In
// `BUILD_MLAS_NO_ONNXRUNTIME` (standalone) mode those two primitives normally
// degrade to a serial loop with a hard thread cap of 1. Rather than fight MLAS
// by re-partitioning at the Rust level, we let MLAS keep its own partitioning
// and give it a *pluggable parallel-for backend*: the vendored standalone
// primitives call the `MlasStandalone*` hooks below, which forward the
// parallel-for onto a real thread pool that Rust drives with Rayon (the same
// global pool the rest of ep-cpu uses, so there is no oversubscription). When
// no backend is registered (e.g. the mlas-sys unit tests, or the ep-cpu default
// build) the hooks run serially — identical to the original spike behaviour.
//
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/mlas/inc/mlas.h"
#include "core/mlas/inc/mlas_qnbit.h"

#include <cstddef>
#include <cstring>
#include <functional>
#include <new>

// ---- Pluggable parallel-for backend (driven from Rust/Rayon) ----------------

extern "C" {

// One unit of MLAS work: run partition `tid`. `task_ctx` is opaque to Rust and
// points back to the C++ trampoline state.
typedef void (*mlas_task_fn)(void* task_ctx, std::ptrdiff_t tid);

// Run `task(task_ctx, tid)` for every tid in [0, iterations) across the backing
// pool. `rust_ctx` is the opaque pointer registered with `mlas_set_threading`.
typedef void (*mlas_parallel_for_fn)(
    void* rust_ctx,
    std::ptrdiff_t iterations,
    mlas_task_fn task,
    void* task_ctx);

// Report the degree of parallelism MLAS may use for its partitioning.
typedef int (*mlas_max_threads_fn)(void* rust_ctx);

}  // extern "C"

namespace {
mlas_parallel_for_fn g_parallel_for = nullptr;
mlas_max_threads_fn g_max_threads = nullptr;
void* g_rust_ctx = nullptr;
}  // namespace

// Register (or clear, with all-null args) the Rust-backed threading backend.
extern "C" void mlas_set_threading(
    mlas_parallel_for_fn parallel_for,
    mlas_max_threads_fn max_threads,
    void* rust_ctx)
{
    g_parallel_for = parallel_for;
    g_max_threads = max_threads;
    g_rust_ctx = rust_ctx;
}

// Hook called by the vendored standalone `MlasGetMaximumThreadCount`.
extern "C" int MlasStandaloneMaxThreads()
{
    if (g_max_threads != nullptr) {
        int n = g_max_threads(g_rust_ctx);
        return n > 0 ? n : 1;
    }
    return 1;
}

// Trampoline that lets the C parallel-for callback invoke a C++
// `std::function<void(ptrdiff_t)>` (the closure MLAS passes to
// `MlasTrySimpleParallel`) without exposing C++ types across FFI.
namespace {
void mlas_std_function_trampoline(void* task_ctx, std::ptrdiff_t tid)
{
    (*static_cast<const std::function<void(std::ptrdiff_t)>*>(task_ctx))(tid);
}
}  // namespace

// Hook called by the vendored standalone `MlasTrySimpleParallel`. Forwards the
// parallel-for onto the registered backend, or runs serially if none is set.
extern "C" void MlasStandaloneParallelFor(std::ptrdiff_t iterations, void* work)
{
    const auto& fn = *static_cast<const std::function<void(std::ptrdiff_t)>*>(work);
    if (g_parallel_for != nullptr && iterations > 1) {
        g_parallel_for(
            g_rust_ctx,
            iterations,
            &mlas_std_function_trampoline,
            const_cast<void*>(static_cast<const void*>(&fn)));
    } else {
        for (std::ptrdiff_t tid = 0; tid < iterations; ++tid) {
            fn(tid);
        }
    }
}

extern "C" void mlas_sgemm(
    int transA,   // 0 = no-transpose, 1 = transpose
    int transB,
    size_t M,
    size_t N,
    size_t K,
    float alpha,
    const float* A,
    size_t lda,
    const float* B,
    size_t ldb,
    float beta,
    float* C,
    size_t ldc)
{
    MLAS_SGEMM_DATA_PARAMS data;
    data.A = A;
    data.lda = lda;
    data.B = B;
    data.ldb = ldb;
    data.C = C;
    data.ldc = ldc;
    data.alpha = alpha;
    data.beta = beta;
    data.BIsPacked = false;

    MlasGemmBatch(
        transA ? CblasTrans : CblasNoTrans,
        transB ? CblasTrans : CblasNoTrans,
        M, N, K,
        &data, 1,
        /*ThreadPool=*/nullptr,
        /*BackendKernelSelectorConfig=*/nullptr);
}

// ---- Pre-packed B variant (mirrors how ORT pre-packs constant weights) ----

extern "C" size_t mlas_sgemm_pack_b_size(int transA, int transB, size_t N, size_t K)
{
    return MlasGemmPackBSize(
        transA ? CblasTrans : CblasNoTrans,
        transB ? CblasTrans : CblasNoTrans,
        N, K, nullptr);
}

extern "C" void mlas_sgemm_pack_b(
    int transA, int transB, size_t N, size_t K,
    const float* B, size_t ldb, void* packed_b)
{
    MlasGemmPackB(
        transA ? CblasTrans : CblasNoTrans,
        transB ? CblasTrans : CblasNoTrans,
        N, K, B, ldb, packed_b, nullptr);
}

// ---- Blocked n-bit quantized GEMM (SQNBitGemm) ------------------------------
//
// Plain-C wrappers over MLAS's templated `MlasQNBitGemmBatch<float>` and its
// pack/query helpers so Rust can drive the int4/int8 blockwise-quantized
// MatMulNBits decode path without binding the C++ template/struct directly.
//
// `comp_type` is the raw `MLAS_QNBIT_GEMM_COMPUTE_TYPE` value:
//   0 = SQNBIT_CompFp32 (fp32 activation, fp32 accumulate)
//   3 = SQNBIT_CompInt8 (int8 activation, int32 accumulate) -- accuracy_level=4.
//
// Threading: like the SGEMM shim, MLAS's own N/M tile partitioning is routed
// through the registered Rust/Rayon parallel-for backend (`MlasStandalone*`
// hooks above). `MlasQNBitGemmBatch` only takes its parallel branch when
// `ThreadPool != nullptr`, so pass a non-null sentinel (the pointer is never
// dereferenced in the standalone build -- `MlasGetMaximumThreadCount` and
// `MlasTrySimpleParallel` both ignore it) to enable multi-threading.

extern "C" int mlas_qnbit_gemm_available(size_t bits, size_t blk_len, int comp_type)
{
    return MlasIsQNBitGemmAvailable(
               bits, blk_len, static_cast<MLAS_QNBIT_GEMM_COMPUTE_TYPE>(comp_type))
               ? 1
               : 0;
}

extern "C" size_t mlas_qnbit_gemm_pack_b_size(
    size_t n, size_t k, size_t bits, size_t blk_len, int has_zp, int comp_type)
{
    return MlasQNBitGemmPackQuantBDataSize(
        n, k, bits, blk_len, has_zp != 0,
        static_cast<MLAS_QNBIT_GEMM_COMPUTE_TYPE>(comp_type),
        /*BackendKernelSelectorConfig=*/nullptr);
}

extern "C" void mlas_qnbit_gemm_pack_b(
    size_t n,
    size_t k,
    size_t bits,
    size_t blk_len,
    int comp_type,
    const void* quant_b_data,
    void* packed_b,
    const void* quant_b_scale,
    int has_zp,
    const void* quant_b_zero_point)
{
    MlasQNBitGemmPackQuantBData(
        n, k, bits, blk_len,
        static_cast<MLAS_QNBIT_GEMM_COMPUTE_TYPE>(comp_type),
        quant_b_data,
        packed_b,
        quant_b_scale,
        has_zp != 0,
        quant_b_zero_point,
        /*ThreadPool=*/nullptr,
        /*BackendKernelSelectorConfig=*/nullptr);
}

extern "C" size_t mlas_qnbit_gemm_workspace_size(
    size_t m, size_t n, size_t k, size_t bits, size_t blk_len, int has_zp, int comp_type)
{
    return MlasQNBitGemmBatchWorkspaceSize(
        m, n, k, /*BatchN=*/1, bits, blk_len, has_zp != 0,
        static_cast<MLAS_QNBIT_GEMM_COMPUTE_TYPE>(comp_type),
        /*BackendKernelSelectorConfig=*/nullptr);
}

extern "C" void mlas_qnbit_gemm(
    size_t m,
    size_t n,
    size_t k,
    size_t bits,
    size_t blk_len,
    int comp_type,
    const float* a,
    size_t lda,
    const void* packed_b,
    const float* quant_b_scale,
    int has_zp,
    const void* quant_b_zero_point,
    const float* bias,
    float* c,
    size_t ldc,
    void* workspace,
    int multithread)
{
    MLAS_QNBIT_GEMM_DATA_PARAMS<float> params;
    params.A = a;
    params.lda = lda;
    params.Bias = bias;
    params.C = c;
    params.ldc = ldc;

    const auto ct = static_cast<MLAS_QNBIT_GEMM_COMPUTE_TYPE>(comp_type);
    if (ct == SQNBIT_CompInt8) {
        // The int8-compute path derives PackedQuantBData / QuantBScale /
        // QuantBBlkSum from the combined workspace produced by
        // MlasQNBitGemmPackQuantBData (which baked scale + zero point into the
        // block sums), so only the workspace pointer is needed here.
        params.QuantBDataWorkspace = packed_b;
    } else {
        // The fp32-compute path repacks only the quantized nibbles; scales and
        // (optional) zero points are consumed at compute time in their original
        // ONNX layout.
        params.PackedQuantBData = static_cast<const std::byte*>(packed_b);
        params.QuantBScale = quant_b_scale;
        params.QuantBZeroPoint = has_zp != 0 ? quant_b_zero_point : nullptr;
    }

    MLAS_THREADPOOL* thread_pool =
        multithread != 0 ? reinterpret_cast<MLAS_THREADPOOL*>(1) : nullptr;

    MlasQNBitGemmBatch<float>(
        m, n, k, /*BatchN=*/1, bits, blk_len, ct, &params, workspace, thread_pool,
        /*BackendKernelSelectorConfig=*/nullptr);
}

extern "C" void mlas_sgemm_packed(
    int transA,
    int transB,
    size_t M,
    size_t N,
    size_t K,
    float alpha,
    const float* A,
    size_t lda,
    const void* packed_b,
    float beta,
    float* C,
    size_t ldc)
{
    MLAS_SGEMM_DATA_PARAMS data;
    data.A = A;
    data.lda = lda;
    data.B = reinterpret_cast<const float*>(packed_b);
    data.ldb = 0;
    data.C = C;
    data.ldc = ldc;
    data.alpha = alpha;
    data.beta = beta;
    data.BIsPacked = true;

    MlasGemmBatch(
        transA ? CblasTrans : CblasNoTrans,
        transB ? CblasTrans : CblasNoTrans,
        M, N, K,
        &data, 1,
        /*ThreadPool=*/nullptr,
        /*BackendKernelSelectorConfig=*/nullptr);
}

// ---- Vectorized logistic (sigmoid) -----------------------------------------
// Exposes MLAS's SIMD sigmoid so callers can build activations such as SiLU
// (`x * sigmoid(x)`) on top of a battle-tested vectorized exp instead of a
// scalar `expf` loop. Single-threaded: the caller drives any outer sharding.
extern "C" void mlas_compute_logistic(
    const float* input,
    float* output,
    size_t n)
{
    MlasComputeLogistic(input, output, n);
}

// ---- Vectorized SiLU -------------------------------------------------------
// MLAS selects its fused AVX-512F implementation at runtime when available,
// with a portable logistic-then-multiply fallback on other architectures.
extern "C" void mlas_compute_silu(
    const float* input,
    float* output,
    size_t n)
{
    MlasComputeSilu(input, output, n);
}

extern "C" void mlas_eltwise_add(
    const float* left,
    const float* right,
    float* output,
    size_t n)
{
    MlasEltwiseAdd<float>(left, right, output, n);
}

extern "C" void mlas_compute_activation(
    int kind,
    float minimum,
    float maximum,
    const float* input,
    float* output,
    size_t n)
{
    if (input != output) {
        std::memcpy(output, input, n * sizeof(float));
    }
    MLAS_ACTIVATION activation{};
    activation.ActivationKind = static_cast<MLAS_ACTIVATION_KIND>(kind);
    activation.Parameters.Clip.minimum = minimum;
    activation.Parameters.Clip.maximum = maximum;
    MlasActivation(&activation, output, nullptr, 1, n, n);
}

// ---- Float convolution -------------------------------------------------------

namespace {
struct mlas_conv_plan {
    MLAS_ACTIVATION activation{};
    MLAS_CONV_PARAMETERS parameters{};
};
}  // namespace

extern "C" void* mlas_conv_prepare(
    size_t dimensions,
    size_t batch_count,
    size_t group_count,
    size_t input_channels_per_group,
    const int64_t* input_shape,
    const int64_t* kernel_shape,
    const int64_t* dilation_shape,
    const int64_t* padding,
    const int64_t* stride_shape,
    const int64_t* output_shape,
    size_t filter_count_per_group,
    size_t* working_buffer_elements)
{
    mlas_conv_plan* plan = nullptr;
    try {
        plan = new mlas_conv_plan();
        plan->activation.ActivationKind = MlasIdentityActivation;
        MLAS_THREADPOOL* thread_pool = reinterpret_cast<MLAS_THREADPOOL*>(1);
        MlasConvPrepare(
            &plan->parameters,
            dimensions,
            batch_count,
            group_count,
            input_channels_per_group,
            input_shape,
            kernel_shape,
            dilation_shape,
            padding,
            stride_shape,
            output_shape,
            filter_count_per_group,
            &plan->activation,
            working_buffer_elements,
            /*ChannelsLast=*/false,
            /*Beta=*/0.0f,
            thread_pool);
        return plan;
    } catch (...) {
        delete plan;
        return nullptr;
    }
}

extern "C" void mlas_conv_run(
    const void* opaque_plan,
    const float* input,
    const float* filter,
    const float* bias,
    float* working_buffer,
    float* output)
{
    const auto* plan = static_cast<const mlas_conv_plan*>(opaque_plan);
    MLAS_THREADPOOL* thread_pool = reinterpret_cast<MLAS_THREADPOOL*>(1);
    MlasConv(
        &plan->parameters,
        input,
        filter,
        bias,
        working_buffer,
        output,
        thread_pool);
}

extern "C" void mlas_conv_plan_destroy(void* opaque_plan)
{
    delete static_cast<mlas_conv_plan*>(opaque_plan);
}

// ---- NCHWc blocked convolution (ORT's fast float Conv path) -----------------
//
// MLAS's NCHWc kernels keep activations in a channels-blocked layout (channels
// split into SIMD-width blocks) and consume a pre-reordered ("pre-packed")
// filter, exactly as ONNX Runtime's nchwc_transformer does. This is the
// high-throughput float convolution path; the plain `MlasConv` above falls back
// to im2col+GEMM. The caller (Rust) mirrors the transformer's per-conv layout
// decisions: reorder the filter once at kernel construction, reorder the input
// per call, run the blocked conv, then reorder the output back to NCHW.

extern "C" size_t mlas_nchwc_block_size()
{
    return MlasNchwcGetBlockSize();
}

extern "C" void mlas_nchwc_reorder_input_nchw(
    const float* source,
    float* dest,
    size_t channels,
    size_t input_size)
{
    MlasReorderInputNchw(source, dest, channels, input_size);
}

extern "C" void mlas_nchwc_reorder_output_nchw(
    const int64_t* output_shape,
    const float* source,
    float* dest)
{
    MlasReorderOutputNchw(output_shape, source, dest, /*ThreadPool=*/nullptr);
}

extern "C" void mlas_nchwc_reorder_filter_bibo(
    const int64_t* filter_shape,
    const float* source,
    float* dest)
{
    MlasReorderFilterOIHWBiBo(filter_shape, source, dest);
}

extern "C" void mlas_nchwc_reorder_filter_bo(
    const int64_t* filter_shape,
    const float* source,
    float* dest)
{
    MlasReorderFilterOIHWBo(filter_shape, source, dest);
}

extern "C" void mlas_nchwc_conv(
    const int64_t* input_shape,
    const int64_t* kernel_shape,
    const int64_t* dilation_shape,
    const int64_t* padding,
    const int64_t* stride_shape,
    const int64_t* output_shape,
    size_t group_count,
    const float* input,
    const float* filter,
    const float* bias,
    float* output,
    int activation_kind,
    float activation_value0,
    float activation_value1,
    int zero_mode)
{
    MLAS_ACTIVATION activation{};
    activation.ActivationKind = static_cast<MLAS_ACTIVATION_KIND>(activation_kind);
    activation.Parameters.Values[0] = activation_value0;
    activation.Parameters.Values[1] = activation_value1;
    MLAS_THREADPOOL* thread_pool = reinterpret_cast<MLAS_THREADPOOL*>(1);
    MlasNchwcConv(
        input_shape,
        kernel_shape,
        dilation_shape,
        padding,
        stride_shape,
        output_shape,
        group_count,
        input,
        filter,
        bias,
        output,
        &activation,
        zero_mode != 0,
        thread_pool,
        /*BackendKernelSelectorConfig=*/nullptr,
        /*UseBf16=*/false);
}

// ---- Float pooling -----------------------------------------------------------

extern "C" void mlas_pool(
    int kind,
    size_t dimensions,
    const int64_t* input_shape,
    const int64_t* kernel_shape,
    const int64_t* padding,
    const int64_t* stride_shape,
    const int64_t* output_shape,
    const float* input,
    float* output)
{
    MLAS_THREADPOOL* thread_pool = reinterpret_cast<MLAS_THREADPOOL*>(1);
    MlasPool(
        static_cast<MLAS_POOLING_KIND>(kind),
        dimensions,
        input_shape,
        kernel_shape,
        padding,
        stride_shape,
        output_shape,
        input,
        output,
        thread_pool);
}

// NCHWc blocked 2-D pooling. `input_shape`/`output_shape` are blocked NCHWc
// shapes (channel dimension rounded up to the SIMD block); MLAS pools each
// channel independently over the blocked buffer, so no NCHW<->NCHWc reorder is
// needed around it. Mirrors ONNX Runtime's NchwcTransformer pool handling.
extern "C" void mlas_nchwc_pool(
    int kind,
    const int64_t* input_shape,
    const int64_t* kernel_shape,
    const int64_t* dilation_shape,
    const int64_t* padding,
    const int64_t* stride_shape,
    const int64_t* output_shape,
    const float* input,
    float* output)
{
    MLAS_THREADPOOL* thread_pool = reinterpret_cast<MLAS_THREADPOOL*>(1);
    MlasNchwcPool(
        static_cast<MLAS_POOLING_KIND>(kind),
        input_shape,
        kernel_shape,
        dilation_shape,
        padding,
        stride_shape,
        output_shape,
        input,
        output,
        thread_pool);
}