.rst .pdf

repository open issue

Contents

Batch-Reduce General Matrix Multiplication

API Reference

General

The batch-reduce General Matrix Multiplication ukernel (BRGeMM) is an operation that computes a small matrix multiplication batch and accumulates their results in the same destination.

\(C = \sum_i A_i \cdot B_i\)

with

• \(A_i\) a set of matrices of dimension \(M \times K\)

• \(B_i\) a set of matrices of dimension \(K \times N\)

• \(C\) matrix of dimension \(M \times N\).

The BRGeMM ukernel also supports accumulation with values already present in \(C\), as well as post-operation and down-conversion to another \(D\) matrix:

\(D = \operatorname{convert}( \operatorname{post\_ops}(C + \sum_i A_i \cdot B_i, post\_ops\_args))\)

Data Types

In general, C represents an accumulation buffer. Hence, when computations are carried in floating-point arithmetic, C shall be of type f32; when computation is carried in integer arithmetic, C should be of type s32.

The BRGeMM ukernel supports the following combinations of data-types.

A	B	C	D
f32	f32	f32	u8, s8, s32, f32, f16, bf16
f16	f16	f32	u8, s8, s32, f32, f16, bf16
bf16	bf16	f32	u8, s8, s32, f32, f16, bf16
u8, s8	u8, s8	s32	u8, s8, s32, f32, f16, bf16

Data Representation

Because of hardware restrictions, the BRGeMM ukernel requires a specific data layout. For x86-64 architecture this layout applies to a B matrix. It is expressed through dnnl::ukernel::pack_type which can be queried by dnnl::ukernel::brgemm::get_B_pack_type call. If the query returns dnnl::ukernel::pack_type::no_trans, then packing is not required. Otherwise, the user is responsible for packing the data appropriately before calling dnnl::ukernel::brgemm::execute, either with custom code, or by using a dedicated set of APIs: dnnl::ukernel::transform::generate for generating a kernel of a transform routine and dnnl::ukernel::transform::execute to run the generated kernel.

Attributes

The following ukernel attributes can be set through dedicated setters.

Type	Operation	Description	Restrictions
Attribute	Scales	Scales the corresponding tensors by given scale factor(s)
Post-op	Eltwise	Applies an Eltwise operation to the result
Post-op	Binary	Applies a Binary operation to the result	General binary post-op restrictions

Note

if zero-points are passed for A/B, fpmath_mode should be set for the computation to happen over floating-point format (so up-conversion to floating-point format would happen before computation). If computation in integer format is needed, BRGeMM ukernel should be configured without zero-point, and the user should prepare a compensation term that will be passed to the binary post-op.

Implementation limitations

BRGeMM ukernel has no known limitations.

Examples

BRGeMM ukernel example

This C++ API example demonstrates how to create and execute a BRGeMM ukernel.

