tests: benchdnn: graph: add cases for sdpa training

Author: ElaineBao

tests/benchdnn/inputs/graph/complex_fusion/harness_mha_all

@@ -11,6 +11,8 @@
 --reset --dt=f32,bf16,f16 --case=complex_fusion/mha/GQA-fp16-v2.json
 --reset --dt=f32,bf16,f16 --case=complex_fusion/mha/sdpa-plain-wo-mask-f16.json
 --reset --dt=f32,bf16,f16 --case=complex_fusion/mha/sdpa-plain-implicit-causal-mask-fp32-bs1.json
+--reset --dt=0:f32+1:f32+10:f32+13:f32+14:f32 --case=complex_fusion/mha/sdpa-plain-training-forward-bf16-f32.json
+--reset --case=complex_fusion/mha/sdpa-plain-training-backward-f32.json
 
 # f16 inputs + f32 intermediates + f16 outputs
 --reset --op-kind=1:Multiply,1:Divide --case=complex_fusion/mha/sdpa-plain-simplified-f16-f32.json
@@ -38,6 +40,8 @@
 --reset --dt=2:f32+5:f32 --case=complex_fusion/mha/sdpa-plain-wo-mask-f16.json
 --reset --dt=2:f32+6:f32 --case=complex_fusion/mha/sdpa-plain-wo-scale-f16-bs1.json
 --reset --case=complex_fusion/mha/sdpa-plain-bottom-right-implicit-causal-mask-f16-f32.json
+--reset --dt=0:f16+1:f16+10:f16+13:f16+14:f16 --case=complex_fusion/mha/sdpa-plain-training-forward-bf16-f32.json
+--reset --dt=16:f16+17:f16+32:f16+33:f16+34:f16+36:f16+44:f16+45:f16+47:f16 --case=complex_fusion/mha/sdpa-plain-training-backward-bf16-f32.json
 # q_seq_len != kv_seq_len
 --reset --in-shapes=1:1x16x128x64+24:1x16x128x64 --case=complex_fusion/mha/sdpa-plain-bottom-right-implicit-causal-mask-f16-f32.json
 
@@ -63,6 +67,8 @@
 --reset --dt=2:f32+5:f32+0:bf16+1:bf16+4:bf16+7:bf16+9:bf16+10:bf16 --case=complex_fusion/mha/sdpa-plain-wo-mask-f16.json
 --reset --dt=2:f32+6:f32+0:bf16+1:bf16+5:bf16+7:bf16+8:bf16+9:bf16  --case=complex_fusion/mha/sdpa-plain-wo-scale-f16-bs1.json
 --reset --dt=0:bf16+1:bf16+4:bf16+22:bf16+24:bf16+25:bf16 --case=complex_fusion/mha/sdpa-plain-bottom-right-implicit-causal-mask-f16-f32.json
+--reset --case=complex_fusion/mha/sdpa-plain-training-forward-bf16-f32.json
+--reset --case=complex_fusion/mha/sdpa-plain-training-backward-bf16-f32.json
 
 # int8 graphs
 --reset --case=complex_fusion/mha/MHA-GPT-inf-int8-bs1.json

tests/benchdnn/inputs/graph/complex_fusion/harness_mha_ci

@@ -12,17 +12,24 @@
 --reset --dt=f32,bf16,f16 --case=complex_fusion/mha/GQA-fp16-v2.json
 --reset --dt=f32,bf16,f16 --case=complex_fusion/mha/sdpa-plain-wo-mask-f16.json
 --reset --dt=f32,bf16,f16 --case=complex_fusion/mha/sdpa-plain-implicit-causal-mask-fp32-bs1.json
+--reset --dt=0:f32+1:f32+10:f32+13:f32+14:f32 --case=complex_fusion/mha/sdpa-plain-training-forward-bf16-f32.json
+--reset --case=complex_fusion/mha/sdpa-plain-training-backward-f32.json
+
 # f16 inputs + f32 intermediates + f16 outputs
 --reset --case=complex_fusion/mha/sdpa-plain-simplified-f16-f32.json
 --reset --dt=4:f32+9:f32+14:f32 --case=complex_fusion/mha/GQA-fp16-v2.json
 --reset --case=complex_fusion/mha/sdpa-plain-bottom-right-implicit-causal-mask-f16-f32.json
 --reset --case=complex_fusion/mha/codegemma-bf16-f32.json
 --reset --case=complex_fusion/mha/gemma2-bf16-f32.json
+--reset --dt=0:f16+1:f16+10:f16+13:f16+14:f16 --case=complex_fusion/mha/sdpa-plain-training-forward-bf16-f32.json
+--reset --dt=16:f16+17:f16+32:f16+33:f16+34:f16+36:f16+44:f16+45:f16+47:f16 --case=complex_fusion/mha/sdpa-plain-training-backward-bf16-f32.json
 
 # bf16 inputs + f32 intermediates + bf16 outputs
 --reset --dt=1:bf16+2:bf16+3:bf16+4:bf16+5:bf16+6:bf16+104:bf16 --case=complex_fusion/mha/sdpa-plain-simplified-f16-f32.json
 --reset --dt=4:f32+9:f32+14:f32+1:bf16+3:bf16+8:bf16+11:bf16+16:bf16+20:bf16+19:bf16 --case=complex_fusion/mha/GQA-fp16-v2.json
 --reset --dt=0:bf16+1:bf16+4:bf16+22:bf16+24:bf16+25:bf16 --case=complex_fusion/mha/sdpa-plain-bottom-right-implicit-causal-mask-f16-f32.json
+--reset --case=complex_fusion/mha/sdpa-plain-training-forward-bf16-f32.json
+--reset --case=complex_fusion/mha/sdpa-plain-training-backward-bf16-f32.json
 
 # int8 graphs
 --reset --case=complex_fusion/mha/MHA-GPT-inf-int8-bs1.json