Enabling FlyDSL in Helion with support for softmax and rms_norm kernels.

[umechand-amd]

No description provided.

[umechand-amd]

@claude Review the PR based on facts, dont make any assumptions

[claude]

Claude finished @umechand-amd's task in 4m 1s —— View job

---

PR Review: Enabling FlyDSL in Helion (softmax + rms_norm)

Fact-based review of the single commit e5cac7c (diff vs its parent). I read all 11 changed files and cross-checked referenced symbols in the tree. I do not have flydsl installed, so I could not execute test/test_flydsl_reduction.py — the observations below are from static reading only.

• Gather context (read changed files)
• Analyze backend, memory_ops, reduction, tracing changes
• Review tests
• Post review

---

Summary

Adds an experimental FlyDSLBackend (ROCm) in helion/_compiler/backend.py (+896), plus flydsl codegen handlers for load/store (memory_ops.py +641), _reduce (reduce_ops.py), _mask_to (_tracing_ops.py), subscript (view_ops.py), a runtime launcher (runtime/__init__.py), reduction-strategy hooks (reduction_strategy.py), and small shared-path edits (device_ir.py, inductor_lowering.py, backend_registry.py). New test module test/test_flydsl_reduction.py. The backend is registered and marked experimental = True. Overall the change is well-contained to a new backend and gated everywhere by backend.name == "flydsl" / importorskip("flydsl"), so blast radius on other backends is low.

Below are concrete, verifiable findings.

---

Findings

1. Inconsistent tail-mask lane stride between _mask_to and load/store (possible W>1 correctness bug)

_flydsl_col_tail_pred is defined with lane_mod: int = 64 (memory_ops.py). The column base it builds is (offset + (thread_idx.x % lane_mod) * vec).

• In memory_ops.py (both load and store) it is called with lane_mod=_tc if _tc > 0 else 64, where _tc is the reduction _thread_count (= 64*W).
• In _tracing_ops.py (_mask_to codegen) it is called with only 4 positional args, so lane_mod defaults to 64:

  _pred = _flydsl_col_tail_pred(
      state, strategy.offset_var(index), _v, state.sympy_expr(env.block_sizes[index].numel)
  )

For the W>1 regime (_tc = 64*W > 64), the tail predicate used inside _mask_to (thread_idx.x % 64) and the predicate used in load/store (thread_idx.x % _tc) disagree about which column each lane owns. Since test_column_tail_w_gt_1 and test_softmax_cross_wavefront exercise reduction_loops=[512] (W>1) with non-multiple N, this path is reachable. Worth confirming whether the _mask_to call should also pass lane_mod=_tc. Fix this →

2. FastMathFlags namespace differs between files

• reduce_ops.py: {r}.addf(..., fastmath=fmath.FastMathFlags.fast) (uses fmath, i.e. flydsl.expr.math).
• backend.py: w.addf(..., fastmath=arith.FastMathFlags.fast) and reduce(ReductionOp.ADD, fastmath=arith.FastMathFlags.fast) (uses arith, i.e. flydsl.expr.arith).

At most one namespace can be the intended source of FastMathFlags. If FastMathFlags lives only on arith, the reduce_ops.py path (hl.reduce sum) will raise AttributeError at trace time. Please confirm both resolve. (I can't verify against the flydsl package since it isn't installed here.)

3. _reduce min uses .minimumf, which the backend elsewhere says does not exist

reduce_ops.py _reduce min emits:

f"{r} = {r}.minimumf({r}.shuffle_xor({off}, {_WARP}))"

But backend.py states in five places that "flydsl Vector has no .minimumf" and consistently uses -max(-a,-b) (minimum_expr, reduction_combine_expr, the minimum op override, reduction_expr). This is a direct internal contradiction; if the comment is correct, the hl.reduce min path in reduce_ops.py is broken. The reduce_ops.py max path likewise uses .maximumf, which the backend does use, so only min is suspect.

