graph: backend: dnnl: support sdpa training forward pattern with larger_partition_kernel

Author: ElaineBao

src/graph/backend/dnnl/kernels/sdp_decomp.cpp

@@ -51,7 +51,7 @@ status_t sdp_decomp_kernel_t<quantized, dt>::compile_impl(
     BACKEND_DNNL_CHECK(set_given_inputs_outputs(subgraph_, inputs, outputs));
 
     // Check if it's supported by decomposition kernel
-    if (!sdp_cfg_.initial_check(subgraph_, inputs))
+    if (!sdp_cfg_.initial_check(subgraph_, inputs, outputs))
         return status::unimplemented;
 
     subgraph_visualizer_t vis(part->id(), [this](const value_t *val) {

src/graph/backend/dnnl/kernels/sdp_decomp_config.cpp

@@ -27,14 +27,17 @@ namespace graph {
 namespace dnnl_impl {
 
 bool sdp_decomp_config_t::initial_check(const std::shared_ptr<subgraph_t> &sg,
-        const std::vector<logical_tensor_t> &inputs) {
+        const std::vector<logical_tensor_t> &inputs,
+        const std::vector<logical_tensor_t> &outputs) {
     // The order of input logical tensors in inputs is not certain, we need
     // to record the input offset in a certain order of ops.
     CHECK_BOOL(record_input_offset(sg, inputs));
     dims src1_user_dims = ltw(inputs[graph_inport[mm1_src]]).vdims();
     ndims = src1_user_dims.size();
     VCHECK_SDP_DECOMP(ndims == 4 || ndims == 5, false,
             "Input dims should be 4 or 5, but got %zu", src1_user_dims.size());
+    VCHECK_SDP_DECOMP(
+            outputs.size() == 1, false, "Doesn't support SDPA training yet");
 
     // Initialize SDP input dimension according to the src of mm1
     int index = 0;

src/graph/backend/dnnl/kernels/sdp_decomp_config.hpp

@@ -156,7 +156,8 @@ struct sdp_decomp_config_t {
     // If the check passes, initialize few members according to inputs
     // If no, return unimplemented status directly and fallback to large kernel
     bool initial_check(const std::shared_ptr<subgraph_t> &sg,
-            const std::vector<logical_tensor_t> &inputs);
+            const std::vector<logical_tensor_t> &inputs,
+            const std::vector<logical_tensor_t> &outputs);
 
     // Used to construct all params that SDP need
     template <bool quantized = false,

src/graph/backend/dnnl/kernels/sdp_primitive.cpp

@@ -59,7 +59,7 @@ status_t sdp_primitive_kernel_t<quantized>::compile_impl(
                     p_engine_, part->get_fpmath_mode(), false, true);
     CHECK(set_given_inputs_outputs(subgraph_, inputs, outputs));
 
-    CHECK(cfg_.initial_check(subgraph_, inputs));
+    CHECK(cfg_.initial_check(subgraph_, inputs, outputs));
 
     subgraph_visualizer_t vis(part->id(), [this](const value_t *val) {
         return this->memory_planner_.get_memory_info(val);

src/graph/backend/dnnl/kernels/sdp_primitive_config.cpp

@@ -169,10 +169,13 @@ 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,
+        const std::vector<logical_tensor_t> &outputs, bool v1_kernel) {
     // At least 3 inputs: Q, K, V
     VCHECK_SDP_PRIMITIVE(inputs.size() >= 3, status::invalid_arguments,
             "At least 3 inputs are required");
+    VCHECK_SDP_PRIMITIVE(outputs.size() == 1, status::unimplemented,
+            "Doesn't support SDPA training yet");
 
     // Ukernel doesn't support f32 datatype now
     VCHECK_SDP_PRIMITIVE(inputs[0].data_type != dnnl_data_type_t::dnnl_f32,

src/graph/backend/dnnl/kernels/sdp_primitive_config.hpp

@@ -84,6 +84,7 @@ struct sdp_primitive_config_t {
     // 3. only support 4-dims tensor
     status_t initial_check(const std::shared_ptr<subgraph_t> &sg,
             const std::vector<logical_tensor_t> &inputs,
+            const std::vector<logical_tensor_t> &outputs,
             bool v1_kernel = false);
 
     // Initialize parameters and primitive.

src/graph/backend/dnnl/kernels/sdp_primitive_v1.cpp

@@ -59,7 +59,7 @@ status_t sdp_primitive_v1_kernel_t::compile_impl(
                     p_engine_, part->get_fpmath_mode(), false, true);
     CHECK(set_given_inputs_outputs(subgraph_, inputs, outputs));
 
-    CHECK(cfg_.initial_check(subgraph_, inputs, true));
+    CHECK(cfg_.initial_check(subgraph_, inputs, outputs, true));
 
     subgraph_visualizer_t vis(part->id(), [this](const value_t *val) {
         return this->memory_planner_.get_memory_info(val);