flagos-ai / flagos-ai/FlagTree
[BUG][mthreads] llc non-termination on giant IR: FFT N=250 naive DFT baseline compile hangs >15min
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
## Environment
- Triton 3.6.0 (FlagTree mthreads3.6), MUSA backend, MTT S5000
- torch_musa 2.7.1, MUSA 4.3.5, Python 3.10
## Symptom
Compiling the original FFT kernel `dft_naive_interleaved` (39_Fast_Fourier_transform,
naive O(N²) DFT path, N=250 non-power-of-2 → functional[3]) hangs the MTGPU backend
at compile time. The Triton frontend completes in seconds (TTIR 12.7KB / TTGIR
16.4KB), but the `llc` subprocess (`-march=mtgpu -mcpu=mp_31 -O2`) spins at 99% CPU
and never returns — `kernel.o` stays 0 bytes, no error, no crash.
The IR reaching llc is a giant 8.3MB single function: the 64×128 fp64 sin/cos block
is scalarized and each lane inlines the LLVM libm implementation
(`__kernel_rem_pio2`, ~12000 inlined copies). Compile time grows super-linearly
(≈quadratic) with sin/cos lane count — 512→12s, 1024→26s, 2048→70s, 4096→228s,
8192 (the N=250 shape)→>900s and still not done — so the N=250 naive path is
effectively uncompilable. It is a path threshold, not an N threshold: power-of-2 N
cases take the small Cooley-Tukey kernel (compiles in milliseconds).
**FlagTree mthreads3.6 output:**
>[control saxpy] DONE rc=0 elapsed=6.85s
>[repro fft N=250 run 1] HANG (90.08s timeout)
>[repro fft N=250 run 2] HANG (90.06s timeout)
>[repro fft N=250 run 3] HANG (90.08s timeout)
**Triton3.7 RTX3070 output:**
>[control saxpy] DONE rc=0 elapsed=5.82s
>[repro fft N=250 run 1] DONE elapsed=11.76s
>[repro fft N=250 run 2] DONE elapsed=3.18s
>[repro fft N=250 run 3] DONE elapsed=3.03s
## Minimal reproducer
```
KERNEL_SRC = '''
import torch
import triton
import triton.language as tl
import math
@triton.jit
def dft_naive_interleaved(
input,
spectrum,
N,
BLOCK_SIZE: tl.constexpr,
RUNNING_BLOCK_SIZE: tl.constexpr,
):
program_id = tl.program_id(0)
offset = program_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offset < N
sum_r = tl.zeros([BLOCK_SIZE], dtype=tl.float64)
sum_i = tl.zeros([BLOCK_SIZE], dtype=tl.float64)
offset_fp64 = offset.to(tl.float64)
for running_offset in range(0, N, RUNNING_BLOCK_SIZE):
offset_tmp = running_offset + tl.arange(0, RUNNING_BLOCK_SIZE)
running_mask = offset_tmp < N
running_offset_orig = 2 * offset_tmp
xr = tl.load(input + running_offset_orig + 0, mask=running_mask, other=0.0).to(
tl.float64
)
xi = tl.load(input + running_offset_orig + 1, mask=running_mask, other=0.0).to(
tl.float64
)
# -2pi * k * n / N
offset_tmp_fp64 = offset_tmp.to(tl.float64)
angle = -(2.0 * math.pi / tl.full([], N, tl.float64)) * (
offset_fp64[:, None] * offset_tmp_fp64[None, :]
)
c = tl.cos(angle)
s = tl.sin(angle)
tr = xr[None, :] * c - xi[None, :] * s
ti = xr[None, :] * s + xi[None, :] * c
valid_addition_mask = mask[:, None] & running_mask[None, :]
tr = tl.where(valid_addition_mask, tr, 0.0)
ti = tl.where(valid_addition_mask, ti, 0.0)
sum_r += tl.sum(tr, axis=1)
sum_i += tl.sum(ti, axis=1)
out_offset = 2 * offset
tl.store(spectrum + out_offset + 0, sum_r.to(tl.float32), mask=mask)
tl.store(spectrum + out_offset + 1, sum_i.to(tl.float32), mask=mask)
# signal and spectrum are tensors on the GPU
def solve(signal: torch.Tensor, spectrum: torch.Tensor, N: int):
BLOCK_SIZE = 64
RUNNING_BLOCK_SIZE = 128
grid = (triton.cdiv(N, BLOCK_SIZE),)
dft_naive_interleaved[grid](signal, spectrum, N, BLOCK_SIZE, RUNNING_BLOCK_SIZE)
'''
CONTROL_SRC = '''
import torch
import triton
import triton.language as tl
@triton.jit
def saxpy(x, y, out, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
m = offs < n
xv = tl.load(x + offs, mask=m)
yv = tl.load(y + offs, mask=m)
tl.store(out + offs, 2.0 * xv + yv, mask=m)
def solve(x, y, out, n):
saxpy[(n // 1024,)](x, y, out, n, BLOCK=1024)
'''
MAIN_CODE = """
import os, sys, time
import torch
def main():
which = sys.argv[1]
print("STAGE init which=" + which, flush=True)
dev = torch.device("musa")
t0 = time.time()
if which == "fft250":
N = 250
signal = torch.zeros(2 * N, dtype=torch.float32, device=dev)
spectrum = torch.empty(2 * N, dtype=torch.float32, device=dev)
print("STAGE compile-start", flush=True)
solve(signal, spectrum, N)
torch.musa.synchronize()
print(f"STAGE done elapsed={time.time()-t0:.2f}s", flush=True)
elif which == "saxpy":
n = 4096
x = torch.randn(n, device=dev)
y = torch.randn(n, device=dev)
out = torch.empty(n, device=dev)
print("STAGE compile-start", flush=True)
solve(x, y, out, n)
torch.musa.synchronize()
print(f"STAGE done elapsed={time.time()-t0:.2f}s", flush=True)
else:
raise SystemExit("unknown mode " + which)
main()
"""
def run_child(which: str, timeout: float, dump_dir: str = ""):
"""Run the kernel compile/launch in a subprocess; returns (verdict, rc, log, elapsed)."""
env = dict(os.environ)
env["MUSA_VISIBLE_DEVICES"] = "4"
if dump_dir:
env["TRITON_DUMP_DIR"] = dump_dir
env["TRITON_KERNEL_DUMP"] = "1"
env["TRITON_ALWAYS_COMPILE"] = "1"
# Triton requires the @triton.jit source to live in a real .py file
# (inspect.getsourcelines), so materialize kernel source + main() directly.
src = KERNEL_SRC if which == "fft250" else CONTROL_SRC
child_file = os.path.join(
tempfile.gettempdir(), f"repro_000007_child_{os.getpid()}.py"
)
with open(child_file, "w") as f:
f.write(src + MAIN_CODE)
proc = subprocess.Popen(
[sys.executable, child_file, which],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
env=env,
text=True,
start_new_session=True, # own process group so we can kill llc children too
)
t0 = time.time()
try:
out, _ = proc.communicate(timeout=timeout)
elapsed = time.time() - t0
return ("DONE", proc.returncode, out, elapsed)
except subprocess.TimeoutExpired:
elapsed = time.time() - t0
# kill the whole process group (python + llc child), then drain
try:
os.killpg(proc.pid, signal.SIGKILL)
except ProcessLookupError:
pass
out, _ = proc.communicate()
return ("HANG", None, out, elapsed)
finally:
try:
os.remove(child_file)
except OSError:
pass
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--times", type=int, default=3, help="repro repetitions (default 3)")
ap.add_argument("--timeout", type=float, default=90.0, help="per-run timeout s (default 90)")
ap.add_argument("--dump-dir", default="", help="TRITON_DUMP_DIR for IR evidence")
args = ap.parse_args()
print("=" * 72)
print("musa-issue-000007 one-click repro: FFT naive DFT llc compile HANG")
print(f"timeout={args.timeout}s times={args.times} dump_dir={args.dump_dir or '(none)'}")
print("=" * 72)
# 1) control: trivial kernel must compile+run fast
verdict, rc, out, ctrl_el = run_child("saxpy", args.timeout)
print(f"[control saxpy] {verdict} rc={rc} elapsed={ctrl_el:.2f}s")
sys.stdout.write(out)
if verdict != "DONE":
print("CONTROL FAILED — environment problem, cannot judge. Aborting.")
return 2
results = []
for i in range(args.times):
verdict, rc, out, el = run_child("fft250", args.timeout, args.dump_dir)
results.append(verdict)
print(f"[repro fft N=250 run {i+1}] {verdict} (timeout={args.timeout}s, elapsed={el:.2f}s)")
sys.stdout.write(out)
n_hang = results.count("HANG")
print("=" * 72)
print(f"RESULT: {n_hang}/{args.times} HANG reproduced "
f"(control saxpy compiled+ran in {ctrl_el:.2f}s within the same {args.timeout}s timeout)")
print("=" * 72)
if n_hang == args.times:
print("musa-issue-000007 REPRODUCED: naive DFT (N=250) compile hangs in llc.")
return 0
print("NOT reproduced this run.")
return 1
if __name__ == "__main__":
sys.exit(main())
```
Contributor guide
Research direction
Start with the provided dft_naive_interleaved reproducer and compare its generated 8.3MB IR with the llc invocation on the MUSA backend; the saxpy control establishes that the environment works. Done means the N=250 naive DFT path no longer hangs in llc, with the supplied timeout-based reproducer completing successfully.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 42/100