4. Large verbatim duplication between load and store codegen

In memory_ops.py, the two @_decorators.codegen(..., "flydsl") functions duplicate ~150 lines each: the row/col block-id resolution, _flydsl_reduction_col_block handling, redcol_vec/redcol_offset/redcol_pred derivation, and the entire setup is None block (buffer/row/logical_divide/copy_atom, including the two nearly-identical is_vec/else branches). This is a maintenance hazard — the tail-mask bug in #1 is exactly the kind of divergence duplication invites. Recommend extracting a shared helper for the setup + index derivation.

5. customize_ast passes parameters via os.environ mutation

FlyDSLBackend.customize_ast sets os.environ["HELION_ONLINE_TO_3PASS_MIN_N"] = "0" around a call to rewrite_online_to_3pass(hf), then restores it in finally. Using a process-global env var to pass an argument into a function is not thread-safe (concurrent compiles race on the var) and is surprising. If rewrite_online_to_3pass can take min_n as a parameter, prefer passing it directly.

6. Temp-file leak in default_flydsl_launcher

runtime/__init__.py writes a JIT wrapper with tempfile.NamedTemporaryFile(..., delete=False) and never removes fname. Every unique cache_key (per kernel/arity/grid/threads) leaves a _flydsl_jit_*.py file in the temp dir for the process lifetime. Also the module name uses abs(hash(cache_key)), which is not collision-free; consider a monotonic counter or the temp filename stem for uniqueness.

7. Stale comment: max_reduction_threads=1024

In reduction_strategy.py the new flydsl branch comment says "(max_reduction_threads=1024 makes the base value chunk-sized...)", but FlyDSLBackend.max_reduction_threads() returns 64, not 1024. The logic still works (the branch overrides _thread_count regardless), but the parenthetical is factually wrong and will mislead future readers.

8. Repo-guideline: heavy getattr/except usage

CLAUDE.md states: "Do NOT add unnecessary error checks via hasattr, getattr, except, etc." The new code uses many getattr(self, "_flydsl_...", default) guards (e.g. _flydsl_warps_per_row, _flydsl_num_threads, _flydsl_helpers_emitted) and a bare except Exception: in _flydsl_reduction_col_block. Some getattrs guard attributes only set in pre_codegen; those would be cleaner as attributes initialized in __init__ so the access is unconditional. The bare except Exception around env.size_hint(...) swallows all errors — please narrow or remove it.

9. Tests are skipped without flydsl (no coverage on standard CI)

test/test_flydsl_reduction.py starts with pytest.importorskip("flydsl"), so on any runner without the AMD/ROCm flydsl package the entire module is skipped. That's the right gate for an experimental ROCm backend, but it means none of this logic (including the paths in #1–#3) is exercised by the default CI matrix. Please confirm there is an AMD CI lane that actually installs flydsl and runs this file; otherwise regressions here will go unnoticed.

---

Minor / positional notes

• backend.py sigmoid is 1/(1+exp(-x)) (not the numerically-stable two-sided form). Fine for the current softmax/rms_norm kernels; noting for future op coverage.
• FlyDSLOpOverrides.__getattr__ raising BackendUnsupported for un-overridden ops is a nice guard against silent Triton fallthrough — good defensive choice.
• _flydsl_copy_cls cleanly rejecting >128-bit widths (fp32 V=8) with BackendUnsupported, and the matching test_fp32_v8_rejected, is solid.
• Gating the cute_vector_widths spec on backend_name in ("cute", "flydsl") in device_ir.py and the inductor_lowering.py broadcast-expand skip both look correctly scoped to not affect other backends.

What I could not verify

• Runtime correctness / numerical results (no flydsl install here). Findings Minor fix to test file name #1–test the CI #3 are static inconsistencies that would need an AMD run or a maintainer's knowledge of the flydsl API to confirm as bugs vs. false alarms.
  · branch umechand/fly-dsl