..  Licensed to the Apache Software Foundation (ASF) under one
    or more contributor license agreements.  See the NOTICE file
    distributed with this work for additional information
    regarding copyright ownership.  The ASF licenses this file
    to you under the Apache License, Version 2.0 (the
    "License"); you may not use this file except in compliance
    with the License.  You may obtain a copy of the License at

..    http://www.apache.org/licenses/LICENSE-2.0

..  Unless required by applicable law or agreed to in writing,
    software distributed under the License is distributed on an
    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.

.. _fusion-arch:

Operator Fusion
===============

Operator fusion is one of the most impactful optimizations in TVM. Instead of launching one kernel
per operator (e.g., conv2d, bias_add, relu), fusion merges multiple operators into a single kernel,
eliminating intermediate memory allocations and kernel launch overhead.

TVM provides two complementary fusion mechanisms:

- **Automatic fusion** (``FuseOps`` + ``FuseTIR``): groups operators based on their computational
  patterns using a post-dominator analysis algorithm.
- **Pattern-based fusion** (``FuseOpsByPattern``): groups operators that match user-defined
  dataflow patterns, typically for offloading to external backends (cuBLAS, CUTLASS, DNNL, etc.).

Both produce the same output: Relax functions marked with ``Primitive=True`` that are later
lowered to fused TIR kernels or dispatched to external libraries.

Overview
--------

Fusion involves three passes:

.. code-block:: text

   IRModule (after LegalizeOps)
        │
        ▼  AnnotateTIROpPattern        ← label each op (elementwise, reduce, etc.)
   IRModule (annotated)
        │
        ▼  FuseOps                     ← group ops into fused Relax functions
   IRModule (with fused functions marked Primitive=True)
        │
        ▼  FuseTIR                     ← merge TIR PrimFuncs inside each group
   IRModule (fused TIR kernels)

In the compilation pipeline, these passes appear in the backend-specific ``legalize_passes``
phase. For example, the CUDA pipeline (``python/tvm/relax/backend/cuda/pipeline.py``) runs:

.. code-block:: python

   LegalizeOps()          # lower Relax ops to call_tir
   AnnotateTIROpPattern() # annotate pattern kinds
   FoldConstant()
   FuseOps()              # group ops
   FuseTIR()              # merge TIR functions


Operator Pattern Classification
-------------------------------

Before fusion, ``AnnotateTIROpPattern`` analyzes each TIR function in the module and assigns
an ``OpPatternKind``. The fusion algorithm uses these pattern kinds to decide which operators
can be fused together.

.. list-table::
   :header-rows: 1
   :widths: 20 10 70

   * - Pattern Kind
     - Value
     - Description
   * - ``kElemWise``
     - 0
     - Elementwise: one-to-one input/output mapping (e.g., ``add``, ``relu``, ``exp``).
   * - ``kBroadcast``
     - 1
     - Broadcasting: output axes map to input axes in order, but some input axes may be
       broadcast (e.g., ``bias_add``). Note: ``transpose`` is **not** broadcast because axes
       are reordered.
   * - ``kInjective``
     - 2
     - Injective: each output element depends on a single input element, but the mapping may
       be non-trivial (e.g., ``reshape``, ``concatenate``, ``transpose``).
   * - ``kCommReduce``
     - 3
     - Communicative reduction: output elements aggregate over input elements
       (e.g., ``sum``, ``max``, ``mean``).
   * - ``kOutEWiseFusable``
     - 4
     - Complex operation whose output can accept elementwise followers, but cannot chain
       with another complex op (e.g., ``conv2d``, ``matmul``, ``dense``).
   * - ``kTuple``
     - 7
     - Tuple node. Can fuse into subsequent injective ops but is treated specially.
   * - ``kOpaque``
     - 8
     - Opaque: cannot be fused (e.g., external function calls, operations with side effects).

These kinds form an ordering: lower values are "simpler" and more fusable. The fusion algorithm
uses ``CombinePattern(lhs, rhs) = max(lhs, rhs)`` when merging patterns along a path.


FuseOps: Automatic Fusion
-------------------------

``FuseOps`` (``src/relax/transform/fuse_ops.cc``) groups bindings in a dataflow block into
new Relax functions. It operates only within ``DataflowBlock``\ s — if your module doesn't have
any, run ``ConvertToDataflow`` first.

Algorithm
~~~~~~~~~

The fusion algorithm addresses diamond-shaped dataflow branches, where a single producer
(e.g., conv2d) has multiple consumers that eventually reconverge:

.. code-block:: text

            conv2d
            /  |  \
           /   |   \
         op    op   op
          \    |    /
           \   |   /
          elemwise add

At the point of ``conv2d``, we don't know if all future paths will merge. The algorithm uses
**post-dominator analysis** to resolve this:

1. **Build forward graph**: construct an ``IndexedForwardGraph`` from the dataflow block.
   Each node has an ``OpPatternKind`` and a list of forward edges.

2. **Build post-dominator tree**: compute the immediate post-dominator of each node using
   Least Common Ancestor (LCA) on the DAG. The post-dominator of a node is the closest
   downstream node where **all** future paths converge.

3. **Fuse groups**: for each node in topological order, check if it can be fused with its
   immediate post-dominator:

   - **CheckPath**: verify that all paths from the node to its post-dominator satisfy the
     fusion conditions (pattern compatibility, depth limits, argument limits).
   - **CommitFuse**: mark all intermediate nodes as belonging to the same group using a
     Union-Find data structure.

4. **Create grouped functions**: extract each group into a new ``relax.Function`` with the
   attribute ``Primitive=True``. Replace the original bindings with a call to the grouped
   function.

Fusion rules
~~~~~~~~~~~~

The key fusion decisions depend on the ``OpPatternKind`` of the source, the path, and the
post-dominator. The algorithm runs in three phases (via ``GraphPartitioner::RunFuse``) so that
higher-complexity ops get a chance to fuse first:

- **Phase 0**: ``kOutEWiseFusable`` ops (e.g., ``conv2d``) can fuse with their elementwise
  post-dominator if all intermediate ops are broadcast or simpler. This enables patterns like
  conv2d + bias_add + relu. Two ``kOutEWiseFusable`` ops cannot fuse together.
- **Phase 1**: ``kInjective`` and ``kTuple`` ops can fuse only when all paths to the
  post-dominator are injective or simpler. This is deferred to phase 1 so that
  ``kOutEWiseFusable`` groups are finalized first.
- **Phase 2**: fuse injective ops into intermediate tuple nodes that have already been absorbed
  by subsequent injective groups.

``kElemWise`` / ``kBroadcast`` ops are processed in **every** phase (not restricted to one):
they can fuse into a post-dominator that is injective or reduction. The sink (final node) may
also be a ``kOutEWiseFusable`` group that was formed in phase 0 — this is how elementwise
producers merge into an existing conv2d fusion group.

Additional constraints:

- **Reduction** (``kCommReduce``) ops never initiate fusion — they act as sinks only. Elementwise
  and broadcast producers can fuse *into* a reduction, but a reduction cannot fuse forward.
- **Opaque** ops are fusion barriers.
- A group cannot exceed ``kMaxFusedOps`` (256) nodes or the maximum function argument count.

Example
~~~~~~~

Given two elementwise ops (``add``, ``exp``) and one injective op (``squeeze``).
The examples below are simplified pseudocode — real TVMScript would reference TIR functions
via ``cls.func_name``:

.. code-block:: python

   # Before FuseOps (simplified)
   @R.function
   def main(x: R.Tensor((10, 20), "float32")):
       with R.dataflow():
           lv0 = R.call_tir(add, (x, const_1), out_sinfo=R.Tensor((10, 20), "float32"))
           lv1 = R.call_tir(exp, (lv0,), out_sinfo=R.Tensor((10, 20), "float32"))
           gv = R.call_tir(squeeze, (lv1,), out_sinfo=R.Tensor((10, 20), "float32"))
           R.output(gv)
       return gv

