Enable yapf for test_line_info.py (#5612)

Signed-off-by: Anatoly Myachev <[email protected]>

.pre-commit-config.yaml

@@ -33,7 +33,6 @@ repos:
     hooks:
       - id: yapf
         args: ["-p", "-i"]
-        exclude: "python/test/unit/language/test_line_info.py"
 
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v19.1.6

pyproject.toml

@@ -9,16 +9,6 @@ each_dict_entry_on_separate_line=false
 split_before_named_assigns = false
 split_complex_comprehension = true
 
-[tool.yapfignore]
-ignore_patterns = [
-    # This exclusion is also specified in .pre-commit-config.yaml.
-    # - We put it here because if you run yapf directly, we want it to skip the
-    #   file.
-    # - We also put it in .pre-commit-config because yapf raises an error if
-    #   pre-commit runs it but all of the files it might touch are ignored!
-    "python/test/unit/language/test_line_info.py"
-]
-
 [tool.ruff]
 line-length = 120
 

python/test/unit/language/test_line_info.py

@@ -9,9 +9,7 @@
 
 
 @triton.jit
-def kernel_single(X,
-                  Y,
-                  BLOCK: tl.constexpr):
+def kernel_single(X, Y, BLOCK: tl.constexpr):
     x = tl.load(X + tl.arange(0, BLOCK))
     tl.store(Y + tl.arange(0, BLOCK), x)
 
@@ -22,9 +20,7 @@ def device_inline(x):
 
 
 @triton.jit
-def kernel_call(X,
-                Y,
-                BLOCK: tl.constexpr):
+def kernel_call(X, Y, BLOCK: tl.constexpr):
     x = tl.load(X + tl.arange(0, BLOCK))
     y = device_inline(x)
     tl.store(Y + tl.arange(0, BLOCK), y)
@@ -154,38 +150,38 @@ def test_line_info(func: str):
     shape = (128, )
     kernel_info = {}
     if func == "single":
-        kernel_info = kernel_single.warmup(torch.float32, torch.float32, BLOCK=shape[0], grid=(1,))
+        kernel_info = kernel_single.warmup(torch.float32, torch.float32, BLOCK=shape[0], grid=(1, ))
     elif func == "call":
-        kernel_info = kernel_call.warmup(torch.float32, torch.float32, BLOCK=shape[0], grid=(1,))
+        kernel_info = kernel_call.warmup(torch.float32, torch.float32, BLOCK=shape[0], grid=(1, ))
     elif func == "call_noinline":
-        kernel_info = kernel_call_noinline.warmup(torch.float32, torch.float32, BLOCK=shape[0], grid=(1,))
+        kernel_info = kernel_call_noinline.warmup(torch.float32, torch.float32, BLOCK=shape[0], grid=(1, ))
     elif func == "autotune":
-        kernel_info = kernel_autotune.warmup(torch.float32, torch.float32, SIZE=shape[0], grid=(1,))[0]
+        kernel_info = kernel_autotune.warmup(torch.float32, torch.float32, SIZE=shape[0], grid=(1, ))[0]
     elif func == "dot_combine":
-        kernel_info = kernel_dot_combine.warmup(20, grid=(1,))
+        kernel_info = kernel_dot_combine.warmup(20, grid=(1, ))
     elif func == "cdiv":
-        kernel_info = kernel_cdiv.warmup(20, grid=(1,))
+        kernel_info = kernel_cdiv.warmup(20, grid=(1, ))
 
     file_lines = extract_file_lines(command, anchor, separator, kernel_info.asm[obj_kind])
     if func == "single":
-        assert (check_file_lines(file_lines, "test_line_info.py", 15))
-        assert (check_file_lines(file_lines, "test_line_info.py", 16))
+        assert (check_file_lines(file_lines, "test_line_info.py", 13))
+        assert (check_file_lines(file_lines, "test_line_info.py", 14))
     elif func == "call":
-        assert (check_file_lines(file_lines, "test_line_info.py", 28))
-        assert (check_file_lines(file_lines, "test_line_info.py", 30))
+        assert (check_file_lines(file_lines, "test_line_info.py", 24))
+        assert (check_file_lines(file_lines, "test_line_info.py", 26))
     elif func == "call_noinline":
-        assert (check_file_lines(file_lines, "test_line_info.py", 42))
-        assert (check_file_lines(file_lines, "test_line_info.py", 35))
-        assert (check_file_lines(file_lines, "test_line_info.py", 37))
+        assert (check_file_lines(file_lines, "test_line_info.py", 38))
+        assert (check_file_lines(file_lines, "test_line_info.py", 31))
+        assert (check_file_lines(file_lines, "test_line_info.py", 31))
     elif func == "autotune":
-        assert (check_file_lines(file_lines, "test_line_info.py", 53))
-        assert (check_file_lines(file_lines, "test_line_info.py", 54))
-        assert (check_file_lines(file_lines, "test_line_info.py", 55))
+        assert (check_file_lines(file_lines, "test_line_info.py", 49))
+        assert (check_file_lines(file_lines, "test_line_info.py", 50))
+        assert (check_file_lines(file_lines, "test_line_info.py", 51))
     elif func == "dot_combine":
-        assert (check_file_lines(file_lines, "test_line_info.py", 65))
-        assert (check_file_lines(file_lines, "test_line_info.py", 66, should_contain=False))
+        assert (check_file_lines(file_lines, "test_line_info.py", 61))
+        assert (check_file_lines(file_lines, "test_line_info.py", 62, should_contain=False))
     elif func == "cdiv":
-        assert (check_file_lines(file_lines, "test_line_info.py", 75))
+        assert (check_file_lines(file_lines, "test_line_info.py", 71))
 
 
 @pytest.mark.interpreter
@@ -201,18 +197,18 @@ def test_line_info_interpreter(func: str):
         expected_def_lineno = 12
     elif func == "call":
         kernel = kernel_call
-        expected_def_lineno = 25
+        expected_def_lineno = 23
     elif func == "call_noinline":
         kernel = kernel_call_noinline
-        expected_def_lineno = 41
+        expected_def_lineno = 37
     elif func == "autotune":
         kernel = kernel_autotune.fn
-        expected_def_lineno = 52
+        expected_def_lineno = 48
     elif func == "dot_combine":
         kernel = kernel_dot_combine
-        expected_def_lineno = 62
+        expected_def_lineno = 58
     elif func == "cdiv":
         kernel = kernel_cdiv
-        expected_def_lineno = 72
+        expected_def_lineno = 68
     kernel.rewrite()
     assert kernel.rewriter.def_file_lineno == expected_def_lineno