feat(jit): two-phase compilation for autotuning compile_hints by fsx950223 · Pull Request #266 · ROCm/FlyDSL

fsx950223 · 2026-03-23T07:09:44Z

Motivation

During autotuning, FlyDSL recompiles the entire MLIR pipeline for each waves_per_eu/maxnreg configuration. Most passes (layout lowering, SCF-to-CF, GPU-to-ROCDL, etc.) produce identical IR regardless of these backend flags — only the final gpu-module-to-binary pass differs. This wastes ~30s per redundant full compilation.

Split out from #265 per review feedback.

Technical Details

Two-phase compilation

Splits MlirCompiler.compile() into two reusable phases:

Phase 1 (compile_pre_binary): all MLIR passes except gpu-module-to-binary. Result is cached as ASM text in _pre_binary_cache.
Phase 2 (compile_binary): only gpu-module-to-binary with the specific waves_per_eu/maxnreg flags from compile_hints.

Both _pre_binary_fragments and _binary_fragment are derived from _pipeline_fragments() (slice off last element) — no pass list duplication.

Cache key changes

_make_cache_key now includes __compile_hints__ so different compiler options produce distinct cache entries.
_make_base_cache_key (new) produces the key WITHOUT compile_hints — used for pre-binary module lookup across autotuning configs.
_pre_binary_cache uses (chip, base_key) composite key for multi-arch safety.
_pre_binary_cache is guarded by use_cache — respects FLYDSL_RUNTIME_ENABLE_CACHE=0.

