Compiler and pipeline
=====================

FlyDSL includes a JIT compiler that traces Python kernel functions into MLIR
and lowers them through the Fly dialect pipeline to GPU binaries.

``@flyc.kernel`` and ``@flyc.jit``
------------------------------------

The primary API for defining and compiling kernels:

.. code-block:: python

   import flydsl.compiler as flyc
   import flydsl.expr as fx

   @flyc.kernel
   def my_kernel(A: fx.Tensor, B: fx.Tensor, n: fx.Constexpr[int]):
       tid = fx.thread_idx.x
       bid = fx.block_idx.x
       # ... kernel body using layout ops ...

   @flyc.jit
   def launch(A: fx.Tensor, B: fx.Tensor, n: fx.Constexpr[int],
              stream: fx.Stream = fx.Stream(None)):
       my_kernel(A, B, n).launch(
           grid=(grid_x, 1, 1),
           block=(256, 1, 1),
           stream=stream,
       )

- ``@flyc.kernel`` compiles the function body into a ``gpu.func`` inside a
  ``gpu.module``. It uses AST rewriting to trace Python code into MLIR IR.
- ``@flyc.jit`` wraps a host-side function that constructs and launches kernels.
  On first call it triggers JIT compilation; subsequent calls with the same type
  signature use a cached compiled artifact.

Compilation flow
-----------------

On first call, ``@flyc.jit`` runs the following pipeline:

1. **AST rewriting**: The Python source is parsed and rewritten to emit MLIR ops.
2. **MLIR module construction**: The kernel body is traced into ``fly``, ``gpu``,
   ``arith``, ``scf``, ``memref``, and ``vector`` dialect ops.
3. **Fly pass pipeline**: The module is lowered through three pass stages,
   defined in ``RocmBackend._pipeline_parts()``
   (``python/flydsl/compiler/backends/rocm.py``). See
   :doc:`../architecture_guide` §3 for the per-pass table.

   A. ``pre_binary_fragments`` (Fly → ROCDL):

      - ``fly-rewrite-func-signature``
      - ``fly-canonicalize``
      - ``fly-layout-lowering``
      - ``fly-int-swizzle-simplify``
      - ``canonicalize``
      - ``fly-convert-atom-call-to-ssa-form``
      - ``fly-promote-regmem-to-vectorssa``
      - ``convert-fly-to-rocdl``
      - ``canonicalize``
      - ``gpu.module(convert-scf-to-cf, cse, convert-gpu-to-rocdl{chipset=gfxNNN ...}, fly-rocdl-cluster-attr)``

   B. ``binary_prep_fragments`` (→ LLVM):

      - ``rocdl-attach-target{chip=gfxNNN ...}``
      - ``convert-scf-to-cf``
      - ``convert-cf-to-llvm``
      - ``gpu-to-llvm{use-bare-pointers-...=true}``
      - ``convert-vector-to-llvm``
      - ``convert-arith-to-llvm``
      - ``convert-func-to-llvm``
      - ``reconcile-unrealized-casts``
      - ``ensure-debug-info-scope-on-llvm-func`` (optional, gated by ``FLYDSL_DEBUG_ENABLE_DEBUG_INFO``)

   C. ``binary_fragment``:

      - ``gpu-module-to-binary{format=fatbin opts="..."}``

4. **Cached artifact**: The compiled binary is cached to disk
   (``~/.flydsl/cache/``) keyed by the compiler toolchain hash and kernel
   type signature.

Tensor arguments
-----------------

Use ``flyc.from_dlpack`` to convert PyTorch tensors into FlyDSL tensor
descriptors with layout metadata:

.. code-block:: python

   import flydsl.compiler as flyc

   tA = flyc.from_dlpack(torch_tensor).mark_layout_dynamic(
       leading_dim=0, divisibility=4
   )
   launch(tA, B, n, stream=torch.cuda.Stream())

ROCDL operations
-----------------

The ``flydsl.expr.rocdl`` module provides AMD-specific operations:

- **fx.rocdl.make_buffer_tensor** -- create buffer resource descriptor from tensor (CDNA buffer copy)
- **fx.rocdl.BufferCopy32b** / **BufferCopy128b** -- buffer copy atoms
- **fx.rocdl.MFMA** -- MFMA instruction atoms (CDNA3/CDNA4; for example, ``MFMA(16, 16, 4, fx.Float32)``)
- **fx.rocdl.WMMA** / **fx.rocdl.WMMAScale** -- wave32 WMMA MMA atoms; ``WMMA`` is arch-dispatched (gfx11 / gfx120x RDNA4 / gfx1250), ``WMMAScale`` is the gfx1250 E8M0 MX-scaled form
- **fx.rocdl.make_tdm_atom** / **fx.rocdl.TDM** -- gfx1250 TDM async Global↔LDS whole-tile copy atom (1–5D; base from the copy operand, per-dim extent/stride/imm_offset/mask as atom state)

fly-opt CLI
------------

The ``fly-opt`` tool is a command-line interface for running MLIR passes on
``.mlir`` files:

.. code-block:: bash

   fly-opt --fly-canonicalize input.mlir
   fly-opt --fly-layout-lowering input.mlir
   fly-opt --help