Skip to content

list(Vector(buffer_load(vec_width>1))) hangs Python tracing (unbounded recursion, no IR emitted) #793

Description

@coderfeli

Summary

Calling list(...) on a Vector that wraps a multi-lane buffer_ops.buffer_load result (vec_width > 1) makes the DSL hang during Python tracing (kernel emission never completes — no MLIR is ever produced). It should either yield the N scalar elements or raise a clear error; instead it spins / recurses without terminating.

Environment

  • FlyDSL @ current main (repro seen on a 5cb28f6c-based build, gfx950)
  • Path: kernels/preshuffle_gemm.py (fp8/int8 epilogue), and a standalone snippet below.

Trigger (the offending pattern)

from flydsl.expr.typing import T
from flydsl.expr.typing import Vector as Vec

rsrc = fx.buffer_ops.create_buffer_resource(arg_scale_a, max_size=True)
v = fx.buffer_ops.buffer_load(rsrc, fx.Int32(idx), vec_width=4, dtype=T.f32)
vals = list(Vec(v))          # <-- hangs tracing (never returns)

Working alternative (no hang)

Keeping it as a Vector and indexing/operating on it compiles normally:

vals = Vec(v).bitcast(fx.Float32)   # then vals[i] / vector ops — compiles fine

Concrete, verified reproducer (in-tree)

In kernels/preshuffle_gemm.py, the fp8/int8 scale_a epilogue load. This version compiles in ~8s:

scale_a_rsrc = fx.buffer_ops.create_buffer_resource(arg_scale_a, max_size=True)
s_a_vals = [
    Vec(fx.buffer_ops.buffer_load(scale_a_rsrc, fx.Int32(bx_m + mi*16 + lane_div_16*4),
        vec_width=4, dtype=T.f32)).bitcast(fx.Float32)
    for mi in range_constexpr(m_repeat)
]

Changing only the wrapper to list(Vec(...)):

s_a_vals = [
    list(Vec(fx.buffer_ops.buffer_load(scale_a_rsrc, fx.Int32(bx_m + mi*16 + lane_div_16*4),
        vec_width=4, dtype=T.f32)))
    for mi in range_constexpr(m_repeat)
]

…makes the same compile hang for 4+ minutes with no IR emitted (killed). Reproduce via:

FLYDSL_RUNTIME_ENABLE_CACHE=0 COMPILE_ONLY=1 FLYDSL_DUMP_IR=1 FLYDSL_DUMP_DIR=/tmp/x \
python3 tests/kernels/test_preshuffle_gemm.py --in_dtype fp8 -M 2048 -N 4096 -K 512 \
  --tile_m 128 --tile_n 128 --tile_k 256 --use_async_copy --num_iters 2 --num_warmup 0 --no_aiter_bench
# working version: /tmp/x/*/00_origin.mlir appears within seconds
# list() version: no stage files ever appear; process must be killed

A stripped-down standalone kernel doing list(Vec(buffer_load(vec_width=4))) similarly fails to terminate (surfaces as RecursionError: maximum recursion depth exceeded when the trace is small), pointing at unbounded recursion in Vector.__iter__/__len__ when the underlying value is a multi-lane buffer_load result.

Expected behavior

list(Vec(buffer_load(..., vec_width=N))) should return N scalar SSA values, or raise a clear, immediate error — not hang / recurse without bound during tracing.

Impact / suggested fix

Silent multi-minute hangs are a sharp usability footgun (they look like a compiler stall, not a user error). Either make Vector.__iter__/__len__ terminate for buffer_load-backed vectors, or raise a descriptive error directing users to index the Vector (or .bitcast(...)) instead of list()-ing it.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions