modify

src/flag_gems/ops/addmm.py

@@ -12,7 +12,7 @@
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("addmm"),
+    configs=runtime.get_tuned_config("addmm"),
     key=["M", "N", "K"],
 )
 @triton.jit(do_not_specialize=["alpha", "beta"])

src/flag_gems/ops/all.py

@@ -21,7 +21,7 @@ def reduce_all(a, b):
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("all"), key=["M", "N"])
+@triton.autotune(configs=runtime.get_tuned_config("all"), key=["M", "N"])
 @triton.jit
 def all_kernel_dim(
     inp,

src/flag_gems/ops/amax.py

@@ -42,7 +42,7 @@ def amax_kernel_2(mid, out, mid_size, BLOCK_MID: tl.constexpr):
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("amax"), key=["M", "N"])
+@triton.autotune(configs=runtime.get_tuned_config("amax"), key=["M", "N"])
 @triton.jit
 def amax_kernel(
     inp,

src/flag_gems/ops/any.py

@@ -21,7 +21,7 @@ def reduce_any(a, b):
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("any"), key=["M", "N"])
+@triton.autotune(configs=runtime.get_tuned_config("any"), key=["M", "N"])
 @triton.jit
 def any_kernel_dim(
     inp,

src/flag_gems/ops/argmax.py

@@ -47,7 +47,7 @@ def argmax_kernel_2(mid_value, mid_index, out, mid_size, BLOCK_MID: tl.constexpr
 
 
 @libentry()
-@triton.heuristics(runtime.get_heuristics_config("argmax"))
+@triton.heuristics(runtime.get_heuristic_config("argmax"))
 @triton.jit
 def argmax_kernel(
     inp,

src/flag_gems/ops/bmm.py

@@ -12,10 +12,10 @@
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("bmm"),
+    configs=runtime.get_tuned_config("bmm"),
     key=["M", "N", "K"],
 )
-@triton.heuristics(runtime.get_heuristics_config("bmm"))
+@triton.heuristics(runtime.get_heuristic_config("bmm"))
 @triton.jit
 def bmm_kernel(
     A,

src/flag_gems/ops/count_nonzero.py

@@ -23,7 +23,7 @@ def count_nonzero_kernel_1(x_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr):
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("count_nonzero"), key=["numel"])
+@triton.autotune(configs=runtime.get_tuned_config("count_nonzero"), key=["numel"])
 @triton.jit
 def count_nonzero_kernel(x_ptr, out_ptr, N, numel, BLOCK_SIZE: tl.constexpr):
     pid_x = tle.program_id(0)
@@ -41,7 +41,7 @@ def count_nonzero_kernel(x_ptr, out_ptr, N, numel, BLOCK_SIZE: tl.constexpr):
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("count_nonzero"), key=["numel"])
+@triton.autotune(configs=runtime.get_tuned_config("count_nonzero"), key=["numel"])
 @triton.jit
 def count_nonzero_combin_kernel_1(x_ptr, out_ptr, N, numel, BLOCK_SIZE: tl.constexpr):
     pid_x = tle.program_id(0)

src/flag_gems/ops/cross_entropy_loss.py

@@ -12,7 +12,7 @@
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss"),
+    configs=runtime.get_tuned_config("cross_entropy_loss"),
     key=["C", "D"],
 )
 @triton.jit(do_not_specialize=["ignore_index"])
@@ -73,7 +73,7 @@ def celoss_indices_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss"),
+    configs=runtime.get_tuned_config("cross_entropy_loss"),
     key=["C", "D"],
 )
 @triton.jit(do_not_specialize=["label_smoothing"])
@@ -132,7 +132,7 @@ def celoss_probability_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss"),
+    configs=runtime.get_tuned_config("cross_entropy_loss"),
     key=["C", "D"],
 )
 @triton.jit(do_not_specialize=["ignore_index", "label_smoothing"])
@@ -213,7 +213,7 @@ def celoss_indices_smooth_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss"),
+    configs=runtime.get_tuned_config("cross_entropy_loss"),
     key=["C", "D"],
 )
 @triton.jit(do_not_specialize=["ignore_index", "mean_num"])
@@ -284,7 +284,7 @@ def celoss_indices_bwd(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss"),
+    configs=runtime.get_tuned_config("cross_entropy_loss"),
     key=["C", "D"],
 )
 @triton.jit(do_not_specialize=["label_smoothing", "mean_num"])
