Skip to content

Commit 627a571

Browse files
AmdSampsanaromero77amd
authored andcommitted
[NO CP] triton sanity check for 2D POI (#2798)
Added a check that includes autotune configs for 2D POI only if their size is big enough. (cherry picked from commit a2b0fd7)
1 parent d235a15 commit 627a571

File tree

1 file changed

+23
-8
lines changed

1 file changed

+23
-8
lines changed

torch/_inductor/runtime/triton_heuristics.py

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1995,15 +1995,30 @@ def pointwise(
19951995
*hinted_configs,
19961996
]
19971997
if torch.version.hip:
1998-
configs += [ # add here
1999-
]
2000-
# bypass triton_config_with_settings -> triton_config logic
20011998
if "x" in size_hints and "y" in size_hints:
2002-
configs += [
2003-
Config({"XBLOCK": 512, "YBLOCK": 8}, num_warps=8), # wrt1/t21 # triton_poi_fused__unsafe_view_add_addmm_cat_clone_permute_split_with_sizes_view_19
2004-
Config({"XBLOCK": 32, "YBLOCK": 128}, num_warps=4), # wrt2: 570us : triton_poi_fused_add_transpose_view_52
2005-
Config({"XBLOCK":64, "YBLOCK": 32}, num_warps=8), # wrt3: 150us: triton_poi_fused__to_copy_add_native_layer_norm_native_layer_norm_backward_permute_view_103
2006-
]
1999+
"""add 2D tiling configs, but don't use triton_config_with_settings function
2000+
as it is buggy and might change the tiling randomly
2001+
"""
2002+
def addConfig__(xblock:int, yblock:int, num_warps:int, num_stages:int):
2003+
# only add a tiling config if size is bigger than the tile
2004+
# check also for grid overflow
2005+
xgrid = (size_hints["x"] + xblock - 1) // xblock
2006+
ygrid = (size_hints["y"] + yblock - 1) // yblock
2007+
if xgrid > 2147483647:
2008+
return
2009+
if ygrid > 65535:
2010+
return
2011+
if size_hints["x"] < xblock:
2012+
return
2013+
if size_hints["y"] < yblock:
2014+
return
2015+
# all good, add the config
2016+
configs.append(Config({"XBLOCK": xblock, "YBLOCK": yblock}, num_warps=num_warps, num_stages=num_stages))
2017+
addConfig__(512, 8, 8,1 ) # wrt1/t21 # triton_poi_fused__unsafe_view_add_addmm_cat_clone_permute_split_with_sizes_view_19
2018+
addConfig__(32, 128, 4, 1) # wrt2: 570us : triton_poi_fused_add_transpose_view_52
2019+
addConfig__(64, 32, 8, 1) # wrt3: 150us: triton_poi_fused__to_copy_add_native_layer_norm_native_layer_norm_backward_permute_view_103
2020+
addConfig__(64, 256, 4, 1) # wri0: 70us: triton_poi_fused_clone_tanh_transpose_19
2021+
addConfig__(512, 64, 8, 1) # wri0: 58us: triton_poi_fused_clone_53
20072022

20082023
if len(size_hints) == 3:
20092024
if disable_pointwise_autotuning(inductor_meta):

0 commit comments

Comments
 (0)