onednn-src 0.1.13

/*******************************************************************************
* Copyright 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.
*******************************************************************************/

#ifndef CPU_MATMUL_REF_GROUPED_GEMM_HPP
#define CPU_MATMUL_REF_GROUPED_GEMM_HPP

#include "oneapi/dnnl/dnnl_config.h"

#if DNNL_EXPERIMENTAL_GROUPED_MEMORY

#include "common/c_types_map.hpp"
#include "common/memory_desc_wrapper.hpp"
#include "common/primitive.hpp"
#include "common/type_helpers.hpp"
#include "common/utils.hpp"

#include "cpu/matmul/cpu_matmul_pd.hpp"
#include "cpu/primitive_attr_postops.hpp"

namespace dnnl {
namespace impl {
namespace cpu {
namespace matmul {

struct ref_grouped_t : public primitive_t {
    struct pd_t : public cpu_matmul_pd_t {
        using cpu_matmul_pd_t::cpu_matmul_pd_t;

        DECLARE_COMMON_PD_T("ref_grouped:any", ref_grouped_t);

        // Weights are 3D: [G, K, N]
        // Override masks to include 0th expert dimension
        int wei_qmask_K() const { return (1 << 0) | (1 << 1); }

        int wei_qmask_N() const { return (1 << 0) | (1 << 2); }