After ``FuseOps``, all three are grouped into a single function:

.. code-block:: python

   # After FuseOps
   @R.function(private=True)
   def fused_add_exp_squeeze(x, p0):
       R.func_attr({"Primitive": True})
       with R.dataflow():
           lv0 = R.call_tir(add, (x, p0), ...)
           lv1 = R.call_tir(exp, (lv0,), ...)
           gv = R.call_tir(squeeze, (lv1,), ...)
           R.output(gv)
       return gv

   @R.function
   def main(x: R.Tensor((10, 20), "float32")):
       with R.dataflow():
           gv = fused_add_exp_squeeze(x, const_1)
           R.output(gv)
       return gv


FuseTIR: Merging TIR Functions
------------------------------

``FuseTIR`` (``src/relax/transform/fuse_tir.cc``) takes the grouped Relax functions produced by
``FuseOps`` and merges their internal TIR ``PrimFunc``\ s into a single TIR function.

Before ``FuseTIR``, a fused group still contains multiple ``R.call_tir`` calls to separate
TIR functions. ``FuseTIR`` inlines and merges them:

.. code-block:: text

   Before FuseTIR:
     fused_add_exp_squeeze:
       call_tir(add, ...)        → separate TIR PrimFunc
       call_tir(exp, ...)        → separate TIR PrimFunc
       call_tir(squeeze, ...)    → separate TIR PrimFunc

   After FuseTIR:
     fused_add_exp_squeeze:      → single merged TIR PrimFunc

The merged function eliminates intermediate buffers — the output of ``add`` is directly consumed
by ``exp`` without writing to and reading from global memory. This is the core performance benefit
of fusion.

Internally, ``FuseTIR`` uses a ``SymbolicMatcher`` to align symbolic shape variables across the
TIR functions being merged, ensuring that dimensions are correctly mapped when combining buffer
accesses.


FuseOpsByPattern: Pattern-Based Fusion
--------------------------------------

While ``FuseOps`` makes fusion decisions automatically based on operator patterns,
``FuseOpsByPattern`` lets you specify exactly which operator combinations to fuse using
the Relax :ref:`Dataflow Pattern Language (DPL) <relax-dpl>`.

This is primarily used for **backend-specific dispatch**: identifying operator subgraphs that
should be offloaded to external libraries like cuBLAS, CUTLASS, cuDNN, or DNNL.

FusionPattern
~~~~~~~~~~~~~

A ``FusionPattern`` (``python/tvm/relax/transform/transform.py``) defines what to match:

.. code-block:: python

   from tvm.relax.dpl import wildcard, is_op
   from tvm.relax.transform import FusionPattern

   # Match: matmul(x, w) + bias
   x = wildcard()
   w = wildcard()
   bias = wildcard()
   matmul = is_op("relax.matmul")(x, w)
   out = is_op("relax.add")(matmul, bias)

   pattern = FusionPattern(
       name="cutlass.matmul_bias",
       pattern=out,
       annotation_patterns={"matmul": matmul, "bias": bias},
       check=my_check_function,  # optional validation
   )

Fields:

- ``name``: pattern identifier, typically prefixed with the backend name (e.g.,
  ``"cutlass.matmul_bias"``).
- ``pattern``: a DFPattern describing the subgraph to match. See the
  :ref:`DPL deep dive <relax-dpl>` for the full pattern language.
- ``annotation_patterns``: a mapping of names to sub-patterns within the main pattern. These
  are extracted during matching and made available to the ``check`` function and
  ``attrs_getter``.
- ``check``: an optional ``Callable[[PatternCheckContext], bool]`` that validates whether
  a match should be accepted. Receives the matched expression, annotated sub-expressions,
  variable usages, and binding information.
