Skip to content

Commit 4db142c

Browse files
committed
Use IEEE precision for conv_transpose2d fp32 dot
1 parent 52aa24e commit 4db142c

1 file changed

Lines changed: 5 additions & 3 deletions

File tree

src/flag_gems/ops/conv_transpose2d.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -429,7 +429,7 @@ def conv_transpose2d_forward_kernel(
429429

430430
input_block = tl.load(curr_input_pointer, mask=input_mask, other=0.0)
431431
weight_block = tl.load(curr_weight_pointer, mask=weight_mask, other=0.0)
432-
accum += tl.dot(input_block, weight_block, allow_tf32=False)
432+
accum += tl.dot(input_block, weight_block, input_precision="ieee")
433433

434434
bias_pointer += (pid_group[None] * out_per_group_c)[None, :] + output_c_offset[
435435
None, :
@@ -571,7 +571,9 @@ def _conv_transpose2d_forward_residue_body(
571571
weight_block = tl.load(
572572
curr_weight_pointer, mask=weight_mask, other=0.0
573573
)
574-
accum += tl.dot(input_block, weight_block, allow_tf32=False)
574+
accum += tl.dot(
575+
input_block, weight_block, input_precision="ieee"
576+
)
575577

576578
bias_pointer += (pid_group[None] * out_per_group_c)[None, :] + output_c_offset[
577579
None, :
@@ -1088,7 +1090,7 @@ def conv_transpose2d_forward_scatter_kernel(
10881090

10891091
input_block = tl.load(curr_input_pointer, mask=input_mask, other=0.0)
10901092
weight_block = tl.load(curr_weight_pointer, mask=weight_mask, other=0.0)
1091-
accum += tl.dot(input_block, weight_block, allow_tf32=False)
1093+
accum += tl.dot(input_block, weight_block, input_precision="ieee")
10921094

10931095
bias_pointer += (pid_group[None] * out_per_group_c)[None, :] + output_c_offset[
10941096
None, :

0 commit comments

Comments
 (0)