File tree Expand file tree Collapse file tree 1 file changed +17
-17
lines changed Expand file tree Collapse file tree 1 file changed +17
-17
lines changed Original file line number Diff line number Diff line change @@ -2955,6 +2955,14 @@ def _persistent_reduction_configs(
29552955 )
29562956 )
29572957
2958+ tiny_configs = [
2959+ triton_config_reduction (
2960+ size_hints ,
2961+ 2 * (256 // rnumel ) if rnumel <= 256 else 1 ,
2962+ rnumel ,
2963+ reduction_hint = reduction_hint ,
2964+ )
2965+ ]
29582966 # defer to more autotuning, initially
29592967 if "y" in size_hints :
29602968 pass
@@ -2981,23 +2989,15 @@ def _persistent_reduction_configs(
29812989 )
29822990 ]
29832991
2984- elif reduction_hint == ReductionHint .OUTER :
2985- configs = configs [- 1 :]
2986- elif reduction_hint == ReductionHint .OUTER_TINY :
2987- tiny_configs = [
2988- triton_config_reduction (
2989- size_hints ,
2990- 2 * (256 // rnumel ) if rnumel <= 256 else 1 ,
2991- rnumel ,
2992- reduction_hint = reduction_hint ,
2993- )
2994- ]
2995- if max_autotune_enabled :
2996- for tconfig in tiny_configs :
2997- if tconfig not in configs :
2998- configs .append (tconfig )
2999- else :
3000- configs = tiny_configs
2992+ elif reduction_hint == ReductionHint .OUTER :
2993+ configs = configs [- 1 :]
2994+ elif reduction_hint == ReductionHint .OUTER_TINY :
2995+ configs = tiny_configs
2996+ else :
2997+ # If autotune is enabled append tiny configs
2998+ for conf in tiny_configs :
2999+ if conf not in configs :
3000+ configs .append (conf )
30013001
30023002 for c in configs :
30033003 # we don't need Rn_BLOCK for persistent reduction
You can’t perform that action at this time.
0 commit comments