- ``attrs_getter``: an optional function that extracts attributes (e.g., transpose flags,
  data types) from the matched expressions to annotate the grouped function.

Applying patterns
~~~~~~~~~~~~~~~~~

.. code-block:: python

   from tvm.relax.transform import FuseOpsByPattern

   mod = FuseOpsByPattern(
       patterns=[pattern1, pattern2, ...],  # ordered by priority
       bind_constants=True,
       annotate_codegen=False,
   )(mod)

Key parameters:

- ``patterns``: a list of ``FusionPattern`` objects, ordered by priority. Higher-priority
  patterns come first — if a subgraph matches multiple patterns, the first match wins.
- ``bind_constants``: if ``True``, constants used by the matched subgraph are captured inside
  the grouped function.
- ``annotate_codegen``: if ``True``, wraps each composite function with an outer function
  annotated with ``"Codegen"`` and ``"global_symbol"`` attributes for external backend dispatch.
  The ``"Codegen"`` value is derived from the pattern name prefix (e.g., ``"dnnl"`` from
  ``"dnnl.conv2d_relu"``).

PatternCheckContext
~~~~~~~~~~~~~~~~~~~

The ``check`` function receives a ``PatternCheckContext`` with:

- ``matched_expr``: the root expression matched by the pattern.
- ``annotated_expr``: a mapping from annotation pattern names to their matched expressions.
- ``matched_bindings``: variable-to-value bindings within the matched subgraph.
- ``var_usages``: a mapping from variable definitions to all their uses in the function.
- ``value_to_bound_var``: reverse mapping from values to the variables they are bound to.

This context enables sophisticated validation logic, such as checking that an intermediate
result is not used outside the fused group, or verifying data type compatibility.


How Backends Use Fusion
-----------------------

The default backend pipelines (CUDA, ROCm, CPU, etc.) all include ``FuseOps`` + ``FuseTIR``
in their ``legalize_passes`` phase for automatic fusion, as shown in the `Overview`_ above.

For external library dispatch (cuBLAS, CUTLASS, cuDNN, DNNL), ``FuseOpsByPattern`` is used
separately. These are **not** included in the default pipeline — users add them explicitly
when building a custom compilation flow. The typical sequence is:

1. **Pattern-based dispatch** (``FuseOpsByPattern``): identify subgraphs that should be
   offloaded to external libraries. For example, CUTLASS patterns match
   matmul+bias+activation combinations (``python/tvm/relax/backend/cuda/cutlass.py``).
   Functions marked by patterns are annotated with ``Composite`` and optionally ``Codegen``
   attributes. See :ref:`external-library-dispatch` for the full BYOC pipeline.

2. **Automatic fusion** (``FuseOps`` + ``FuseTIR``): remaining operators that were not
   matched by backend patterns are fused automatically based on their pattern kinds.


Source Code Map
---------------

.. list-table::
   :header-rows: 1
   :widths: 50 50

   * - Path
     - Contents
   * - ``src/relax/transform/fuse_ops.cc``
     - FuseOps and FuseOpsByPattern implementation
   * - ``src/relax/analysis/graph_partitioner.h``
     - IndexedForwardGraph, DominatorTree, GraphPartitioner (Union-Find)
   * - ``src/relax/transform/fuse_tir.cc``
     - FuseTIR implementation, SymbolicMatcher
   * - ``include/tvm/relax/op_attr_types.h``
     - ``OpPatternKind`` enum definition
   * - ``python/tvm/relax/transform/transform.py``
     - Python API: FuseOps, FuseTIR, FuseOpsByPattern, FusionPattern
   * - ``python/tvm/relax/dpl/``
     - Dataflow Pattern Language (DFPattern, is_op, wildcard, etc.)
   * - ``python/tvm/relax/backend/cuda/cutlass.py``
     - Example: CUTLASS fusion patterns
   * - ``python/tvm/relax/backend/cuda/cublas.py``
     - Example: cuBLAS fusion patterns