        status_t init(engine_t *engine) {
            using namespace data_type;
            const auto src_type = src_md(0)->data_type;
            const auto wei_type = weights_md(0)->data_type;
            const auto dst_type = dst_md(0)->data_type;

            memory_desc_wrapper src_d(src_md());
            memory_desc_wrapper wei_d(weights_md(0));
            memory_desc_wrapper dst_d(dst_md());

            // Supported configurations: grouped src/dst, dense 3D weights
            VDISPATCH_MATMUL(src_d.is_grouped_desc() && dst_d.is_grouped_desc(),
                    VERBOSE_UNSUPPORTED_SPARSE_CFG);
            VDISPATCH_MATMUL(wei_d.is_blocking_desc() && wei_d.ndims() == 3,
                    VERBOSE_UNSUPPORTED_SPARSE_CFG);

            // Supported data types: fp and int for src/wei
            const bool is_fp_src = utils::one_of(
                    src_type, f32, bf16, f16, f8_e5m2, f8_e4m3, f4_e2m1);
            const bool is_fp_wei = utils::one_of(
                    wei_type, f32, bf16, f16, f8_e5m2, f8_e4m3, f4_e2m1);
            const bool is_int_src = utils::one_of(src_type, u8, s8);
            const bool is_int_wei = utils::one_of(wei_type, u8, s8, s4, u4);

            // Supported configurations: fp src + int wei (weight-only quantization),
            // int src + int wei, fp src + fp wei
            VDISPATCH_MATMUL(is_fp_src || is_int_src, VERBOSE_UNSUPPORTED_DT);
            VDISPATCH_MATMUL(is_fp_wei || is_int_wei, VERBOSE_UNSUPPORTED_DT);
            VDISPATCH_MATMUL(utils::one_of(dst_type, f32, bf16, f16),
                    VERBOSE_UNSUPPORTED_DT);
            VDISPATCH_MATMUL(IMPLICATION(is_int_src, is_int_wei),
                    VERBOSE_UNSUPPORTED_DT_CFG);
            // WOQ requires weight scales and fpmath with apply_to_int
            VDISPATCH_MATMUL(IMPLICATION(is_fp_src && is_int_wei,
                                     !attr()->scales_.has_default_values(
                                             DNNL_ARG_WEIGHTS)),
                    VERBOSE_UNSUPPORTED_DT_CFG);
            VDISPATCH_MATMUL(IMPLICATION(is_fp_src && is_int_wei,
                                     attr()->fpmath_.apply_to_int_),
                    VERBOSE_UNSUPPORTED_DT_CFG);

            // Check for supported quantization schemes
            const auto &attr_scales = attr()->scales_;
            if (!attr_scales.has_default_values(DNNL_ARG_SRC)) {
                const int src_mask = attr_scales.get_mask(DNNL_ARG_SRC);
                const int rowwise_mask = src_qmask_M();
                const int blocked_mask = src_qmask_M() | src_qmask_K();
                // Allow row-wise or blocked (K-grouping) scales for src
                VDISPATCH_MATMUL(
                        src_mask == rowwise_mask || src_mask == blocked_mask,
                        VERBOSE_UNSUPPORTED_SCALES_CFG);
                VDISPATCH_MATMUL(
                        utils::one_of(attr_scales.get_data_type(DNNL_ARG_SRC),
                                f32, bf16, f16, e8m0, f8_e4m3, f8_e5m2),
                        VERBOSE_UNSUPPORTED_SCALES_CFG);
                if (!attr_scales.get(DNNL_ARG_SRC).has_default_groups()) {
                    // K-grouped src scales supported with int and fp types
                    VDISPATCH_MATMUL(is_int_src || is_fp_src,
                            VERBOSE_UNSUPPORTED_SCALES_CFG);
                    const auto gM = attr_scales.get_group(DNNL_ARG_SRC, -2);
                    VDISPATCH_MATMUL(gM == 1, VERBOSE_UNSUPPORTED_SCALES_CFG);
                    const auto gK = attr_scales.get_group(DNNL_ARG_SRC, -1);
                    VDISPATCH_MATMUL(gK > 1, VERBOSE_UNSUPPORTED_SCALES_CFG);
                    VDISPATCH_MATMUL(
                            K() % gK == 0, VERBOSE_UNSUPPORTED_SCALES_CFG);
                }
            }
            if (!attr_scales.has_default_values(DNNL_ARG_WEIGHTS)) {
                const int wei_mask = attr_scales.get_mask(DNNL_ARG_WEIGHTS);
                const int colwise_mask = wei_qmask_N();
                const int blocked_mask = wei_qmask_K() | wei_qmask_N();
                // Allow column-wise or blocked (K grouping) scales for weights
                VDISPATCH_MATMUL(
                        utils::one_of(
                                attr_scales.get_data_type(DNNL_ARG_WEIGHTS),
                                f32, bf16, f16, e8m0, f8_e4m3, f8_e5m2),
                        VERBOSE_UNSUPPORTED_SCALES_CFG);
                VDISPATCH_MATMUL(
                        wei_mask == colwise_mask || wei_mask == blocked_mask,
                        VERBOSE_UNSUPPORTED_SCALES_CFG);
                if (!attr_scales.get(DNNL_ARG_WEIGHTS).has_default_groups()) {
                    // K-grouped wei scales supported with int and fp types
                    VDISPATCH_MATMUL(is_int_wei || is_fp_wei,
                            VERBOSE_UNSUPPORTED_SCALES_CFG);
                    const auto gK = attr_scales.get_group(DNNL_ARG_WEIGHTS, -2);
                    VDISPATCH_MATMUL(gK > 1, VERBOSE_UNSUPPORTED_SCALES_CFG);
                    VDISPATCH_MATMUL(
                            K() % gK == 0, VERBOSE_UNSUPPORTED_SCALES_CFG);
                    const auto gN = attr_scales.get_group(DNNL_ARG_WEIGHTS, -1);
                    VDISPATCH_MATMUL(gN == 1, VERBOSE_UNSUPPORTED_SCALES_CFG);
                }
            }
            VDISPATCH_MATMUL(attr_scales.has_default_values(DNNL_ARG_DST),
                    VERBOSE_UNSUPPORTED_SCALES_CFG);

            // Zero-points are supported for wei: for WOQ (fp src) and for int arithmetic (int src)
            const auto &attr_zps = attr()->zero_points_;
            VDISPATCH_MATMUL(attr_zps.has_default_values(DNNL_ARG_SRC),
                    VERBOSE_UNSUPPORTED_ZP_CFG);
            VDISPATCH_MATMUL(attr_zps.has_default_values(DNNL_ARG_DST),
                    VERBOSE_UNSUPPORTED_ZP_CFG);

            // Allow column-wise or blocked (K grouping) zps for weights
            if (!attr_zps.has_default_values(DNNL_ARG_WEIGHTS)) {
                VDISPATCH_MATMUL(is_int_wei, VERBOSE_UNSUPPORTED_ZP_CFG);
                VDISPATCH_MATMUL(
                        utils::one_of(attr_zps.get_data_type(DNNL_ARG_WEIGHTS),
                                u8, s8, u4, s4, s32),
                        VERBOSE_UNSUPPORTED_ZP_CFG);
                const int zp_mask = attr_zps.get_mask(DNNL_ARG_WEIGHTS);
                const int colwise_mask = wei_qmask_N();
                const int blocked_mask = wei_qmask_K() | wei_qmask_N();
                VDISPATCH_MATMUL(
                        zp_mask == colwise_mask || zp_mask == blocked_mask,
                        VERBOSE_UNSUPPORTED_ZP_CFG);
                if (!attr_zps.get(DNNL_ARG_WEIGHTS).has_default_groups()) {
                    const auto gK = attr_zps.get_group(DNNL_ARG_WEIGHTS, -2);
                    VDISPATCH_MATMUL(gK > 1, VERBOSE_UNSUPPORTED_ZP_CFG);
                    VDISPATCH_MATMUL(K() % gK == 0, VERBOSE_UNSUPPORTED_ZP_CFG);
                    const auto gN = attr_zps.get_group(DNNL_ARG_WEIGHTS, -1);
                    VDISPATCH_MATMUL(gN == 1, VERBOSE_UNSUPPORTED_ZP_CFG);
                }
            }

            // for src/wei scales group size in case of K grouping,
            // one must be a multiple of the other
            if (!attr_scales.has_default_values(DNNL_ARG_SRC)
                    && !attr_scales.get(DNNL_ARG_SRC).has_default_groups()
                    && !attr_scales.has_default_values(DNNL_ARG_WEIGHTS)
                    && !attr_scales.get(DNNL_ARG_WEIGHTS)
                                .has_default_groups()) {
                const auto src_gK = attr_scales.get_group(DNNL_ARG_SRC, -1);
                const auto wei_gK = attr_scales.get_group(DNNL_ARG_WEIGHTS, -2);
                VDISPATCH_MATMUL(src_gK % wei_gK == 0 || wei_gK % src_gK == 0,
                        VERBOSE_UNSUPPORTED_SCALES_CFG);
            }

            // Scales and ZPs groups must match
            if (!attr_scales.has_default_values(DNNL_ARG_WEIGHTS)
                    && !attr_zps.has_default_values(DNNL_ARG_WEIGHTS)) {
                const auto scale_gK
                        = attr_scales.get_group(DNNL_ARG_WEIGHTS, -2);
                const auto zp_gK = attr_zps.get_group(DNNL_ARG_WEIGHTS, -2);
                VDISPATCH_MATMUL(scale_gK == zp_gK, VERBOSE_INCONSISTENT_DIM,
                        "wei_scale_group_k", (int)scale_gK, "wei_zp_group_k",
                        (int)zp_gK);
            }

            // Post-ops: only binary mul is supported (for NVFP4 global scale)
            const auto &po = attr()->post_ops_;
            bool po_ok = true;
            for (int i = 0; i < po.len(); ++i) {
                const auto &e = po.entry_[i];
                po_ok = po_ok && e.is_binary()
                        && e.binary.alg == alg_kind::binary_mul;
            }
            VDISPATCH_MATMUL(po_ok, VERBOSE_UNSUPPORTED_POSTOP);
            VDISPATCH_MATMUL(attr_.post_ops_.set_default_formats(dst_md(0))
                            == status::success,
                    VERBOSE_UNSUPPORTED_POSTOP);

            return status::success;
        }
    };

    ref_grouped_t(const pd_t *apd) : primitive_t(apd) {}

    status_t init(engine_t *engine) override {
        ref_post_ops
                = utils::make_unique<ref_post_ops_t>(pd()->attr()->post_ops_);
        if (!ref_post_ops) return status::out_of_memory;
        CHECK(ref_post_ops->init(pd()->dst_md()));
        return status::success;
    }

    status_t execute(const exec_ctx_t &ctx) const override;

private:
    const pd_t *pd() const { return (const pd_t *)primitive_t::pd().get(); }
    std::unique_ptr<ref_post_ops_t> ref_post_ops;
};

} // namespace matmul
} // namespace cpu
} // namespace impl
} // namespace dnnl

#endif // DNNL_EXPERIMENTAL_GROUPED_MEMORY
#endif // CPU_MATMUL_REF_GROUPED_GEMM_HPP