graph: backend: dnnl: merge 2 sdpa primitive kernel

xiang1guo · xiang1guo · commit 5b6dac7e2362 · 2025-06-15T20:18:02.000-07:00
diff --git a/src/graph/backend/dnnl/kernels/sdp.hpp b/src/graph/backend/dnnl/kernels/sdp.hpp
@@ -27,7 +27,6 @@
 #include "graph/backend/dnnl/kernels/large_partition.hpp"
 #include "graph/backend/dnnl/kernels/sdp_decomp.hpp"
 #include "graph/backend/dnnl/kernels/sdp_primitive.hpp"
-#include "graph/backend/dnnl/kernels/sdp_primitive_v1.hpp"
 
 #include "graph/backend/dnnl/dnnl_partition_impl.hpp"
 
@@ -66,15 +65,7 @@ struct sdp_base_t : public kernel_base_t {
 
         status_t ret = status::unimplemented;
 
-        // SDPA Ukernel v1 with fused internal sdpa solution. Support fload sdpa
-        // only.
-        // TODO(GX): Support quantized sdpa and merge with sdp_primitive_kernel_t.
         if (enable_ukernel) {
-            kernel = std::make_shared<sdp_primitive_v1_kernel_t<quantized>>();
-            ret = kernel->compile_impl(part, g_engine, inputs, outputs);
-        }
-
-        if (ret != status::success && enable_ukernel) {
             kernel = std::make_shared<sdp_primitive_kernel_t<quantized>>();
             ret = kernel->compile_impl(part, g_engine, inputs, outputs);
         }
diff --git a/src/graph/backend/dnnl/kernels/sdp_primitive.cpp b/src/graph/backend/dnnl/kernels/sdp_primitive.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2024-2025 Intel Corporation
+* 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.
@@ -49,6 +49,7 @@ status_t sdp_primitive_kernel_t<quantized>::compile_impl(
 #if defined(DNNL_WITH_SYCL) && DNNL_GPU_VENDOR != DNNL_VENDOR_INTEL
     return status::unimplemented;
 #endif
+
     p_engine_ = make_dnnl_engine(*g_engine);
     g_alloc_
             = reinterpret_cast<graph::allocator_t *>(g_engine->get_allocator());
@@ -68,7 +69,6 @@ status_t sdp_primitive_kernel_t<quantized>::compile_impl(
 
     BACKEND_DNNL_ADD_PASS(pipeline, lower_down);
     BACKEND_DNNL_ADD_PASS(pipeline, fuse_implicit_causal_mask);
-    BACKEND_DNNL_ADD_PASS(pipeline, fuse_reshape_for_gqa);
     if (quantized) {
         BACKEND_DNNL_ADD_PASS(pipeline, lift_up_typecast);
         BACKEND_DNNL_ADD_PASS(pipeline, lift_up_quantize);
@@ -92,44 +92,39 @@ status_t sdp_primitive_kernel_t<quantized>::compile_impl(
 
     pipeline.reset_visualize_arg(true, false);
     BACKEND_DNNL_ADD_PASS(pipeline, infer_shape);
+    BACKEND_DNNL_ADD_PASS(pipeline, fuse_sdpa);
     BACKEND_DNNL_ADD_PASS(pipeline, fuse_dst_transpose_to_predecessor);
+    BACKEND_DNNL_ADD_PASS(pipeline, fuse_reshape_for_gqa_gpu);
+    BACKEND_DNNL_ADD_PASS(pipeline, insert_reshape_for_sdpa);
     BACKEND_DNNL_ADD_PASS(pipeline, layout_propagation);
 
-    // bind the memory for each op
+    // bind the memory for each op`
     auto memory_plan = [&](std::shared_ptr<subgraph_t> &sg) {
         return memory_planner_.run(sg);
     };
     pipeline.reset_visualize_arg(true, true);
     BACKEND_DNNL_ADD_PASS(pipeline, memory_plan);
+    BACKEND_DNNL_ADD_PASS(pipeline, compile_ops);
 
-    auto modify_subgraph = [&] {
-        // Run the added passes
-        CHECK(pipeline.run(subgraph_));
-
-        // fill information for inputs logical tensors
-        for (size_t i = 0; i < inputs.size(); i++) {
-            auto &in = const_cast<logical_tensor_t &>(inputs[i]);
-            in = subgraph_->ins_[i];
-        }
+    // Run the added passes
+    BACKEND_DNNL_CHECK(pipeline.run(subgraph_));
 
-        // fill information for outputs logical tensors
-        for (size_t i = 0; i < outputs.size(); i++) {
-            auto &out = const_cast<logical_tensor_t &>(outputs[i]);
-            out = subgraph_->outs_[i];
-        }
+    // fill information for inputs logical tensors
+    for (size_t i = 0; i < inputs.size(); i++) {
+        auto &in = const_cast<logical_tensor_t &>(inputs[i]);
+        in = subgraph_->ins_[i];
+    }
 
-        return status::success;
-    };
+    // fill information for outputs logical tensors
+    for (size_t i = 0; i < outputs.size(); i++) {
+        auto &out = const_cast<logical_tensor_t &>(outputs[i]);
+        out = subgraph_->outs_[i];
+    }
 
     resource_ctor_ = [this]() {
         return this->memory_planner_.get_exec_args_set().clone();
     };
 
-    CHECK(modify_subgraph());
-
-    cfg_.quantized_ = quantized;
-    CHECK(cfg_.init(subgraph_, p_engine_, inputs, outputs));
-
     return status::success;
 }
 
@@ -145,67 +140,13 @@ void sdp_primitive_kernel_t<quantized>::prepare_args_set(
         mem_idx.first.set_data_handle(
                 outputs[mem_idx.second].get_data_handle());
     }
-}
 
-template <bool quantized>
-status_t sdp_primitive_kernel_t<quantized>::get_prim_exec_args(
-        exec_args_t &args, memory (&mem_storage)[10],
-        const execution_args_set_t *res) const {
-    bool ok = res->find_value_mem_map(cfg_.q_.get(), mem_storage[0])
-            && res->find_value_mem_map(cfg_.k_.get(), mem_storage[1])
-            && res->find_value_mem_map(cfg_.v_.get(), mem_storage[2])
-            && res->find_value_mem_map(cfg_.dst_.get(), mem_storage[3]);
-
-    if (cfg_.scale_)
-        ok = ok && res->find_value_mem_map(cfg_.scale_.get(), mem_storage[4]);
-    if (cfg_.attn_mask_)
-        ok = ok
-                && res->find_value_mem_map(
-                        cfg_.attn_mask_.get(), mem_storage[5]);
-    if (quantized && !(cfg_.k_scale_ || cfg_.v_scale_))
-        return status::invalid_arguments;
-    if (cfg_.k_scale_)
-        ok = ok && res->find_value_mem_map(cfg_.k_scale_.get(), mem_storage[6]);
-    if (cfg_.v_scale_)
-        ok = ok && res->find_value_mem_map(cfg_.v_scale_.get(), mem_storage[7]);
-
-    if (cfg_.k_zero_points_)
-        ok = ok
-                && res->find_value_mem_map(
-                        cfg_.k_zero_points_.get(), mem_storage[8]);
-    if (cfg_.v_zero_points_)
-        ok = ok
-                && res->find_value_mem_map(
-                        cfg_.v_zero_points_.get(), mem_storage[9]);
-
-    VCONDCHECK(graph, exec, check, sdp_primitive_kernel, ok,
-            status::runtime_error,
-            "sdp_primitive_kernel get_prim_exec_args failed");
-
-    memory_arg_t mem_arg_q = {mem_storage[0].get(), true};
-    memory_arg_t mem_arg_k = {mem_storage[1].get(), true};
-    memory_arg_t mem_arg_v = {mem_storage[2].get(), true};
-    memory_arg_t mem_arg_dst = {mem_storage[3].get(), false};
-    memory_arg_t mem_arg_scale = {mem_storage[4].get(true), true};
-    memory_arg_t mem_arg_mask = {mem_storage[5].get(true), true};
-    memory_arg_t mem_arg_k_scale = {mem_storage[6].get(true), true};
-    memory_arg_t mem_arg_v_scale = {mem_storage[7].get(true), true};
-    memory_arg_t mem_arg_k_zero_points = {mem_storage[8].get(true), true};
-    memory_arg_t mem_arg_v_zero_points = {mem_storage[9].get(true), true};
-
-    args.clear();
-    args[DNNL_ARG_QUERIES] = mem_arg_q;
-    args[DNNL_ARG_KEYS] = mem_arg_k;
-    args[DNNL_ARG_VALUES] = mem_arg_v;
-    args[DNNL_ARG_DST] = mem_arg_dst;
-    args[DNNL_ARG_SCALE] = mem_arg_scale;
-    args[DNNL_ARG_ATTN_MASK] = mem_arg_mask;
-    args[DNNL_ARG_ATTR_SCALES | DNNL_ARG_KEYS] = mem_arg_k_scale;
-    args[DNNL_ARG_ATTR_SCALES | DNNL_ARG_VALUES] = mem_arg_v_scale;
-    args[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_KEYS] = mem_arg_k_zero_points;
-    args[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_VALUES] = mem_arg_v_zero_points;
+    grantor_t var_grantor = memory_planner_.internal_temporary_grantor(
+            scratchpad.get_buffer());
 
-    return status::success;
+    for (auto &mem_offkey : res->get_mems_use_internal_temporary()) {
+        mem_offkey.first.set_data_handle(var_grantor.get(mem_offkey.second));
+    }
 }
 
 template <bool quantized>
@@ -218,17 +159,16 @@ status_t sdp_primitive_kernel_t<quantized>::execute_impl(
     execution_args_set_t *res = res_cache.get_or_add(
             reinterpret_cast<size_t>(this), resource_ctor_);
 
-    // Micro kernel doesn't use scratchpad memory, here we force-set size as
-    // zero to avoid redundant memory allocation and deallocation.
-    temporary_scratchpad_t scratchpad(0, p_engine_, *g_alloc_);
+    temporary_scratchpad_t scratchpad(
+            memory_planner_.total_internal_temporary_size(), p_engine_,
+            *g_alloc_);
     prepare_args_set(res, inputs, outputs, scratchpad);
 
-    memory mem_storage[10];
-    exec_args_t args;
-    CHECK(get_prim_exec_args(args, mem_storage, res));
-    exec_ctx_t ctx(p_stream.get(), std::move(args));
+    for (size_t i = 0; i < subgraph_->execs_.size(); i++) {
+        subgraph_->execs_[i]->execute(p_stream, res->get_exec_args()[i]);
+    }
 
-    return cfg_.sdpa_prim_->execute(ctx);
+    return status::success;
 }
 
 #ifdef DNNL_WITH_SYCL
@@ -242,42 +182,31 @@ status_t sdp_primitive_kernel_t<quantized>::sycl_execute_impl(
 #if DNNL_GPU_VENDOR != DNNL_VENDOR_INTEL
     return status::unimplemented;
 #endif
+    auto deps = sycl_deps;
+    ::sycl::event returned_event;
+
     dnnl::stream p_stream = make_dnnl_stream(p_engine_, *g_stream);
 
     thread_local_cache_t<execution_args_set_t> res_cache;
     execution_args_set_t *res = res_cache.get_or_add(
             reinterpret_cast<size_t>(this), resource_ctor_);
 
-    // Micro kernel doesn't use scratchpad memory, here we force-set size as
-    // zero to avoid redundant memory allocation and deallocation.
-    temporary_scratchpad_t scratchpad(0, p_engine_, *g_alloc_);
+    temporary_scratchpad_t scratchpad(
+            memory_planner_.total_internal_temporary_size(), p_engine_,
+            *g_alloc_);
     prepare_args_set(res, inputs, outputs, scratchpad);
 
-    memory mem_storage[10];
-    exec_args_t args;
-    CHECK(get_prim_exec_args(args, mem_storage, res));
-    exec_ctx_t ctx(p_stream.get(), std::move(args));
-
-    // Relying on the library's internals here. Since graph API is currently
-    // enabled only for the Intel vendor it is fine to cast stream to
-    // gpu::intel::sycl::stream_t unconditionally.
-    auto *sycl_stream = dnnl::impl::utils::downcast<
-            dnnl::impl::gpu::intel::sycl::stream_t *>(p_stream.get());
-
-    sycl_stream->before_exec_hook();
-
-    if (!sycl_deps.empty()) sycl_stream->sycl_ctx().set_deps(sycl_deps);
-
-    auto status = cfg_.sdpa_prim_->execute(ctx);
-
-    auto return_event = sycl_stream->get_output_event();
-
-    scratchpad.set_deps(return_event);
-    if (sycl_event) *sycl_event = return_event;
+    for (size_t i = 0; i < subgraph_->execs_.size(); i++) {
+        if (subgraph_->is_constant_[i]) continue;
+        returned_event = subgraph_->execs_[i]->execute_sycl(
+                p_stream, res->get_exec_args()[i], deps);
+        deps = {returned_event};
+    }
 
-    sycl_stream->after_exec_hook();
+    scratchpad.set_deps(returned_event);
+    if (sycl_event) *sycl_event = returned_event;
 
-    return status;
+    return status::success;
 }
 #endif
 
@@ -287,50 +216,31 @@ status_t sdp_primitive_kernel_t<quantized>::ocl_execute_impl(
         const stream_t *g_stream, const std::vector<tensor_t> &inputs,
         const std::vector<tensor_t> &outputs,
         const std::vector<cl_event> &cl_deps, cl_event *ret_event) {
+    auto deps = cl_deps;
+    cl_event returned_event {};
 
     dnnl::stream p_stream = make_dnnl_stream(p_engine_, *g_stream);
 
     thread_local_cache_t<execution_args_set_t> res_cache;
     execution_args_set_t *res = res_cache.get_or_add(
             reinterpret_cast<size_t>(this), resource_ctor_);
 
-    // Micro kernel doesn't use scratchpad memory, here we force-set size as
-    // zero to avoid redundant memory allocation and deallocation.
-    temporary_scratchpad_t scratchpad(0, p_engine_, *g_alloc_);
+    temporary_scratchpad_t scratchpad(
+            memory_planner_.total_internal_temporary_size(), p_engine_,
+            *g_alloc_);
     prepare_args_set(res, inputs, outputs, scratchpad);
 
-    memory mem_storage[10];
-    exec_args_t args;
-    CHECK(get_prim_exec_args(args, mem_storage, res));
-    exec_ctx_t ctx(p_stream.get(), std::move(args));
-
-    // TODO (pc): refactor
-    auto *ocl_stream = dnnl::impl::utils::downcast<gpu::intel::ocl::stream_t *>(
-            p_stream.get());
-
-    ocl_stream->before_exec_hook();
-
-    if (!cl_deps.empty()) {
-        std::vector<xpu::ocl::wrapper_t<cl_event>> events(cl_deps.size());
-        for (size_t i = 0; i < cl_deps.size(); i++)
-            events[i] = xpu::ocl::wrapper_t<cl_event>(cl_deps[i], true);
-        ocl_stream->ocl_ctx().set_deps(events);
+    for (size_t i = 0; i < subgraph_->execs_.size(); i++) {
+        if (subgraph_->is_constant_[i]) continue;
+        returned_event = subgraph_->execs_[i]->execute_ocl(
+                p_stream, res->get_exec_args()[i], deps);
+        deps = {returned_event};
     }
 
-    auto status = cfg_.sdpa_prim_->execute(ctx);
-
-    cl_event return_event = nullptr;
-    if ((ocl_stream->flags() & stream_flags::in_order) == 0) {
-        auto last = ocl_stream->get_output_event();
-        return_event = last.release();
-    }
+    scratchpad.set_deps(returned_event);
+    if (ret_event) *ret_event = returned_event;
 
-    scratchpad.set_deps(return_event);
-    if (ret_event) *ret_event = return_event;
-
-    ocl_stream->after_exec_hook();
-
-    return status;
+    return status::success;
 }
 #endif
 
diff --git a/src/graph/backend/dnnl/kernels/sdp_primitive.hpp b/src/graph/backend/dnnl/kernels/sdp_primitive.hpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2024 Intel Corporation
+* 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.
@@ -73,9 +73,6 @@ struct sdp_primitive_kernel_t : public kernel_base_t {
             const std::vector<tensor_t> &outputs,
             const scratchpad_t &scratchpad);
 
-    status_t get_prim_exec_args(exec_args_t &args, memory (&mem_storage)[10],
-            const execution_args_set_t *res) const;
-
     status_t execute_impl(const stream_t *g_stream,
             const std::vector<tensor_t> &inputs,
             const std::vector<tensor_t> &outputs) override;
diff --git a/src/graph/backend/dnnl/kernels/sdp_primitive_config.cpp b/src/graph/backend/dnnl/kernels/sdp_primitive_config.cpp
@@ -169,7 +169,7 @@ status_t sdp_primitive_config_t::locate_io(std::shared_ptr<subgraph_t> &sg,
 
 status_t sdp_primitive_config_t::initial_check(
         const std::shared_ptr<subgraph_t> &sg,
-        const std::vector<logical_tensor_t> &inputs, bool v1_kernel) {
+        const std::vector<logical_tensor_t> &inputs) {
     // At least 3 inputs: Q, K, V
     VCHECK_SDP_PRIMITIVE(inputs.size() >= 3, status::invalid_arguments,
             "At least 3 inputs are required");
@@ -302,15 +302,6 @@ status_t sdp_primitive_config_t::initial_check(
     VCHECK_SDP_PRIMITIVE(q_id != -1 && k_id != -1 && v_id != -1,
             status::unimplemented, "Q, K, V are not found");
 
-    // Note: sdpa_primitive_v1 kernel accept 5D GQA pattern, and will reshape to
-    // 4D in later compilation pass.
-    if (!v1_kernel) {
-        VCHECK_SDP_PRIMITIVE(ltw(inputs[q_id]).vdims().size() == 4
-                        && ltw(inputs[k_id]).vdims().size() == 4
-                        && ltw(inputs[v_id]).vdims().size() == 4,
-                status::unimplemented, "Q, K, V should be 4-dims");
-    }
-
     // sdp_primitive only supports single scale value.
     if (scale) {
         const auto &s = scale->get_input_value(1)->get_logical_tensor();
diff --git a/src/graph/backend/dnnl/kernels/sdp_primitive_config.hpp b/src/graph/backend/dnnl/kernels/sdp_primitive_config.hpp
@@ -83,8 +83,7 @@ struct sdp_primitive_config_t {
     // 2. only support fp16 data type
     // 3. only support 4-dims tensor
     status_t initial_check(const std::shared_ptr<subgraph_t> &sg,
-            const std::vector<logical_tensor_t> &inputs,
-            bool v1_kernel = false);
+            const std::vector<logical_tensor_t> &inputs);
 
     // Initialize parameters and primitive.
     status_t init(std::shared_ptr<subgraph_t> &sg, const dnnl::engine &p_engine,
diff --git a/src/graph/backend/dnnl/kernels/sdp_primitive_v1.cpp b/src/graph/backend/dnnl/kernels/sdp_primitive_v1.cpp
diff --git a/src/graph/backend/dnnl/kernels/sdp_primitive_v1.hpp b/src/graph/backend/dnnl/kernels/sdp_primitive_v1.hpp
