Repository: uxlfoundation/oneDNN
Pull request #2980: benchdnn: graph: support input displace for scalar
Author: login: wzt1997
Labels: name: component:graph-api
description: Codeowner: @oneapi-src/onednn-graph
name: component:tests
description: Codeowner: @oneapi-src/onednn-arch
Created: 2025-03-31T02:31:41Z | Merged: 2025-04-11T07:02:44Z

## General

To support bottom-right casual mask( [RFC ](https://github.com/uxlfoundation/oneDNN/blob/yixin/rfc/graph-api-bottom-right-causal-mask/rfcs/20250314-graph-api-bottom-right-causal-mask/README.md)), Graph API will introduce scalar logical tensor. In current SDPA scenarios, the scalar values are relevant with sequence lengths, hence it requires input displacement with shape information.

The work in the PR includes: 
1. Support scalar displacement for `Add`->`Sub`->`GE` operation chain.
2. Resolve the TODO( [link ](https://github.com/uxlfoundation/oneDNN/blob/633978a82ec14ccc01a9cccaf173e89c25cbbc89/tests/benchdnn/graph/input_displacer.hpp#L37)) in input displacer by moving `fill_cfg` into the displacement arguments.


Diff:
diff --git a/tests/benchdnn/graph/input_displacer.cpp b/tests/benchdnn/graph/input_displacer.cpp
index 72c2828c20a..a029e5f2e9e 100644
--- a/tests/benchdnn/graph/input_displacer.cpp
+++ b/tests/benchdnn/graph/input_displacer.cpp
@@ -73,12 +73,24 @@ partition_data_displacer_t::partition_data_displacer_t(
                     lt = &parent_op->in_lts_[0]) {
                 parent_op = &dg_->get_op_by_out_lt(lt->id_);
                 if (parent_op->empty()) {
-                    if (aop.kind_ == "Divide" || aop.kind_ == "Multiply") {
+                    if (aop.kind_ == "Divide") {
+                        // Division has values > 1.f to reduce final values.
+                        static const std::vector<float> user_set {
+                                2.f, 4.f, 8.f};
                         // There's a special case for Divide, when second (user)
                         // input should be displaced with power-of-2 values.
-                        quantize_displace_.emplace(lt->id_,
-                                std::make_tuple(
-                                        aop, i, *lt, filling_type_t::pow2));
+                        displace_args_.emplace(lt->id_,
+                                displace_args_t {aop, i, *lt,
+                                        filling_type_t::fixed_setting,
+                                        {user_set, "Div displacer"}});
+                    } else if (aop.kind_ == "Multiply") {
+                        // Multiplication has values <= 1.f to reduce final values.
+                        static const std::vector<float> user_set {
+                                0.25f, 0.5f, 1.f};
+                        displace_args_.emplace(lt->id_,
+                                displace_args_t {aop, i, *lt,
+                                        filling_type_t::fixed_setting,
+                                        {user_set, "Mul displacer"}});
                     }
                     break;
                 }
@@ -105,13 +117,12 @@ partition_data_displacer_t::partition_data_displacer_t(
                                         == f8_main_op_kind.end()))
                             break;
 
-                        quantize_displace_.emplace(parent_op_in_lt.id_,
-                                std::make_tuple(aop, i, parent_op_in_lt,
-                                        filling_type_t::quantization));
+                        displace_args_.emplace(parent_op_in_lt.id_,
+                                displace_args_t {aop, i, parent_op_in_lt,
+                                        filling_type_t::quantization});
                         break;
                     }
-                }
-                if (parent_op->kind_ == "StaticReshape") {
+                } else if (parent_op->kind_ == "StaticReshape") {
                     // StaticReshape is accepted when the pattern is
                     // "StaticReshape + Matmul" and it doesn't have any
                     // predecessors in the partition
@@ -122,9 +133,9 @@ partition_data_displacer_t::partition_data_displacer_t(
                             || op_ids_set_.find(prev_parent_op.id_)
                                     == op_ids_set_.end()) {
                         if (aop.kind_ == "MatMul") {
-                            quantize_displace_.emplace(parent_op_in_lt.id_,
-                                    std::make_tuple(aop, i, parent_op_in_lt,
-                                            filling_type_t::quantization));
+                            displace_args_.emplace(parent_op_in_lt.id_,
+                                    displace_args_t {aop, i, parent_op_in_lt,
+                                            filling_type_t::quantization});
                         }
                         break;
                     }
@@ -203,6 +214,8 @@ partition_data_displacer_t::partition_data_displacer_t(
             if (!causal_mask_lt) break;
 
             filling_type_t filling_type = filling_type_t::undef;
+            std::string cfg_name;
+            float user_set_value = 0.f;
             if (aop.kind_ == "Add") {
                 const auto ndims = causal_mask_lt->shape_.size();
                 if (ndims < 2) {
@@ -219,7 +232,8 @@ partition_data_displacer_t::partition_data_displacer_t(
                     // are computed. To avoid numerical instabilities, a zero
                     // mask can be applied without compromising validation
                     // capabilities.
-                    filling_type = filling_type_t::zero;
+                    filling_type = filling_type_t::fixed_setting;
+                    cfg_name = "Explicit_padding_mask";
                 } else {
                     // This is a look-ahead (or causal) mask case, when future
                     // tokens (row < col) are set to infinity to remove all
@@ -229,17 +243,21 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -INFINITY;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
@@ -247,11 +265,66 @@ partition_data_displacer_t::partition_data_displacer_t(
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {{user_set_value}, cfg_name}});
+            } else if (filling_type == filling_type_t::causal_mask) {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
+            break;
+        }
+
+        // Fill proper data for bottom-right implicit casual mask
+        while (aop.kind_ == "Add") {
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+            static constexpr int seq_len_q = 0;
+            static constexpr int seq_len_kv = 1;
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_seq_len_displace_args =
+                    [&](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == seq_len_q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {{user_set_value}, cfg_name}});
+                            }
+                        }
+                    };
+
+            set_seq_len_displace_args(&aop, seq_len_q);
+            set_seq_len_displace_args(child_op, seq_len_kv);
             break;
         }
     }
@@ -264,16 +337,16 @@ int partition_data_displacer_t::displace_input_data(
         return FAIL;
     }
 
-    if (quantize_displace_.find(lt_id) == quantize_displace_.end()) {
+    if (displace_args_.find(lt_id) == displace_args_.end()) {
         // no need to displace the data of this tensor
         return OK;
     }
-    const displace_t &displace = quantize_displace_.at(lt_id);
-
-    const auto &main_op = ::std::get<0>(displace);
-    auto main_op_offset = ::std::get<1>(displace);
-    auto tensor = ::std::get<2>(displace);
-    const auto filling_type = ::std::get<3>(displace);
+    const displace_args_t &d_args = displace_args_.at(lt_id);
+    const auto &main_op = d_args.main_op_;
+    const auto &main_op_offset = d_args.main_op_offset_;
+    const auto &tensor = d_args.tensor_;
+    const auto &fill_cfg = d_args.fill_cfg_;
+    const auto filling_type = d_args.filling_type_;
 
     auto opkind = opstr2kind(main_op.kind_);
     int main_op_arg = get_prim_arg_name_from_graph_op_input_offset(
@@ -287,38 +360,9 @@ int partition_data_displacer_t::displace_input_data(
         SAFE(gen_quantize_filling(
                      main_op, main_op_arg, mem_replace, tensor.data_type_, res),
                 WARN);
-    } else if (filling_type == filling_type_t::pow2) {
-        // This custom displace serves a shrinking the output purpose. Values
-        // are picked in the way to have more chances of ending with exact lower
-        // data type representative in the output, which in turn decreases the
-        // number of points with non-zero absolute diff.
-
-        // Division has values > 1.f to reduce final values.
-        static const std::vector<float> pow2_div_vals {2.f, 4.f, 8.f};
-        // Multiplication has values <= 1.f to reduce final values.
-        static const std::vector<float> pow2_mul_vals {0.25f, 0.5f, 1.f};
-        // Guard set.
-        static const std::vector<float> dummy {};
-
-        const bool is_div = main_op.kind_ == "Divide";
-        const bool is_mul = main_op.kind_ == "Multiply";
-        const auto &user_set
-                = is_div ? pow2_div_vals : (is_mul ? pow2_mul_vals : dummy);
-        fill_cfg_t fill_cfg(user_set, "Mul/Div displacer");
-        SAFE(gen_fixed_set_filling(mem_replace, mem.md_, fill_cfg, res), WARN);
     } else if (filling_type == filling_type_t::causal_mask) {
         SAFE(gen_causal_mask_filling(mem_replace, mem.md_, res), WARN);
-    } else if (filling_type == filling_type_t::minus_infinity) {
-        static const std::vector<float> user_set {-INFINITY};
-        fill_cfg_t fill_cfg(user_set, "Implicit_causal_mask");
-        SAFE(gen_fixed_set_filling(mem_replace, mem.md_, fill_cfg, res), WARN);
-    } else if (filling_type == filling_type_t::zero) {
-        static const std::vector<float> user_set {0.f};
-        fill_cfg_t fill_cfg(user_set, "Explicit_padding_mask");
-        SAFE(gen_fixed_set_filling(mem_replace, mem.md_, fill_cfg, res), WARN);
-    } else if (filling_type == filling_type_t::one) {
-        static const std::vector<float> user_set {1.f};
-        fill_cfg_t fill_cfg(user_set, "Explicit_padding_mask");
+    } else if (filling_type == filling_type_t::fixed_setting) {
         SAFE(gen_fixed_set_filling(mem_replace, mem.md_, fill_cfg, res), WARN);
     } else {
         assert(!"unexpected filling type");
@@ -400,8 +444,7 @@ int partition_data_displacer_t::displace_input_data(
 
         mem_replace = ::std::move(
                 const_cast<dnn_mem_t &>(ref_prim.get_arg(DNNL_ARG_DST)));
-        tensor = op.out_lts_[0];
-        parent_op = &dg_->get_op_by_out_lt(tensor.id_);
+        parent_op = &dg_->get_op_by_out_lt(op.out_lts_[0].id_);
     }
 
     if (backward_path_launched) {
diff --git a/tests/benchdnn/graph/input_displacer.hpp b/tests/benchdnn/graph/input_displacer.hpp
index ec273f62a13..6a870f83787 100644
--- a/tests/benchdnn/graph/input_displacer.hpp
+++ b/tests/benchdnn/graph/input_displacer.hpp
@@ -27,29 +27,31 @@ namespace graph {
 enum class filling_type_t {
     undef = 0,
     quantization,
-    // Floating-point power-of-2 values for precise disivision/multiplication.
-    pow2,
-    // Explicit causal mask from SDPA pattern.
     causal_mask,
-    // Implicit causal mask free input.
-    minus_infinity,
-    // Explicit padding mask (1D case) through the Add op from SDPA pattern.
-    zero,
-    // TODO: `pow2`, `minus_infinity`, `zero` and `one` types can be replaced
-    // with the one defined by non empty fill_cfg which will become a member
-    // of displacer.
-    // Explicit padding mask (1D case) through the Select op from SDPA pattern.
-    one,
+    // Fill pre-defined fixed values for data filling, such as 0, 1, -inf, and
+    // specified shape information for scalar input.
+    fixed_setting,
 };
+struct displace_args_t {
 
-// tuple<
-//     main op,
-//     main op offset,
-//     the tensor as a displace starting point,
-//     filling_type
-// >
-using displace_t = ::std::tuple<::graph::deserialized_op_t, size_t,
-        ::graph::deserialized_lt_t, filling_type_t>;
+public:
+    displace_args_t() = default;
+    displace_args_t(const deserialized_op_t &op, size_t offset,
+            const deserialized_lt_t &lt, filling_type_t type,
+            fill_cfg_t cfg = {})
+        : main_op_(op)
+        , main_op_offset_(offset)
+        , tensor_(lt)
+        , filling_type_(type)
+        , fill_cfg_(cfg) {}
+
+    deserialized_op_t main_op_;
+    size_t main_op_offset_;
+    //the tensor as a displace starting point
+    deserialized_lt_t tensor_;
+    filling_type_t filling_type_;
+    fill_cfg_t fill_cfg_;
+};
 
 class partition_data_displacer_t {
 public:
@@ -63,7 +65,7 @@ class partition_data_displacer_t {
     // A set of op_id values from a partition came to a displacer. Used to
     // identify at displacement stage if Deq is the starting point or not.
     std::unordered_set<size_t> op_ids_set_;
-    ::std::unordered_map<size_t, displace_t> quantize_displace_;
+    ::std::unordered_map<size_t, displace_args_t> displace_args_;
 
     int gen_quantize_filling(const ::graph::deserialized_op_t &main_op, int arg,
             dnn_mem_t &mem, const ::std::string &dt, res_t *res);


Reviews:
user: login: dzarukin
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T17:53:58Z
user: login: ElaineBao
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T02:01:52Z
user: login: wzt1997
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-02T02:54:06Z
user: login: ElaineBao
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T03:31:09Z
user: login: wzt1997
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-02T05:42:45Z
user: login: ElaineBao
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T05:49:05Z
user: login: ElaineBao
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T15:55:16Z
user: login: dzarukin
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T16:36:40Z

Review comments:
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == SEQ_LEN_Q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                const std::vector<float> user_set {
+                                        user_set_value};
+
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {user_set, cfg_name}});
+                            }
+                        }
+                    };
+
+            set_scalar_displace_args(&aop, SEQ_LEN_Q);
+            set_scalar_displace_args(child_op, SEQ_LEN_KV);
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: Please use constexpr int constants in here instead macro in a header file for better readability.
created_at: 2025-04-01T17:46:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -264,16 +343,16 @@ int partition_data_displacer_t::displace_input_data(
         return FAIL;
     }
 
-    if (quantize_displace_.find(lt_id) == quantize_displace_.end()) {
+    if (displace_args_.find(lt_id) == displace_args_.end()) {
         // no need to displace the data of this tensor
         return OK;
     }
-    const displace_t &displace = quantize_displace_.at(lt_id);
-
-    const auto &main_op = ::std::get<0>(displace);
-    auto main_op_offset = ::std::get<1>(displace);
-    auto tensor = ::std::get<2>(displace);
-    const auto filling_type = ::std::get<3>(displace);
+    const displace_args_t &d_args = displace_args_.at(lt_id);
+    const auto &main_op = d_args.main_op_;
+    auto main_op_offset = d_args.main_op_offset_;
+    auto tensor = d_args.tensor_;
+    auto &fill_cfg = d_args.fill_cfg_;
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: Two members above seem to fit `const` category as well.
```suggestion
    const auto &fill_cfg = d_args.fill_cfg_;
```
created_at: 2025-04-01T17:47:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -203,6 +217,8 @@ partition_data_displacer_t::partition_data_displacer_t(
             if (!causal_mask_lt) break;
 
             filling_type_t filling_type = filling_type_t::undef;
+            std::string cfg_name {};
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: ```suggestion
            std::string cfg_name;
```
created_at: 2025-04-01T17:49:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: Why change to a specific value?
created_at: 2025-04-01T17:49:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: I'd do `else if` to avoid random errors.
created_at: 2025-04-01T17:50:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == SEQ_LEN_Q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                const std::vector<float> user_set {
+                                        user_set_value};
+
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {user_set, cfg_name}});
+                            }
+                        }
+                    };
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: Is host-scalar only identified by the pattern and not by the property of Op?
created_at: 2025-04-01T17:53:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -74,11 +74,26 @@ partition_data_displacer_t::partition_data_displacer_t(
                 parent_op = &dg_->get_op_by_out_lt(lt->id_);
                 if (parent_op->empty()) {
                     if (aop.kind_ == "Divide" || aop.kind_ == "Multiply") {
+                        // Division has values > 1.f to reduce final values.
+                        static const std::vector<float> pow2_div_vals {
+                                2.f, 4.f, 8.f};
+                        // Multiplication has values <= 1.f to reduce final values.
+                        static const std::vector<float> pow2_mul_vals {
+                                0.25f, 0.5f, 1.f};
+                        // Guard set.
+                        static const std::vector<float> dummy {};
+                        const bool is_div = aop.kind_ == "Divide";
+                        const bool is_mul = aop.kind_ == "Multiply";
+                        const auto &user_set = is_div
+                                ? pow2_div_vals
+                                : (is_mul ? pow2_mul_vals : dummy);
path: tests/benchdnn/graph/input_displacer.cpp
user: login: ElaineBao
body: nit: is it clearer to just split the data filling for Divide and Multiply, so we don't need to check op_kind again
created_at: 2025-04-02T01:52:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
path: tests/benchdnn/graph/input_displacer.cpp
user: login: ElaineBao
body: It seems to be an unintentional change, please remove
created_at: 2025-04-02T01:53:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
path: tests/benchdnn/graph/input_displacer.cpp
user: login: ElaineBao
body: How does this condition co-work with the same check of Add's op kind in L222: https://github.com/uxlfoundation/oneDNN/blob/aa6722e9137ea1c4b4b62926ce4929dd6a34aa02/tests/benchdnn/graph/input_displacer.cpp#L222C15-L222C38
created_at: 2025-04-02T02:00:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -74,11 +74,26 @@ partition_data_displacer_t::partition_data_displacer_t(
                 parent_op = &dg_->get_op_by_out_lt(lt->id_);
                 if (parent_op->empty()) {
                     if (aop.kind_ == "Divide" || aop.kind_ == "Multiply") {
+                        // Division has values > 1.f to reduce final values.
+                        static const std::vector<float> pow2_div_vals {
+                                2.f, 4.f, 8.f};
+                        // Multiplication has values <= 1.f to reduce final values.
+                        static const std::vector<float> pow2_mul_vals {
+                                0.25f, 0.5f, 1.f};
+                        // Guard set.
+                        static const std::vector<float> dummy {};
+                        const bool is_div = aop.kind_ == "Divide";
+                        const bool is_mul = aop.kind_ == "Multiply";
+                        const auto &user_set = is_div
+                                ? pow2_div_vals
+                                : (is_mul ? pow2_mul_vals : dummy);
path: tests/benchdnn/graph/input_displacer.cpp
user: login: wzt1997
body: Good suggestion. I split them to make it more clear.
created_at: 2025-04-02T02:29:30Z
author_association: NONE
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
path: tests/benchdnn/graph/input_displacer.cpp
user: login: wzt1997
body: Thanks for reminding. The change is removed.
created_at: 2025-04-02T02:29:52Z
author_association: NONE
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
path: tests/benchdnn/graph/input_displacer.cpp
user: login: wzt1997
body: As we have a check in [L165](https://github.com/uxlfoundation/oneDNN/blob/aa6722e9137ea1c4b4b62926ce4929dd6a34aa02/tests/benchdnn/graph/input_displacer.cpp#L165) that handles the op chain `Add`/`Select`->`Softmax`. This part handles the input displacement for:
1. Explicit casual mask with Add/Softmax.
2. `-INF` filling for Softmax with implicit casual mask.
( NOTE: for explicit casual mask, the other input of Select( besides MM1 ouput ) will not be handled specially as MM1  output will always be chosen. )

While for bottom-right implicit casual mask, it contains an op chain `Add`->`Sub`->`GE`->`Select`, so it will not interfere with the above handling.
created_at: 2025-04-02T02:48:58Z
author_association: NONE
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == SEQ_LEN_Q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                const std::vector<float> user_set {
+                                        user_set_value};
+
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {user_set, cfg_name}});
+                            }
+                        }
+                    };
+
+            set_scalar_displace_args(&aop, SEQ_LEN_Q);
+            set_scalar_displace_args(child_op, SEQ_LEN_KV);
path: tests/benchdnn/graph/input_displacer.cpp
user: login: wzt1997
body: Changed. Thanks.
created_at: 2025-04-02T02:49:10Z
author_association: NONE
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == SEQ_LEN_Q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                const std::vector<float> user_set {
+                                        user_set_value};
+
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {user_set, cfg_name}});
+                            }
+                        }
+                    };
path: tests/benchdnn/graph/input_displacer.cpp
user: login: wzt1997
body: For current request, it can be identified in both ways. For checking with porperty, it can only be supported after Yixin's update for bottom-right implicit casual mask is merged.
created_at: 2025-04-02T02:51:52Z
author_association: NONE
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == SEQ_LEN_Q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                const std::vector<float> user_set {
+                                        user_set_value};
+
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {user_set, cfg_name}});
+                            }
+                        }
+                    };
path: tests/benchdnn/graph/input_displacer.cpp
user: login: ElaineBao
body: I think checking pattern here can cover two scenarios, first is s_kv and s_q are passed as host_scalar, second is s_kv and s_q are passed as device memory. @wzt1997 can you confirm the implementation applies to two scenarios?
created_at: 2025-04-02T03:31:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +246,91 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -1e30;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
+                        const size_t ndims = op->out_lts_[0].shape_.size();
+                        const size_t seq_len_idx = (which_seq_len == SEQ_LEN_Q)
+                                ? ndims - 2
+                                : ndims - 1;
+
+                        for (size_t i = 0; i < op->in_lts_.size(); i++) {
+                            auto *parent_op = &dg_->get_op_by_out_lt(
+                                    op->in_lts_[i].id_);
+                            // For add->sub->ge, we consider the inputs of add
+                            // and sub as scalars if they have no parent
+                            // tensors.
+                            if (parent_op->empty()) {
+                                float user_set_value = static_cast<float>(
+                                        op->in_lts_[1 - i].shape_[seq_len_idx]);
+                                const std::vector<float> user_set {
+                                        user_set_value};
+
+                                displace_args_.emplace(op->in_lts_[i].id_,
+                                        displace_args_t {*op, i, op->in_lts_[i],
+                                                filling_type_t::fixed_setting,
+                                                {user_set, cfg_name}});
+                            }
+                        }
+                    };
path: tests/benchdnn/graph/input_displacer.cpp
user: login: wzt1997
body: Verified that checking through pattern can support both scenarios: with/without `host_scalar` property.
created_at: 2025-04-02T05:42:44Z
author_association: NONE
diff_hunk: @@ -229,29 +243,92 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -INFINITY;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else if (filling_type == filling_type_t::causal_mask) {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
+            break;
+        }
+
+        // Fill scalar for bottom-right implicit casual mask
+        if (aop.kind_ == "Add") {
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+            static constexpr int seq_len_q = 0;
+            static constexpr int seq_len_kv = 1;
+
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_scalar_displace_args =
path: tests/benchdnn/graph/input_displacer.cpp
user: login: ElaineBao
body: nit: as this pattern doesn't necessarily have to be scalars, please change the variable name here.
created_at: 2025-04-02T05:48:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +243,92 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -INFINITY;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: It would be better to fill a vector above instead of a single value and pass it directly to the `emplace`.
created_at: 2025-04-10T16:30:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,29 +243,92 @@ partition_data_displacer_t::partition_data_displacer_t(
             } else if (aop.kind_ == "Select") {
                 if (select_cond_has_parent) {
                     // Implicit causal mask case.
-                    filling_type = filling_type_t::minus_infinity;
+                    filling_type = filling_type_t::fixed_setting;
+                    user_set_value = -INFINITY;
+                    cfg_name = "Implicit_causal_mask";
                 } else {
                     // Padding mask.
                     assert(qk_data_offset == 1 || qk_data_offset == 2);
                     // Fill condition depending on qk values tensor to use only
                     // its values, which is equivalent of not using a mask.
+                    filling_type = filling_type_t::fixed_setting;
                     if (qk_data_offset == 1) {
-                        filling_type = filling_type_t::one;
+                        user_set_value = 1.f;
                     } else if (qk_data_offset == 2) {
-                        filling_type = filling_type_t::zero;
+                        user_set_value = 0.f;
                     }
+                    cfg_name = "Explicit_padding_mask";
                 }
             }
 
             if (filling_type == filling_type_t::undef) {
                 BENCHDNN_PRINT(
                         7, "%s\n", "[DISPLACE]: Filling type was not set");
                 break;
+            } else if (filling_type == filling_type_t::fixed_setting) {
+                const std::vector<float> user_set {user_set_value};
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {aop, offset, *causal_mask_lt,
+                                filling_type, {user_set, cfg_name}});
+            } else if (filling_type == filling_type_t::causal_mask) {
+                // Casual mask filling
+                displace_args_.emplace(causal_mask_lt->id_,
+                        displace_args_t {
+                                aop, offset, *causal_mask_lt, filling_type});
             }
