[feat] Kernel-level fusion

[HobbitQia]

FuseKernelPass is an MLIR optimization pass that merges neura.kernel operations.

• Step 1: Pattern Detection
  • Producer-Consumer: kernel_A → kernel_B (data dependency)
  • Sibling: kernel_A || kernel_B (shared inputs, no dependency)
• Step 2: Profitability Analysis
  • Creates temporary test modules with cloned kernel bodies. Then runs full Neura transformation pipeline (assign-accelerator → lower-arith → canonicalize → transform-ctrl-to-dataflow)
  • Computes real metrics: RecMII, ResMII, max fanout, operation count
  • Estimates final MII using formula: MII_est = ⌈(1 + α×ops/tiles) × (1 + β×max(fanout-4, 0)) × max(RecMII, ResMII)⌉. Compares fused vs. separate execution: fusion only proceeds if MII_fused ≤ max(MII_k1, MII_k2). Note that this equation is designed to take the resources and congestion (fanout) into consideration and stay tuned.
• Step 3: Kernel Fusion
  • Producer-Consumer: Merges producer body into consumer, eliminates intermediate buffer
  • Sibling: Combines both kernel bodies, deduplicates shared inputs
  • Iteratively applies fusion until no more profitable opportunities exist. Here we prioritize Producer-Consumer fusion.

An example is shown below, where the two loops are wrapped in the same kernel:

// Before fusion
func.func @test_producer_consumer_fusion(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: memref<?xf32>) {
  %cst = arith.constant 2.000000e+00 : f32
  affine.for %arg4 = 0 to 64 {
    %0 = memref.load %arg0[%arg4] : memref<?xf32>
    %1 = memref.load %arg1[%arg4] : memref<?xf32>
    %2 = arith.addf %0, %1 : f32
    memref.store %2, %arg2[%arg4] : memref<?xf32>
  }
  affine.for %arg4 = 0 to 64 {
    %0 = memref.load %arg2[%arg4] : memref<?xf32>
    %1 = arith.mulf %0, %cst : f32
    memref.store %1, %arg3[%arg4] : memref<?xf32>
  }
  return
}
// After fusion
func.func @test_producer_consumer_fusion(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: memref<?xf32>) {
    %cst = arith.constant 2.000000e+00 : f32
    neura.kernel ins(%arg0, %arg1, %arg2, %cst, %arg3 : memref<?xf32>, memref<?xf32>, memref<?xf32>, f32, memref<?xf32>) attributes {kernel_name = "fused_sibling"} {
    ^bb0(%arg4: memref<?xf32>, %arg5: memref<?xf32>, %arg6: memref<?xf32>, %arg7: f32, %arg8: memref<?xf32>):
      affine.for %arg9 = 0 to 64 {
        %0 = memref.load %arg0[%arg9] : memref<?xf32>
        %1 = memref.load %arg1[%arg9] : memref<?xf32>
        %2 = arith.addf %0, %1 : f32
        memref.store %2, %arg2[%arg9] : memref<?xf32>
      }
      affine.for %arg9 = 0 to 64 {
        %0 = memref.load %arg2[%arg9] : memref<?xf32>
        %1 = arith.mulf %0, %cst : f32
        memref.store %1, %arg3[%arg9] : memref<?xf32>
      }
    }
    return
  }

[tancheng]

How do we use this file for testing?

[HobbitQia]

This file can be used for testing with Polygeist as front-end, which can lower C++ to IRs in affine. Here the IRs after lowering are provided in test.mlir and kernel.cpp is only for reference.

[tancheng]

Then can we really leverage Polygeist to compile it, instead of just a reference?

[tancheng]

Aren't the 3 *_PIPE_OCCUPY overlapping with each other?

[HobbitQia]

Actually 3 means the multi-cycle op will not occupies the input and output ports of the tile so we can map other operations onto this tile, which is inclusive execution we proposed before in our DATE paper.

However, I have not finished the implementation and test of inclusive execution so far. Here I just copied some content from CGRA-Mapper. Will tune it in the future.

[tancheng]

So IN_PIPE_OCCUPY does not include start and end, right?