[mlir][Affine] Add support for multi-store producer fusion
authorDiego Caballero <diego.caballero@intel.com>
Wed, 20 Jan 2021 01:24:57 +0000 (03:24 +0200)
committerDiego Caballero <diego.caballero@intel.com>
Wed, 20 Jan 2021 17:03:07 +0000 (19:03 +0200)
commit7dd198852b4db52ae22242dfeda4eccda83aa8b2
treec936a7c378b498f3335276da3a97ad4e1fe1f3c8
parentfd70f70d1e02752f411fcf923fddda31cce376ae
[mlir][Affine] Add support for multi-store producer fusion

This patch adds support for producer-consumer fusion scenarios with
multiple producer stores to the AffineLoopFusion pass. The patch
introduces some changes to the producer-consumer algorithm, including:

* For a given consumer loop, producer-consumer fusion iterates over its
producer candidates until a fixed point is reached.

* Producer candidates are gathered beforehand for each iteration of the
consumer loop and visited in reverse program order (not strictly guaranteed)
to maximize the number of loops fused per iteration.

In general, these changes were needed to simplify the multi-store producer
support and remove some of the workarounds that were introduced in the past
to support more fusion cases under the single-store producer limitation.

This patch also preserves the existing functionality of AffineLoopFusion with
one minor change in behavior. Producer-consumer fusion didn't fuse scenarios
with escaping memrefs and multiple outgoing edges (from a single store).
Multi-store producer scenarios will usually (always?) have multiple outgoing
edges so we couldn't fuse any with escaping memrefs, which would greatly limit
the applicability of this new feature. Therefore, the patch enables fusion for
these scenarios. Please, see modified tests for specific details.

Reviewed By: andydavis1, bondhugula

Differential Revision: https://reviews.llvm.org/D92876
mlir/include/mlir/Analysis/AffineStructures.h
mlir/include/mlir/Analysis/Utils.h
mlir/include/mlir/Transforms/LoopFusionUtils.h
mlir/include/mlir/Transforms/Passes.td
mlir/lib/Analysis/AffineStructures.cpp
mlir/lib/Analysis/Utils.cpp
mlir/lib/Transforms/LoopFusion.cpp
mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
mlir/test/Transforms/loop-fusion.mlir