tests: benchdnn: update post-ops to support binary select op

avmanerikar · avmanerikar · commit 7db6ca1fab83 · 2025-02-03T10:27:32.000-08:00
diff --git a/tests/benchdnn/dnn_types.cpp b/tests/benchdnn/dnn_types.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2017-2024 Intel Corporation
+* Copyright 2017-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.
@@ -582,6 +582,10 @@ std::vector<std::pair<int, int>> attr_t::post_ops_t::get_po_masks(
         const auto &e = this->entry[idx];
         int mask = -1;
         int arg = DNNL_ARG_UNDEF;
+
+        int mask2 = -1;
+        int arg2 = DNNL_ARG_UNDEF;
+
         if (e.is_binary_kind()) {
             using mask_input_t = entry_t::binary_t::mask_input_t;
             auto mask_input = e.binary.mask_input;
@@ -590,6 +594,12 @@ std::vector<std::pair<int, int>> attr_t::post_ops_t::get_po_masks(
                     : policy2mask(
                             DNNL_ARG_SRC_1, e.binary.policy, prim_kind, ndims);
             arg = DNNL_ARG_SRC_1;
+
+            if (e.is_binary_kind_with_ternary_op()) {
+                mask2 = e.binary.mask;
+                arg2 = DNNL_ARG_SRC_2;
+            }
+
         } else if (e.is_prelu_kind()) {
             mask = attr_t::get_default_mask(e.prelu.policy);
             arg = DNNL_ARG_WEIGHTS;
@@ -599,6 +609,10 @@ std::vector<std::pair<int, int>> attr_t::post_ops_t::get_po_masks(
         assert(mask >= 0);
         v_masks.emplace_back(std::make_pair(
                 DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | arg, mask));
+
+        if (e.is_binary_kind_with_ternary_op())
+            v_masks.emplace_back(std::make_pair(
+                    DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | arg2, mask2));
     }
     return v_masks;
 }
@@ -630,11 +644,10 @@ bool attr_t::post_ops_t::entry_t::is_eltwise_kind() const {
     return kind > ELTWISE_START && kind < ELTWISE_END;
 }
 bool attr_t::post_ops_t::entry_t::is_binary_kind() const {
-    // binary select is a ternary operation and not currently
-    // supported in post-ops for the binary primitive
-    // TODO: add post-ops support for binary select operation
-    return kind > pk_t::BINARY_START && kind < pk_t::BINARY_END
-            && kind != pk_t::SELECT;
+    return kind > pk_t::BINARY_START && kind < pk_t::BINARY_END;
+}
+bool attr_t::post_ops_t::entry_t::is_binary_kind_with_ternary_op() const {
+    return kind == pk_t::SELECT;
 }
 bool attr_t::post_ops_t::entry_t::is_prelu_kind() const {
     return kind == PRELU;
@@ -1059,6 +1072,22 @@ int attr_args_t::prepare_post_ops_mds(const attr_t &attr, int ndims,
             mds.emplace((DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx)
                                 | po_rhs_tensor_entry.arg_attr_mask),
                     std::move(rhs_tensor_desc));
+
+            if (e.is_binary_kind_with_ternary_op()) {
+
+                const post_ops_rhs_tensor_entry_t
+                        post_op_rhs_select_tensor_entry
+                        = {e.binary.src2_dt, 0, e.binary.tag, DNNL_ARG_SRC_2};
+
+                auto rhs_select_tensor_desc = dnn_mem_t::init_md(ndims, dims,
+                        e.binary.src2_dt, post_op_rhs_select_tensor_entry.tag);
+
+                mds.emplace((DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx)
+                                    | post_op_rhs_select_tensor_entry
+                                              .arg_attr_mask),
+                        std::move(rhs_select_tensor_desc));
+            }
+
         } else if (e.is_convolution_kind()) {
             // Update dims for post operations appended after conv_dw
             conv_dw_fusion::get_fused_conv_dst_dims(ndims, e, dims, dims);
@@ -1162,9 +1191,16 @@ dnnl_primitive_attr_t create_dnnl_attr(
             } else if (e.is_binary_kind()) {
                 const auto &src1_md = attr_args.get_md(
                         (DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | DNNL_ARG_SRC_1));
+                const auto &src2_md = attr_args.get_md(
+                        (DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | DNNL_ARG_SRC_2));
                 assert(query_md_ndims(src1_md) != 0);
+
+                if (e.is_binary_kind_with_ternary_op())
+                    assert(query_md_ndims(src2_md) != 0);
+
                 DNN_SAFE_V(dnnl_post_ops_append_binary(
-                        ops, e.binary.alg, src1_md));
+                        ops, e.binary.alg, src1_md, src2_md));
+
             } else if (e.is_prelu_kind()) {
                 const auto &policy = e.prelu.policy;
                 const auto mask = attr_t::get_default_mask(policy);
@@ -1585,6 +1621,9 @@ float compute_eltwise_bwd(
 
 float compute_binary(pk_t kind, float src0, float src1, bool src2) {
     // don't compute on nan, propagate it
+
+    printf("src0: %f,\tsrc1: %f,\tsrc2: %d\n", src0, src1, src2);
+
     if (std::isnan(src0) || std::isnan(src1)) return NAN;
 
     if (kind == pk_t::ADD) {
@@ -1672,7 +1711,15 @@ void maybe_post_ops(const attr_t &attr, float &val, float sum_val,
             const auto &b = e.eltwise.beta;
             val = compute_eltwise_fwd(e.kind, val, a, b);
         } else if (e.is_binary_kind()) {
-            val = compute_binary(e.kind, val, *it_po, false);
+
+            auto src1_val = *it_po;
+            bool src2_val = false;
+
+            if (e.is_binary_kind_with_ternary_op()) {
+                it_po++;
+                src2_val = *it_po;
+            }
+            val = compute_binary(e.kind, val, src1_val, src2_val);
             it_po++;
         } else if (e.is_prelu_kind()) {
             val = val > 0 ? val : val * (*it_po);
diff --git a/tests/benchdnn/dnn_types.hpp b/tests/benchdnn/dnn_types.hpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2017-2024 Intel Corporation
+* Copyright 2017-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.
@@ -323,6 +323,7 @@ struct attr_t {
 
                 dnnl_alg_kind_t alg = dnnl_alg_kind_undef;
                 dnnl_data_type_t src1_dt = dnnl_data_type_undef;
+                dnnl_data_type_t src2_dt = dnnl_data_type_undef;
                 policy_t policy = policy_t::COMMON;
                 int64_t mask = -1;
                 mask_input_t mask_input = mask_input_t::none;
@@ -336,6 +337,7 @@ struct attr_t {
             bool is_convolution_kind() const;
             bool is_eltwise_kind() const;
             bool is_binary_kind() const;
+            bool is_binary_kind_with_ternary_op() const;
             bool is_prelu_kind() const;
         };
 
diff --git a/tests/benchdnn/dnnl_common.cpp b/tests/benchdnn/dnnl_common.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2017-2024 Intel Corporation
+* Copyright 2017-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.
@@ -1644,15 +1644,25 @@ int init_ref_memory_args_default_case(int exec_arg, dnn_mem_t &mem,
                     = exec_arg / DNNL_ARG_ATTR_MULTIPLE_POST_OP_BASE - 1;
             assert(bin_po_idx < attr.post_ops.len());
             const auto alg = attr.post_ops.entry[bin_po_idx].kind;
+            const bool is_src2_mem
+                    = ((exec_arg & DNNL_ARG_SRC_2) == DNNL_ARG_SRC_2);
+
             // Binary post-op filling.
-            fill_cfg_t def_binary_cfg(mem.dt(), -16.f, 16.f, /* int = */ true,
-                    alg, "def_binary_post_op");
-            const auto it = fill_cfg_map.find(DNNL_ARG_SRC_1);
-            const bool has_external_cfg = it != fill_cfg_map.end();
-            const fill_cfg_t &binary_fill_cfg
-                    = has_external_cfg ? (*it).second : def_binary_cfg;
-            TIME_FILL(SAFE(fill_random_real(mem, ref_mem, res, binary_fill_cfg),
-                    WARN));
+            if (!is_src2_mem
+                    || attr.post_ops.entry[bin_po_idx]
+                               .is_binary_kind_with_ternary_op()) {
+                const int f_min = is_src2_mem ? 0 : -16.f;
+                fill_cfg_t def_binary_cfg(mem.dt(), f_min, 16.f,
+                        /* int = */ true, alg, "def_binary_post_op");
+                const auto it = fill_cfg_map.find(
+                        is_src2_mem ? DNNL_ARG_SRC_2 : DNNL_ARG_SRC_1);
+                const bool has_external_cfg = it != fill_cfg_map.end();
+                const fill_cfg_t &binary_fill_cfg
+                        = has_external_cfg ? (*it).second : def_binary_cfg;
+                TIME_FILL(SAFE(
+                        fill_random_real(mem, ref_mem, res, binary_fill_cfg),
+                        WARN));
+            }
         } else if (exec_arg & DNNL_ARG_WEIGHTS) {
             // Prelu post-op filling.
             fill_cfg_t def_prelu_fill_cfg(mem.dt(), -2.f, 2.f, /* int = */ true,
diff --git a/tests/benchdnn/dnnl_common.hpp b/tests/benchdnn/dnnl_common.hpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2017-2024 Intel Corporation
+* Copyright 2017-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.
@@ -861,9 +861,13 @@ void init_memory_args(dnn_mem_map_t &mem_map, const prb_t *prb,
     for (int idx = 0; idx < dnnl_post_ops_len(const_po); ++idx) {
         if (dnnl_post_ops_get_kind(const_po, idx) != dnnl_binary) continue;
 
-        int po_arg = DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | DNNL_ARG_SRC_1;
-        const auto &po_md = query_md(const_pd, po_arg);
-        mem_map.emplace(po_arg, dnn_mem_t(po_md, test_engine));
+        int po_arg1 = DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | DNNL_ARG_SRC_1;
+        const auto &po_md1 = query_md(const_pd, po_arg1);
+        mem_map.emplace(po_arg1, dnn_mem_t(po_md1, test_engine));
+
+        int po_arg2 = DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | DNNL_ARG_SRC_2;
+        const auto &po_md2 = query_md(const_pd, po_arg2);
+        mem_map.emplace(po_arg2, dnn_mem_t(po_md2, test_engine));
     }
 
     // Prelu post-op.
diff --git a/tests/benchdnn/self/common.cpp b/tests/benchdnn/self/common.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2017-2024 Intel Corporation
+* Copyright 2017-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.
@@ -344,6 +344,9 @@ void append_binary(attr_t::post_ops_t &po, pk_t akind, dnnl_data_type_t src_dt1,
     attr_t::post_ops_t::entry_t e(akind);
     e.binary.alg = attr_t::post_ops_t::kind2dnnl_kind(akind);
     e.binary.src1_dt = src_dt1;
+
+    if (e.is_binary_kind_with_ternary_op()) e.binary.src2_dt = dnnl_s8;
+
     e.binary.mask_input = mask_input;
     e.binary.mask = mask;
     e.binary.policy = policy;
diff --git a/tests/benchdnn/utils/parser.cpp b/tests/benchdnn/utils/parser.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2019-2024 Intel Corporation
+* Copyright 2019-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.
@@ -217,6 +217,9 @@ attr_t::post_ops_t parse_attr_post_ops_func(const std::string &s) {
         } else if (e.is_binary_kind()) {
             const auto dt_str = get_substr(subs, subs_pos, ':');
             e.binary.src1_dt = str2dt(dt_str.c_str());
+
+            if (e.is_binary_kind_with_ternary_op()) e.binary.src2_dt = dnnl_s8;
+
             if (e.binary.src1_dt == dnnl_data_type_undef) {
                 BENCHDNN_PRINT(0, "%s \'%s\' %s\n",
                         "Error: binary post-op data type", dt_str.c_str(),
