Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

graph: fix the intermediate data types in SDPA patterns #2894

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

graph: fix the intermediate data types in SDPA patterns #2894

Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
7a0285f
graph: interface: op: matmul supports mixed data types
TaoLv Mar 5, 2025
d32cdd2
graph: interface: op: softmax supports mixed data types
TaoLv Mar 12, 2025
79f3e18
graph: interface: op: binary ops support mixed data types
TaoLv Mar 12, 2025
d55039e
examples: graph: sdpa: define with f32 intermediate data type
TaoLv Mar 12, 2025
fdea8ad
graph: backend: dnnl: ukernel sdpa only supports f32 intermediates
TaoLv Mar 13, 2025
abe08ff
graph: backend: dnnl: pattern: sdpa: remove xf16 check from gpu pattern
TaoLv Mar 21, 2025
80f4f02
benchdnn: inputs: graph: add sdpa cases with f32 intermediate type
TaoLv Mar 13, 2025
df61ffb
benchdnn: inputs: graph: test f32 intermediates for implicit mask
TaoLv Mar 21, 2025
76f2cf9
doc: graph: op: update supported data types
TaoLv Mar 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 16 additions & 11 deletions src/graph/interface/op_def.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,17 @@ DNNL_GRAPH_OP_SCHEMA(Add, 1,
.set_num_inputs(2)
.set_num_outputs(1)
.set_commutative_inputs()
.set_input(0, "src_0", "T")
.set_input(1, "src_1", "T")
.set_output(0, "dst", "T")
.set_input(0, "src_0", "T1")
.set_input(1, "src_1", "T2")
.set_output(0, "dst", "T3")
.set_attr(op_attr::auto_broadcast, false, attribute_kind::s,
"numpy", {"none", "numpy"})
.set_type_constraints(
"T", {data_type::f32, data_type::bf16, data_type::f16})
"T1", {data_type::f32, data_type::bf16, data_type::f16})
.set_type_constraints(
"T2", {data_type::f32, data_type::bf16, data_type::f16})
.set_type_constraints(
"T3", {data_type::f32, data_type::bf16, data_type::f16})
.set_shape_inference_function(
infer_elemwise_arithmetic_output_shape))

Expand Down Expand Up @@ -791,9 +795,6 @@ DNNL_GRAPH_OP_SCHEMA(MishBackward, 1,
"T", {data_type::f32, data_type::bf16, data_type::f16})
.set_shape_inference_function(infer_identity_output_shape))

// TODO(Yixin): for Multiply. input and output needs to have the same dtypes
// But in current pytorch bridge's type promotion system, there's no
// such constraints. So this feature is postponed.
DNNL_GRAPH_OP_SCHEMA(Multiply, 1,
op_schema_t()
.set_num_inputs(2)
Expand Down Expand Up @@ -1127,13 +1128,17 @@ DNNL_GRAPH_OP_SCHEMA(Subtract, 1,
op_schema_t()
.set_num_inputs(2)
.set_num_outputs(1)
.set_input(0, "src_0", "T")
.set_input(1, "src_1", "T")
.set_output(0, "dst", "T")
.set_input(0, "src_0", "T1")
.set_input(1, "src_1", "T2")
.set_output(0, "dst", "T3")
.set_attr(op_attr::auto_broadcast, false, attribute_kind::s,
"numpy", {"none", "numpy"})
.set_type_constraints(
"T", {data_type::f32, data_type::bf16, data_type::f16})
"T1", {data_type::f32, data_type::bf16, data_type::f16})
.set_type_constraints(
"T2", {data_type::f32, data_type::bf16, data_type::f16})
.set_type_constraints(
"T3", {data_type::f32, data_type::bf16, data_type::f16})
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This requires some documentation about type promotion as users might wonder what happens for example with f16 <- f16 + bf16.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In fact, we don't allow f16 + bf16. It's mentioned in the "supported data types" section in the op document. When src0 and src1 have different data types, one of them should be f32 and the other one (f16 or bf16) will be promoted to f32 for calculation.

.set_shape_inference_function(
infer_elemwise_arithmetic_output_shape))

Expand Down