Review feedback addressed (from #265)

Feedback	Resolution
DRY: pass list duplicated 3 times	`_pre_binary_fragments = _pipeline_fragments()[:-1]`
DRY: `_make_base_cache_key` is a copy	`_make_cache_key` builds on `_make_base_cache_key`
`_pre_binary_cache` has no `use_cache` guard	Guarded by `if use_cache`
`base_key` doesn't include `chip`	Uses `(chip, base_key)` composite key

Test Plan

Autotuning with multiple waves_per_eu configs: Phase 1 compiles once, Phase 2 runs per config
FLYDSL_RUNTIME_ENABLE_CACHE=0: no pre-binary caching, full recompilation each call
FLYDSL_RUNTIME_ENABLE_CACHE=1 (default): pre-binary cached, fast Phase 2 only for subsequent configs
No compile_hints (standard path): falls through to original MlirCompiler.compile()

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Split MlirCompiler.compile() into two phases: - Phase 1 (compile_pre_binary): all MLIR passes except gpu-module-to-binary - Phase 2 (compile_binary): only gpu-module-to-binary with waves_per_eu/maxnreg This allows the autotuner to compile Phase 1 once and reuse the pre-binary module across different compile_hints configs, avoiding redundant full recompilations (~30s each). Also adds compile_hints to the JIT cache key so different compiler options produce distinct cache entries. Addresses review feedback from PR #265: - _pre_binary_fragments derived from _pipeline_fragments (no duplication) - _make_cache_key builds on shared _make_base_cache_key helper - _pre_binary_cache guarded by use_cache flag - chip included in pre-binary cache key for multi-arch safety Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Copilot

Pull request overview

This PR optimizes FlyDSL autotuning by splitting JIT compilation into a reusable “pre-binary” phase (all MLIR lowering passes) and a per-config “binary” phase (final gpu-module-to-binary), reducing redundant recompilation when only backend compile hints change.

Changes:

Adds MlirCompiler.compile_pre_binary() and MlirCompiler.compile_binary() to support two-phase compilation driven by CompilationContext compile hints.
Introduces _make_base_cache_key() and extends _make_cache_key() to include __compile_hints__, enabling correct artifact caching across different tuning configurations.
Adds an in-memory _pre_binary_cache keyed by (chip, base_cache_key) to reuse phase-1 results across autotuning configs when runtime caching is enabled.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-23T07:14:35Z

+        all_frags = MlirCompiler._pipeline_fragments(chip=chip)
+        if all_frags and isinstance(all_frags[-1], str) and all_frags[-1].startswith('gpu-module-to-binary'):
+            return all_frags[:-1]
+        return all_frags


_pre_binary_fragments silently falls back to returning the full pipeline when the last fragment isn’t gpu-module-to-binary. That would make “pre-binary” compilation accidentally include the binary pass (and then Phase 2 would run it again), which is hard to debug. Consider raising (or at least logging) if the expected last fragment isn’t gpu-module-to-binary, mirroring _binary_fragment’s strictness.

Suggested change

return all_frags

raise RuntimeError('Last pipeline fragment is not gpu-module-to-binary')

Copilot · 2026-03-23T07:14:35Z

+            # Two-phase compilation: reuse pre-binary module for different compile_hints.
+            # This avoids redundant full recompilations during autotuning when only
+            # waves_per_eu/maxnreg differ.
+            _hints = CompilationContext.get_compile_hints()


The two-phase path in call bypasses MlirCompiler.compile(), so debug behaviors implemented there (e.g. FLYDSL_DUMP_IR stage dumps and PassManager.enable_ir_printing via env.debug.print_after_all) won’t run when compile_hints are set. If those debug flags are expected to work during autotuning/debugging, consider disabling two-phase compilation when env.debug.dump_ir/print_after_all is enabled, or extending compile_pre_binary/compile_binary to honor the same debug pathways.

Copilot · 2026-03-23T07:14:35Z

+    def compile_binary(cls, pre_binary_asm: str, *, chip: str = None) -> ir.Module:
+        """Phase 2: Run gpu-module-to-binary on a pre-binary module ASM."""
+        if chip is None:
+            chip = env.compile.arch or get_rocm_arch()
+        module = ir.Module.parse(pre_binary_asm)
+        frag = cls._binary_fragment(chip=chip)
+        pipeline = f"builtin.module({frag})"
+        pm = PassManager.parse(pipeline)
+        pm.enable_verifier(env.debug.enable_verifier)
+        pm.run(module.operation)


compile_binary parses pre_binary_asm and runs passes without verifying the parsed module first. Since pre_binary_asm may come from an in-memory cache (or could be corrupted/unsupported across versions), adding module.operation.verify() after parsing would make failures deterministic and easier to diagnose.

Copilot · 2026-03-23T07:14:36Z

        self._call_state_cache = {}  # cache_key -> CallState
        self._sig = None  # lazy: set on first call
        self._mem_cache = {}
+        self._pre_binary_cache = {}  # (chip, base_cache_key) -> (pre_binary_asm, chip)


_pre_binary_cache stores (pre_binary_asm, chip) as the value, but chip is already part of the dict key. Keeping only pre_binary_asm would avoid redundant state and prevent inconsistencies if these ever diverge.

Copilot · 2026-03-23T07:14:36Z

+    def _make_cache_key(self, bound_args):
+        """Cache key including compile_hints for full artifact caching."""
+        key_parts = list(self._make_base_cache_key(bound_args))
+        from .kernel_function import CompilationContext
+        hints = CompilationContext.get_compile_hints()
+        if hints:
+            key_parts.append(('__compile_hints__', tuple(sorted(hints.items()))))
+        return tuple(key_parts)


This PR changes cache key semantics by adding compile_hints and introduces a new “base key” used for pre-binary reuse. There doesn’t appear to be a unit test asserting (1) base key is unchanged by CompilationContext.compile_hints, and (2) full cache key changes when hints change. Adding a small pytest for these invariants would help prevent regressions without requiring a GPU run.

Copilot AI review requested due to automatic review settings March 23, 2026 07:09

Copilot started reviewing on behalf of fsx950223 March 23, 2026 07:11

View session

Copilot AI reviewed Mar 23, 2026

View reviewed changes

coderfeli closed this Apr 1, 2026

jhinpan mentioned this pull request Jun 2, 2026

Discussion: make autotune a supported and broadly usable tuning path #612

Closed

This was referenced Jun 30, 2026

[Feature] Make autotune a safe, adopted path #770

Open

[1/5] autotune: harden cache key + add restore_value (#770) #783

Merged

Navigation Menu

Repository navigation

feat(jit): two-phase compilation for autotuning compile_hints - #266

feat(jit): two-phase compilation for autotuning compile_hints#266
fsx950223 wants to merge 1 commit into
mainfrom
feat/two-phase-compilation

fsx950223 commented Mar 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	return all_frags
	raise RuntimeError('Last pipeline fragment is not gpu-module-to-binary')

Conversation

fsx950223 commented Mar 23, 2026

Motivation

Technical Details

Two-phase compilation

Cache key changes

Review feedback addressed (from #265)

Test Plan

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants