Fusion Analysis Interface for Compute Operations

MLIR
,
1

As part of our downstream MLIR compiler pipeline, we are utilizing the tile and fuse transform, specifically the greedy utility function tileConsumerAndFuseProducersUsingSCF. This transformation has been particularly helpful in applying tiling and controlling the fusion of various compute operations that we aim to combine within a single scf.for loop structure.

However, we are currently missing some essential analysis information that would allow us to plan the tile-and-fuse phase more effectively and potentially build profitability heuristics around it. To address this, we are considering a pre-analysis step that could inform our tiling decisions.

Problem

Once we select a tile size for a root operation (typically a consumer), the tile sizes of all upstream producers are expected to be inferred accordingly. This inference effectively determines the resulting loop structure. Therefore, having a good understanding of how tile sizes propagate through the computation chain is crucial.

Initially, we attempted to use the TilingInterface methods:

  • getIterationDomainTileFromOperandTile
  • getIterationDomainTileFromResultTile

for the purpose of this analysis. However, we observed that these methods:

  • Depend heavily on implementation details of the tiling interface.
  • Often modify or clone IR, introducing inaccuracies in our analysis and leading to potentially misleading conclusions.

For purely analytical purposes, we need a way to query the relationship between the iteration domain and the shapes of operands/results — without altering the IR.

Proposed Need

We are looking for a mapping that connects the iteration space dimensions of an operation with the corresponding dimensions of each operand/result. Such a mapping would enable:

  • A static analysis of how a given tile size in the root operation’s iteration space affects the shapes of its operands and results.
  • A predictive model to estimate the impact of tiling on all operations in the fusion group.
  • Avoidance of side effects caused by actual IR transformation during the analysis phase.

This data would be extremely helpful in designing profitable fusion heuristics, enabling us to control fusion not only based on structure, but also based on cost models and performance estimates.

Suggested Fusion Interface

#ifndef FUSION_INTERFACE_DEFS
#define FUSION_INTERFACE_DEFS

include "mlir/IR/OpBase.td"

def FusionInterface : OpInterface<"FusionOpInterface"> {
  let description = [{
    The `FusionInterface` facilitates fusion analysis by providing methods
    to retrieve affine indexing maps for operands and results, along with
    iteration space properties.

    Each affine map specifies, for a given operand or result, how its
    indices are derived from the operation’s iteration domain.

    For example, if a `test.foo` operation implements the following
    indexing maps:

      #first_operand_map  = affine_map<(m, n, k) -> (m, k)>
      #second_operand_map = affine_map<(m, n, k) -> (k, n)>
      #first_result_map   = affine_map<(m, n, k) -> (m, n)>

    This indicates that:
      - the first operand is accessed using (m, k),
      - the second operand is accessed using (k, n),
      - the result is indexed by (m, n).

    These maps capture how loop indices determine operand and result
    accesses, and are essential for accurate tiling, fusion analysis,
    and profitability estimation without altering the IR.
  }];
  let cppNamespace = "::mlir";

  let methods = [
    //===------------------------------------------------------------------===//
    // Interface methods for fusion analysis.
    //===------------------------------------------------------------------===//

    InterfaceMethod<
      /*desc=*/[{
        Populates a vector of `AffineMap`s corresponding to the indexing maps
        of each operand. If an operand does not have an indexing map representation,
        a `nullptr` (or an empty AffineMap) should be inserted in its place.
      }],
      /*retTy=*/"void",
      /*methodName=*/"getIndexingMapsForOperands",
      /*args=*/(ins "::mlir::SmallVectorImpl<::mlir::AffineMap>&" : $indexingMapsRes),
      /*methodBody=*/""
    >,

    InterfaceMethod<
      /*desc=*/[{
        Populates a vector of `AffineMap`s corresponding to the indexing maps
        of each result. If a result does not have an indexing map representation,
        a `nullptr` (or an empty AffineMap) should be inserted in its place.
      }],
      /*retTy=*/"void",
      /*methodName=*/"getIndexingMapsForResults",
      /*args=*/(ins "::mlir::SmallVectorImpl<::mlir::AffineMap>&" : $indexingMapsRes),
      /*methodBody=*/""
    >
  ];
}

#endif // FUSION_INTERFACE_DEFS

A similar interface for the mentioned above is also being used within the downstream IREE project (you may refer to iree/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.td at 0781072a70e3aafbcbde6c22ca687e4d75b084fd · iree-org/iree

1 Like