Weird AST constructor issue with mode="max-autotune" with python 3.11

[Chillee]

🐛 Describe the bug

Traceback (most recent call last):
  File "/data/users/chilli/pytorch/torch/_inductor/triton_heuristics.py", line 383, in _precompile_config
    binary = triton.compile(*compile_args, **compile_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/chilli/triton_source/python/triton/compiler/compiler.py", line 268, in compile
    module = src.make_ir(options, codegen_fns, context)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/chilli/triton_source/python/triton/compiler/compiler.py", line 112, in make_ir
    return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/chilli/triton_source/python/triton/runtime/jit.py", line 553, in parse
    raise e
  File "/data/users/chilli/triton_source/python/triton/runtime/jit.py", line 551, in parse
    tree = ast.parse(self.src)
           ^^^^^^^^^^^^^^^^^^^
SystemError: AST constructor recursion depth mismatch (before=45, after=42)
backend='inductor' raised:

Repro:

import torch
import os
import sys
import shutil

directory_path = '/tmp/torchinductor_chilli/'
if os.path.exists(directory_path) and os.path.isdir(directory_path):
    shutil.rmtree(directory_path)
import torch._inductor.config
torch.set_default_device('cuda')

torch._dynamo.config.automatic_dynamic_shapes = False
# Needed since changing args to function causes recompiles
torch._dynamo.config.cache_size_limit = 1000

@torch.compile(mode="max-autotune-no-cudagraphs")
def f(a, b):
    return torch.mm(a, b)

try:
    for N in range(512, 1024, 16):
        print(N)
        a = torch.randn(N, N, dtype=torch.bfloat16)
        b = torch.randn(N, N, dtype=torch.bfloat16)
        f(a, b)
except Exception as e:
    print(e)
    sys.exit(1)
sys.exit(0)

I bisected back to this commit: #124030

cc: @eellison

cc: @cpuhrsch

Setting TORCHINDUCTOR_COMPILE_THREADS=1 also fixes it.

Versions

python=3.11

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang

[eellison]

I just setup a pytorch-311 build and can't repro this. Python 3.11.9. Do you have to run multiple times to repro or does it happen consistently?

[Chillee]

@eellison hmm.. this script repros it consistently for me. Like, when you run it, it goes through compilation for all the shapes without erroring?

[eellison]

That's correct

[cpuhrsch]

I tried to get it to repo on CI in #125331 but it seems to have green / with an unrelated failure. The error is definitely real though, because we see it in AO here: pytorch/ao#197

[eellison]

What python version are you specifically ?

[eellison]

There is an AST Parsing bug that was fixed in 3.11.8: pytest-dev/pytest#11724 (comment).

I'm using 3.11.9 here so maybe that explains it.

Can @Chillee or @cpuhrsch give your 3.11 versions ?

[Chillee]

Yeah i'm on 3.11.7

[eellison]

I guess we can disable precompilation from 3.11.0 to 3.11.7 and print a warning.

which has been around in that form in python3.11 releases since 3.11.0b4 up to and including 3.11.7.