[Fix] fix some issues with JIT decorators existing in the examples (#681)

* [Fix] fix some issues with JIT decorators existing in the examples

* format

* Uses PassConfigKey instand of str

---------
Co-authored-by: Cunxiao <nicunxiao@bytedance.com>

examples/convolution/example_convolution.py

@@ -25,6 +25,7 @@ def ref_program(stride, padding, dilation):
     return main
 
 
+@tilelang.jit(out_idx=[2])
 def convolution(N,
                 C,
                 H,
@@ -116,8 +117,7 @@ def main(argv=None):
     block_k = 32
     num_stages = 3
     threads = 256
-    program = convolution(N, C, H, W, F, K, S, D, P, block_m, block_n, block_k, num_stages, threads)
-    kernel = tilelang.compile(program, out_idx=[2])
+    kernel = convolution(N, C, H, W, F, K, S, D, P, block_m, block_n, block_k, num_stages, threads)
 
     out_c = kernel(a, b)
     ref_c = ref_program(S, P, D)(a, b)

examples/convolution/example_convolution_autotune.py

@@ -32,10 +32,7 @@ def ref_program(stride, padding, dilation):
 
 def get_configs(N, C, H, W, F, K, S, D, P, with_roller=False, topk=15):
     if with_roller:
-        if torch.version.hip is not None:
-            arch=CDNA("hip")
-        else:
-            arch = CUDA("cuda")
+        arch = CDNA("hip") if torch.version.hip is not None else CUDA("cuda")
         carve_template = ConvTemplate(
             N=N,
             C=C,
@@ -102,6 +99,7 @@ def get_configs(N, C, H, W, F, K, S, D, P, with_roller=False, topk=15):
 
 def get_best_config(N, C, H, W, F, K, S, D, P, ref_prog, with_roller=False):
 
+    @tilelang.jit(out_idx=[2])
     def kernel(
         block_M=None,
         block_N=None,
@@ -212,6 +210,7 @@ def get_heuristic_config() -> dict:
         }
 
 
+@tilelang.jit(out_idx=[2])
 def convolution(N,
                 C,
                 H,
@@ -302,7 +301,7 @@ def main(n: int = 128,
         kernel = result.kernel
     else:
         config = get_heuristic_config()
-        kernel = tilelang.compile(convolution(N, C, H, W, F, K, S, D, P, **config), out_idx=[2])
+        kernel = convolution(N, C, H, W, F, K, S, D, P, **config)
 
     profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Auto)
     tilelang_latency = profiler.do_bench()