+            break;
+        }
+
+        // Fill proper data for bottom-right implicit casual mask
+        while (aop.kind_ == "Add") {
+            auto *aop_out_lt = &aop.out_lts_[0];
+            auto *child_op = &dg_->get_op_by_in_lt(aop_out_lt->id_);
+            if (child_op->kind_ != "Subtract") break;
+
+            auto *child_op_out_lt = &child_op->out_lts_[0];
+            auto *next_child_op = &dg_->get_op_by_in_lt(child_op_out_lt->id_);
+            if (next_child_op->kind_ != "GreaterEqual") break;
+
+            const std::string cfg_name = "Bottom_right_implicit_padding_mask";
+            static constexpr int seq_len_q = 0;
+            static constexpr int seq_len_kv = 1;
 
-            quantize_displace_.emplace(causal_mask_lt->id_,
-                    std::make_tuple(
-                            aop, offset, *causal_mask_lt, filling_type));
+            // The following subtract and greaterEqual must also be a part of
+            // the partition.
+            if (op_ids_set_.find(child_op->id_) == op_ids_set_.end()
+                    || op_ids_set_.find(next_child_op->id_)
+                            == op_ids_set_.end())
+                break;
+
+            const auto set_seq_len_displace_args =
+                    [&, this](const deserialized_op_t *op, int which_seq_len) {
path: tests/benchdnn/graph/input_displacer.cpp
user: login: dzarukin
body: ```suggestion
                    [&](const deserialized_op_t *op, int which_seq_len) {
```
created_at: 2025-04-10T16:34:13Z
author_association: CONTRIBUTOR

PR comments:
user: login: wzt1997
created_at: 2025-04-02T07:56:01Z
body: make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph
author_association: NONE
user: login: wzt1997
created_at: 2025-04-11T02:09:38Z
body: make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph
author_association: NONE