@@ -369,7 +369,7 @@ def celoss_probability_bwd(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss"),
+    configs=runtime.get_tuned_config("cross_entropy_loss"),
     key=["C", "D"],
 )
 @triton.jit(do_not_specialize=["ignore_index", "label_smoothing", "mean_num"])
@@ -463,7 +463,7 @@ def celoss_indices_smooth_bwd(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("cross_entropy_loss_sum_and_scale"),
+    configs=runtime.get_tuned_config("cross_entropy_loss_sum_and_scale"),
     key=[
         "N",
     ],

src/flag_gems/ops/dropout.py

@@ -13,7 +13,7 @@
 from ..runtime import torch_device_fn
 
 
-@triton.heuristics(runtime.get_heuristics_config("dropout"))
+@triton.heuristics(runtime.get_heuristic_config("dropout"))
 @triton.jit(do_not_specialize=["p", "philox_seed", "philox_offset"])
 def dropout_forward_kernel(
     X,
@@ -65,7 +65,7 @@ def dropout_forward_kernel(
     tl.store(Y + off_3, y3, mask=off_3 < N, eviction_policy="evict_first")
 
 
-@triton.heuristics(runtime.get_heuristics_config("dropout"))
+@triton.heuristics(runtime.get_heuristic_config("dropout"))
 @triton.jit(do_not_specialize=["p", "philox_seed", "philox_offset"])
 def dropout_backward_kernel(
     DY,

src/flag_gems/ops/exponential_.py

@@ -13,7 +13,7 @@
 from ..runtime import torch_device_fn
 
 
-@triton.heuristics(runtime.get_heuristics_config("exponential_"))
+@triton.heuristics(runtime.get_heuristic_config("exponential_"))
 @triton.jit(do_not_specialize=["philox_seed", "philox_offset", "N"])
 def fused_exponential_kernel(
     out_ptr,

src/flag_gems/ops/gather.py

@@ -36,7 +36,7 @@ def generate_gather_kernel(
     code.writeline("@libentry()")
     code.writeline("@triton.heuristics(")
     with code.indent():
-        code.writeline("runtime.get_heuristics_config('gather')")
+        code.writeline("runtime.get_heuristic_config('gather')")
     code.writeline(")")
     code.writeline("@triton.jit")
 

src/flag_gems/ops/index_select.py

@@ -10,7 +10,7 @@
 
 
 @libentry()
-@triton.heuristics(runtime.get_heuristics_config("index_select"))
+@triton.heuristics(runtime.get_heuristic_config("index_select"))
 @triton.jit
 def index_select_kernel(
     inp, out, M, N, index, index_len, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr

src/flag_gems/ops/instancenorm.py

@@ -22,7 +22,7 @@ def prev_multiple_of(a, b):
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("instancenorm"),
+    configs=runtime.get_tuned_config("instancenorm"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -72,7 +72,7 @@ def instance_norm_persistent_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("instancenorm"),
+    configs=runtime.get_tuned_config("instancenorm"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -126,7 +126,7 @@ def instance_norm_persistent_kernel_multiline(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("instance_norm_loop"),
+    configs=runtime.get_tuned_config("instance_norm_loop"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -218,7 +218,7 @@ def instance_norm_loop_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("instancenorm"),
+    configs=runtime.get_tuned_config("instancenorm"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -312,7 +312,7 @@ def update_running_stats_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("instance_norm_backward"),
+    configs=runtime.get_tuned_config("instance_norm_backward"),
     key=["M", "N", "C"],
 )
 @triton.jit
@@ -379,7 +379,7 @@ def instance_norm_backward_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("instance_norm_weight_bias_backward"),
+    configs=runtime.get_tuned_config("instance_norm_weight_bias_backward"),
     key=["N", "B", "C"],
 )
 @triton.jit

src/flag_gems/ops/layernorm.py

@@ -20,7 +20,7 @@ def prev_multiple_of(a, b):
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("layer_norm_persistent"),
+    configs=runtime.get_tuned_config("layer_norm_persistent"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -69,7 +69,7 @@ def layer_norm_persistent_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("layer_norm_persistent"),
+    configs=runtime.get_tuned_config("layer_norm_persistent"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -123,7 +123,7 @@ def layer_norm_persistent_kernel_multiline(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("layer_norm_loop"),
+    configs=runtime.get_tuned_config("layer_norm_loop"),
     key=["M", "N"],
 )
 @triton.jit(do_not_specialize=["eps"])
@@ -220,7 +220,7 @@ def layer_norm_loop_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("layer_norm_backward"),
+    configs=runtime.get_tuned_config("layer_norm_backward"),
     key=["M", "N"],
 )
 @triton.jit
@@ -288,7 +288,7 @@ def layer_norm_backward_kernel(
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("weight_bias_backward"),
+    configs=runtime.get_tuned_config("weight_bias_backward"),
     key=["N"],
 )
 @triton.jit

src/flag_gems/ops/log_softmax.py

@@ -54,7 +54,7 @@ def log_softmax_kernel(
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("log_softmax"), key=["M", "N"])
+@triton.autotune(configs=runtime.get_tuned_config("log_softmax"), key=["M", "N"])
 @triton.jit
 def log_softmax_backward_kernel(
     out_ptr,

src/flag_gems/ops/masked_fill.py

@@ -10,7 +10,7 @@
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("masked_fill"), key=["N"])
+@triton.autotune(configs=runtime.get_tuned_config("masked_fill"), key=["N"])
 @triton.jit
 def masked_fill_kernel(inp, expand_mask, value, out, N, BLOCK_SIZE: tl.constexpr):
     pid = tle.program_id(axis=0)
@@ -24,7 +24,7 @@ def masked_fill_kernel(inp, expand_mask, value, out, N, BLOCK_SIZE: tl.constexpr
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("masked_fill"), key=["N"])
+@triton.autotune(configs=runtime.get_tuned_config("masked_fill"), key=["N"])
 @triton.jit
 def masked_fill_kernel_self(inp, expand_mask, value, N, BLOCK_SIZE: tl.constexpr):
     pid = tle.program_id(axis=0)

src/flag_gems/ops/masked_select.py

@@ -11,7 +11,7 @@
 
 
 @libentry()
-@triton.autotune(configs=runtime.get_triton_config("masked_select"), key=["n_elements"])
+@triton.autotune(configs=runtime.get_tuned_config("masked_select"), key=["n_elements"])
 @triton.jit
 def masked_select_kernel(
     inp_ptr,

src/flag_gems/ops/max.py

@@ -47,7 +47,7 @@ def heur_block_n(args):
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("max"),
+    configs=runtime.get_tuned_config("max"),
     key=[
         "M",
         "N",

src/flag_gems/ops/mean.py

@@ -60,7 +60,7 @@ def mean(inp, *, dtype=None):
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("mean"),
+    configs=runtime.get_tuned_config("mean"),
     key=["M", "N"],
 )
 @triton.jit

src/flag_gems/ops/min.py

@@ -47,7 +47,7 @@ def heur_block_n(args):
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("min"),
+    configs=runtime.get_tuned_config("min"),
     key=[
         "M",
         "N",

src/flag_gems/ops/mm.py

@@ -12,10 +12,10 @@
 
 @libentry()
 @libtuner(
-    configs=runtime.get_triton_config("mm"),
+    configs=runtime.get_tuned_config("mm"),
     key=["M", "N", "K"],
 )
-@triton.heuristics(runtime.get_heuristics_config("mm"))
+@triton.heuristics(runtime.get_heuristic_config("mm"))
 @triton.jit
 def mm_kernel(
     A,

src/flag_gems/ops/multinomial.py

@@ -11,7 +11,7 @@
 
 
 @libentry()
-@triton.heuristics(runtime.get_heuristics_config("multinomial"))
+@triton.heuristics(runtime.get_heuristic_config("multinomial"))
 @triton.jit(do_not_specialize=["K", "N", "philox_seed", "philox_offset"])
 def multinomial_with_replacement(
     cdf_ptr, out_ptr, K, N, philox_seed, philox_offset, NBLOCK: tl.constexpr

src/flag_gems/ops/mv.py

@@ -12,7 +12,7 @@
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("mv"),
+    configs=runtime.get_tuned_config("mv"),
     key=["M", "N"],
 )
 @triton.jit

src/flag_gems/ops/nonzero.py

@@ -12,7 +12,7 @@
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("nonzero"),
+    configs=runtime.get_tuned_config("nonzero"),
     key=[
         "n_elements",
     ],

src/flag_gems/ops/prod.py

@@ -70,7 +70,7 @@ def heur_block_n(args):
 
 @libentry()
 @triton.autotune(
-    configs=runtime.get_triton_config("prod"),
+    configs=runtime.get_tuned_config("prod"),
     key=[
         "M",
         "N",