/*******************************************************************************
* Copyright 2024-2025 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/


#include <algorithm>
#include <cmath>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

#include "example_utils.hpp"
#include "oneapi/dnnl/dnnl_ukernel.hpp"

using namespace dnnl;
using namespace dnnl::ukernel;

void brgemm_example() {

    // Create execution dnnl::engine. Needed for reorders to operate over input
    // data.
    dnnl::engine engine(engine::kind::cpu, 0);

    // Create dnnl::stream. Needed for reorders for the same reason.
    dnnl::stream engine_stream(engine);

    // ukernel dimensions.
    // K is for a whole tensor, K_blk is for a single ukernel.
    const memory::dim M = 8, K = 128, K_blk = 64, N = 48;
    if (K % K_blk != 0) {
        printf("K_blk must divide K.\n");
        return;
    }
    const memory::dim n_calls = K / K_blk;

    memory::data_type a_dt = memory::data_type::u8;
    memory::data_type b_dt = memory::data_type::s8;
    memory::data_type c_dt = memory::data_type::s32; // Accumulator data type.
    memory::data_type d_dt = memory::data_type::f32; // Output data type.

    // Query the packing requirement from the ukernel. It's enough to query
    // packing requirements once for multiple ukernel objects.
    const auto pack = brgemm::get_B_pack_type(a_dt, b_dt);

    // If the value is `pack_type::undef`, ukernel API is not supported on the
    // target system.
    if (pack == pack_type::undef) {
        printf("Kernel is not supported on this platform.\n");
        return;
    }

    // Packing is required if the returned value is different from
    // `pack_type::no_pack`.
    // If packing is required, specific `ldb` value can be used ahead, since
    // transform has a limited set of supported values.
    bool need_pack = pack != pack_type::no_trans;

    const memory::dim lda = K;
    // `ldb` for `need_pack = true` must be one of 16, 32, 48, or 64. This
    // example doesn't explore options for dividing N into blocks which would
    // likely happen for N > 64.
    // const memory::dim ldb = need_pack ? N_block : N;
    const memory::dim ldb = N;
    const memory::dim ldc = N; // Leading dimension for accumulator.
    const memory::dim ldd = N; // Leading dimension for an actual output.
    const memory::dim batch_size = n_calls - 1;

    // A, B, and C tensors dimensions.
    memory::dims A_dims = {M, K};
    memory::dims B_dims = {K, N};
    memory::dims C_dims = {M, N};
    memory::dims D_dims = {M, N};
    memory::dims binary_add_dims = {1, 1};
    memory::dims B_scales_dims = {1, N};

    // Allocate buffers with user data.
    std::vector<float> A_user_data(product(A_dims));
    std::vector<float> B_user_data(product(B_dims));
    std::vector<float> binary_add_user_data(product(binary_add_dims));
    std::vector<float> B_scales_user_data(product(B_scales_dims));
    std::vector<float> D_data(product(D_dims)); // For reference comparison
    std::vector<float> D_user_data(product(D_dims)); // For reference comparison

    // Initialize A.
    std::generate(A_user_data.begin(), A_user_data.end(), []() {
        static int i = 0;
        return i++ % 4;
    });
    // Initialize B.
    std::generate(B_user_data.begin(), B_user_data.end(), []() {
        static int i = 6;
        static int sign_gen = 0;
        int sign = (sign_gen++ % 2) ? -1 : 1;
        float val = sign * (i++ % 5);
        return val;
    });
    // Initialize binary_add.
    std::generate(
            binary_add_user_data.begin(), binary_add_user_data.end(), []() {
                static int i = 3;
                return i++ % 6;
            });
    // Initialize B scales.
    std::generate(B_scales_user_data.begin(), B_scales_user_data.end(), []() {
        static int i = 4;
        return (float)(i++ % 16) / 8.f;
    });

    // Create f32 memories. They are used as data holders and reorder into
    // memories passed to the ukernel.
    auto A_f32_md = memory::desc(
            A_dims, memory::data_type::f32, memory::format_tag::ab);
    auto B_f32_md = memory::desc(
            B_dims, memory::data_type::f32, memory::format_tag::ab);
    auto binary_add_f32_md = memory::desc(
            binary_add_dims, memory::data_type::f32, memory::format_tag::ab);
    auto B_scales_f32_md = memory::desc(
            B_scales_dims, memory::data_type::f32, memory::format_tag::ab);
    auto D_f32_md = memory::desc(
            D_dims, memory::data_type::f32, memory::format_tag::ab);

    auto A_f32_mem = memory(A_f32_md, engine, A_user_data.data());
    auto B_f32_mem = memory(B_f32_md, engine, B_user_data.data());
    auto binary_add_f32_mem
            = memory(binary_add_f32_md, engine, binary_add_user_data.data());
    auto B_scales_f32_mem
            = memory(B_scales_f32_md, engine, B_scales_user_data.data());
    auto D_f32_mem = memory(D_f32_md, engine, D_user_data.data());

    // Create ukernel memories in requested data types.
    // Note that all formats are `ab`.
    auto A_md = memory::desc(A_dims, a_dt, memory::format_tag::ab);
    auto B_md = memory::desc(B_dims, b_dt, memory::format_tag::ab);
    auto binary_add_md = memory::desc(
            binary_add_dims, memory::data_type::f32, memory::format_tag::ab);
    auto B_scales_md = memory::desc(
            B_scales_dims, memory::data_type::f32, memory::format_tag::ab);
    auto C_md = memory::desc(C_dims, c_dt, memory::format_tag::ab);
    auto D_md = memory::desc(D_dims, d_dt, memory::format_tag::ab);

    auto A_mem = memory(A_md, engine);
    auto B_mem = memory(B_md, engine);
    auto binary_add_mem = memory(binary_add_md, engine);
    auto B_scales_mem = memory(B_scales_md, engine);
    auto C_mem = memory(C_md, engine);
    auto D_mem = memory(D_md, engine);

    const auto *A_ptr = reinterpret_cast<uint8_t *>(A_mem.get_data_handle());
    auto *B_ptr = reinterpret_cast<uint8_t *>(B_mem.get_data_handle());

    const size_t a_dt_size
            = memory::data_type_size(A_mem.get_desc().get_data_type());
    const size_t b_dt_size
            = memory::data_type_size(B_mem.get_desc().get_data_type());

    // Reorder user data into buffers passed to ukernels in target data types.
    reorder(A_f32_mem, A_mem).execute(engine_stream, A_f32_mem, A_mem);
    reorder(B_f32_mem, B_mem).execute(engine_stream, B_f32_mem, B_mem);
    reorder(binary_add_f32_mem, binary_add_mem)
            .execute(engine_stream, binary_add_f32_mem, binary_add_mem);
    reorder(B_scales_f32_mem, B_scales_mem)
            .execute(engine_stream, B_scales_f32_mem, B_scales_mem);
    reorder(D_f32_mem, D_mem).execute(engine_stream, D_f32_mem, D_mem);
    // Prepare C buffer. Needed to use a single ukernel in the example with
    // `set_add_C(true)`.
    // Note: to avoid this step, the first ukernel should run
    // `set_add_C(false)`, and it will initialize C buffer with intermediate
    // values.
    float *C_ptr = reinterpret_cast<float *>(C_mem.get_data_handle());
    for (memory::dim i = 0; i < M * N; i++) {
        C_ptr[i] = 0;
    }

    // Create ukernel post-ops (ReLU + Add).
    // It reuses `primitive_attr` abstraction.
    post_ops brgemm_ops;
    brgemm_ops.append_eltwise(
            algorithm::eltwise_relu, /* alpha = */ 0.f, /* beta = */ 0.f);
    brgemm_ops.append_binary(algorithm::binary_add, binary_add_md);

    // Create BRGeMM ukernel objects.
    // There are two objects:
    // * `brg` is the basic one which operates over K dimension divided into
    //   blocks. It utilizes `set_add_C(true)` to accumulate into the same
    //   buffer. It also uses `batch_size` to process as much as the number of
    //   blocks over K minus one.
    // * `brg_po` is the ukernel that would be called the last in the chain
    //   since it has attributes attached to the object and those will execute
    //   after all accumulation over K dimension is done.
    brgemm brg, brg_po;
    if (batch_size > 0) {
        // Construct a basic brgemm object.
        // `allow_empty` makes the interface to return an empty `brg` object
        // in case of the critical error.
        brg = brgemm(M, N, K_blk, batch_size, lda, ldb, ldc, a_dt, b_dt, c_dt,
                /* allow_empty = */ true);
        if (!brg) {
            printf("Error: brg object was not constructed.\n");
            return;
        }

        // Instruct the ukernel to append the result to the C tensor.
        brg.set_add_C(true);

        // Finalize the initialization.
        // Successful completion returns `true`. Otherwise, `brg` object can't
        // be used due to lack of support or non-compatible settings. The
        // specific reason may be found by using `ONEDNN_VERBOSE=all` env var.
        const bool ok = brg.finalize();
        if (!ok) {
            printf("Kernel is not supported on this platform.\n");
            return;
        }

        // Generate the executable code.
        brg.generate();
    }

    // Construct a brgemm object with post-ops.
    brg_po = brgemm(M, N, K_blk, 1, lda, ldb, ldc, a_dt, b_dt, c_dt,
            /* allow_empty = */ true);
    if (!brg_po) {
        printf("Error: brg_po object was not constructed.\n");
        return;
    }

    // Instruct the kernel to append the result to the C tensor computed by
    // `brg` ukernel.
    brg_po.set_add_C(true);
    // Specify post-ops.
    brg_po.set_post_ops(ldd, d_dt, brgemm_ops);
    // Specify quantization scales for B.
    if (b_dt == memory::data_type::s8 || b_dt == memory::data_type::u8) {
        brg_po.set_B_scales(/* mask = */ 2);
    }

    // Finalize the initialization.
    const bool ok = brg_po.finalize();
    if (!ok) {
        printf("Kernel is not supported on this platform.\n");
        return;
    }

    // Generate the executable code.
    brg_po.generate();

    // Query a scratchpad size and initialize a scratchpad buffer if the ukernel
    // is expecting it. This is a service space needed, has nothing in common
    // with accumulation buffer.
    size_t scratchpad_size = brg_po.get_scratchpad_size();
    std::vector<uint8_t> scratchpad(scratchpad_size);

    uint8_t *B_blocked = nullptr;
    void *B_base_ptr = B_ptr;
    size_t blocked_B_size = 0;

    // If packing is needed, create a dedicated object for data transformation.
    if (need_pack) {
        // Transform kernel for tensor B. The ukernel expects B passed in a
        // special VNNI format for low precision data types, e.g., bfloat16_t
        // or int8.
        // Note: the routine doesn't provide a `batch_size` argument in the
        // constructor as it can be either incorporated into `K` dimension, or
        // manually iterated over in a for-loop on the user side.
        transform pack_B(/* K = */ K_blk * n_calls, /* N = */ N,
                /* in_pack_type = */ pack_type::no_trans, /* in_ld = */ N,
                /* out_ld = */ ldb, /* in_dt = */ b_dt, /* out_dt = */ b_dt);

        // Size of the packed tensor.
        blocked_B_size = ldb * K_blk * memory::data_type_size(b_dt);

        B_blocked = new uint8_t[blocked_B_size * n_calls];
        B_base_ptr = B_blocked;

        // Generate the executable code.
        pack_B.generate();

        // Pack B routine execution.
        // Note: usually should be split to process only a part of B that the
        // ukernel will execute.
        pack_B.execute(B_ptr, B_blocked);
    }

    // ukernel execution section.
    //
    // Prepare buffers for execution.
    std::vector<std::pair<memory::dim, memory::dim>> A_B_offsets(batch_size);
    for (memory::dim i = 0; i < batch_size; i++) {
        const memory::dim A_offset_i = i * K_blk * a_dt_size;
        const memory::dim B_offset_i
                = need_pack ? i * blocked_B_size : i * N * K_blk * b_dt_size;
        A_B_offsets[i] = std::make_pair(A_offset_i, B_offset_i);
    }

    if (brg) {
        // A call to initialize hardware features. For example, prepare AMX
        // unit.
        brg.set_hw_context();

        // An execute call. `A_B_offsets` is a vector of pairs of offsets to A
        // and packed B tensors. `C_ptr` is a pointer to an accumulator buffer.
        brg.execute(A_ptr, B_base_ptr, A_B_offsets, C_ptr, scratchpad.data());
    }

    // Same set of operations for a ukernel with post-ops.
    std::vector<std::pair<memory::dim, memory::dim>> A_B_po_offsets;
    const memory::dim A_offset_po = batch_size * K_blk * a_dt_size;
    const memory::dim B_offset_po = need_pack
            ? batch_size * blocked_B_size
            : batch_size * N * K_blk * b_dt_size;
    A_B_po_offsets.emplace_back(A_offset_po, B_offset_po);

    // This object also requires this call since ukernel with post-ops may
    // require differently initialized internals underneath. If basic ukernel
    // was used and they share the same internals, this call will be optimized.
    brg_po.set_hw_context();

    // Prepare post-ops arguments and put them in a vector to make sure pointers
    // are sitting side by side.
    std::vector<const void *> bin_po_ptrs;
    bin_po_ptrs.push_back(binary_add_mem.get_data_handle());

    // Setting post-ops arguments into an attributes arguments storage.
    attr_params params;
    params.set_post_ops_args(bin_po_ptrs.data());
    params.set_B_scales(B_scales_mem.get_data_handle());

    // An execute call. The difference here is when post operations are
    // requested, an additional D tensor pointer to store final output result
    // after finishing accumulation and post-ops application is required.
    // Additionally, a special `params` object with post operations handles
    // is required.
    //
    // If post operations are not defined, the call is invalid, and a special
    // API checks its validity.
    if (brg_po.is_execute_postops_valid()) {
        brg_po.execute(A_ptr, B_base_ptr, A_B_po_offsets, C_ptr,
                D_mem.get_data_handle(), scratchpad.data(), params);
    } else {
        brg_po.execute(
                A_ptr, B_base_ptr, A_B_po_offsets, C_ptr, scratchpad.data());
    }

    // Once all computations are done and there are no more calls to ukernels
    // until they delegate control to the application, need to release the
    // hardware context.
    brgemm::release_hw_context();

    // Clean up an extra buffer.
    delete B_blocked;

    // Used for verification results, need unconditional reorder.
    auto user_D_mem = memory(D_f32_md, engine, D_data.data());
    reorder(D_mem, user_D_mem).execute(engine_stream, D_mem, user_D_mem);

    // Skip the check by default as data filling doesn't help with proper
    // verification of the result. Negative result doesn't necessarily mean
    // the functionality is broken. This is just a general sanity check.
    if (true) return;

    // A simplified fast verification that ukernel returned expected results.
    // Note: potential off-by-1 or 2 errors may pop up. This could be solved
    // with more sparse filling.
    bool to_throw = false;
    for (int m = 0; m < M; m++) {
        for (int n = 0; n < N; n++) {
            D_user_data[m * N + n] = 0;
            for (int k = 0; k < K; k++) {
                D_user_data[m * N + n]
                        += A_user_data[m * K + k] * B_user_data[k * N + n];
            }
            // B scales ref
            D_user_data[m * N + n] *= B_scales_user_data[n];
            // Relu post-op ref
            D_user_data[m * N + n] = std::max(D_user_data[m * N + n], 0.f);
            // Binary post-op ref
            D_user_data[m * N + n] += binary_add_user_data[0];

            const float diff
                    = fabsf(D_user_data[m * N + n] - D_data[m * N + n]);
            if (diff > 1.19e-7) {
                to_throw = true;
                if (true) {
                    printf("Error: [%3d:%3d] Ref:%12g Got:%12g Diff:%12g\n", m,
                            n, D_user_data[m * N + n], D_data[m * N + n], diff);
                }
            }
        }
    }
    if (to_throw) { throw status::runtime_error; }
}

int main(int argc, char **argv) {
    return handle_example_errors({dnnl::engine::kind::cpu}, brgemm_example);
}

Basic Concepts Data transformation