[DispatchCreation] Fuse reshape op chains along with set_encoding ops (iree-org#21365)

There can be reshape ops in between set_encoding ops and their producer
dispatch region ops when we fuse encoding ops into dispatches, so this
PR pulls any reshape op chains into the producer dispatch along with the
set_encoding op. This enables more producer fusions for set_encoding in
the data tiling fusion path.

---------

Signed-off-by: Max Dawkins <[email protected]>
Signed-off-by: keshavvinayak01 <[email protected]>

Author: Max191

compiler/src/iree/compiler/DispatchCreation/FuseEncodingOpsIntoDispatchRegions.cpp

@@ -8,6 +8,7 @@
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
+#include "iree/compiler/DispatchCreation/FusionUtils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
@@ -83,16 +84,18 @@ struct FuseEncodingOpsIntoDispatchRegionsPass
 
     for (IREE::Encoding::SetEncodingOp encodingOp : encodingOps) {
       OpOperand &operand = encodingOp.getSourceMutable();
-      auto producerDispatch =
-          operand.get().getDefiningOp<IREE::Flow::DispatchRegionOp>();
+      std::optional<std::pair<OpResult, SmallVector<Operation *>>>
+          producerChain = getProducerDispatchValueAndOpChain(operand.get());
       // Nothing to fuse with, so wrap the `encodingOp` in its own dispatch.
-      if (!producerDispatch) {
+      if (!producerChain) {
         continue;
       }
 
       // Find producer operation inside of the dispatch region to determine if
       // fusion is possible.
-      auto result = cast<OpResult>(operand.get());
+      OpResult result = producerChain->first;
+      auto producerDispatch =
+          result.getDefiningOp<IREE::Flow::DispatchRegionOp>();
       auto dispatchReturnOp = cast<IREE::Flow::ReturnOp>(
           producerDispatch.getBody().front().getTerminator());
       auto producerInRegion = dyn_cast<OpResult>(
@@ -107,10 +110,18 @@ struct FuseEncodingOpsIntoDispatchRegionsPass
           !isFusableWithSetEncoding(producerInRegion.getOwner())) {
         continue;
       }
-      // Fuse the `encodingOp` into the producer dispatch region.
-      if (failed(moveFollowingOpIntoDispatchRegion(rewriter, encodingOp,
-                                                   producerDispatch))) {
-        return signalPassFailure();
+      // Fuse the `encodingOp` and the producer chain into the dispatch.
+      SmallVector<Operation *> dispatchConsumers(
+          llvm::reverse(producerChain->second));
+      dispatchConsumers.push_back(encodingOp);
+      for (Operation *consumer : dispatchConsumers) {
+        FailureOr<IREE::Flow::DispatchRegionOp> fusedDispatch =
+            moveFollowingOpIntoDispatchRegion(rewriter, consumer,
+                                              producerDispatch);
+        if (failed(fusedDispatch)) {
+          return signalPassFailure();
+        }
+        producerDispatch = fusedDispatch.value();
       }
     }
 

compiler/src/iree/compiler/DispatchCreation/FusionUtils.cpp

@@ -9,10 +9,12 @@
 
 #include "compiler/src/iree/compiler/DispatchCreation/FusionUtils.h"
 #include "compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
+#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/RegionUtils.h"
 
 namespace mlir::iree_compiler::DispatchCreation {
@@ -177,4 +179,53 @@ LogicalResult moveOperandDefs(RewriterBase &rewriter,
   return success();
 }
 
+std::optional<std::pair<OpResult, SmallVector<Operation *>>>
+getProducerDispatchValueAndOpChain(Value operand) {
+  auto operandType = dyn_cast<RankedTensorType>(operand.getType());
+  if (!operandType || operandType.getRank() == 0) {
+    return std::nullopt;
+  }
+
+  SmallVector<Operation *> opChain;
+  auto producerValue = dyn_cast<OpResult>(operand);
+  while (producerValue &&
+         !isa<IREE::Flow::DispatchRegionOp>(producerValue.getOwner())) {
+    if (!llvm::hasSingleElement(producerValue.getUses())) {
+      return std::nullopt;
+    }
+
+    // If it is an operation that we want to look past, add it to the chain
+    // and update the `producerValue`.
+    Operation *currOperation = producerValue.getOwner();
+    if (isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp>(currOperation)) {
+      opChain.push_back(currOperation);
+      producerValue = dyn_cast<OpResult>(currOperation->getOperand(0));
+      continue;
+    }
+
+    // Conservative, bail out.
+    return std::nullopt;
+  }
+
+  if (!producerValue) {
+    return std::nullopt;
+  }
+
+  auto producerDispatch =
+      dyn_cast<IREE::Flow::DispatchRegionOp>(producerValue.getOwner());
+  // TODO(MaheshRavishankar): Multi-result producer dispatches can be supported.
+  // Will require to move the consumer dispatch immediately after the producer
+  // instead of what is done below and move other operands of the consumer
+  // dispatch before the producer dispatch.
+  if (!producerDispatch ||
+      !llvm::hasSingleElement(producerDispatch.getBody()) ||
+      producerDispatch->getNumResults() != 1) {
+    return std::nullopt;
+  }
+  if (!llvm::hasSingleElement(producerValue.getUses())) {
+    return std::nullopt;
+  }
+  return std::make_pair(producerValue, opChain);
+}
+
 } // namespace mlir::iree_compiler::DispatchCreation

compiler/src/iree/compiler/DispatchCreation/FusionUtils.h

@@ -34,4 +34,11 @@ LogicalResult moveOperandDefs(RewriterBase &rewriter,
                               DominanceInfo &dominanceInfo,
                               ArrayRef<Operation *> ignoreOperations = {});
 
+/// Returns the closest producer dispatch region op result and the chain of
+/// operations being looked past during the traversal to find the producer
+/// dispatch. Returns std::nullopt if the dispatch or any ops in the chain have
+/// multiple uses.
+std::optional<std::pair<OpResult, SmallVector<Operation *>>>
+getProducerDispatchValueAndOpChain(Value operand);
+
 } // namespace mlir::iree_compiler::DispatchCreation

compiler/src/iree/compiler/DispatchCreation/SetEncoding.cpp

@@ -16,6 +16,7 @@
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/DispatchCreation/FusionUtils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
@@ -310,59 +311,6 @@ struct FoldFillWithSetEncoding final
 //===---------------------------------------------------------------------===//
 // Set padding encodings
 //===---------------------------------------------------------------------===//
-
-// Utility to return the producer dispatch region op result and the chain of
-// operations being looked past during the traversal to find the producer
-// dispatch.
-static std::optional<std::pair<OpResult, SmallVector<Operation *>>>
-getProducerDispatchValueAndOpChain(Value operand) {
-  auto operandType = dyn_cast<RankedTensorType>(operand.getType());
-  if (!operandType || operandType.getRank() == 0) {
-    return std::nullopt;
-  }
-
-  SmallVector<Operation *> opChain;
-  auto producerValue = dyn_cast<OpResult>(operand);
-  while (producerValue &&
-         !isa<IREE::Flow::DispatchRegionOp>(producerValue.getOwner())) {
-    if (!llvm::hasSingleElement(producerValue.getUses())) {
-      return std::nullopt;
-    }
-
-    // If it is an operation that we want to look past, add it to the chain
-    // and update the `producerValue`.
-    Operation *currOperation = producerValue.getOwner();
-    if (isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp>(currOperation)) {
-      opChain.push_back(currOperation);
-      producerValue = dyn_cast<OpResult>(currOperation->getOperand(0));
-      continue;
-    }
-
-    // Conservative, bail out.
-    return std::nullopt;
-  }
-
-  if (!producerValue) {
-    return std::nullopt;
-  }
-
-  auto producerDispatch =
-      dyn_cast<IREE::Flow::DispatchRegionOp>(producerValue.getOwner());
-  // TODO(MaheshRavishankar): Multi-result producer dispatches can be supported.
-  // Will require to move the consumer dispatch immediately after the producer
-  // instead of what is done below and move other operands of the consumer
-  // dispatch before the producer dispatch.
-  if (!producerDispatch ||
-      !llvm::hasSingleElement(producerDispatch.getBody()) ||
-      producerDispatch->getNumResults() != 1) {
-    return std::nullopt;
-  }
-  if (!llvm::hasSingleElement(producerValue.getUses())) {
-    return std::nullopt;
-  }
-  return std::make_pair(producerValue, opChain);
-}
-
 struct PaddedValue {
   Value paddedValue;
   SmallVector<Value> dynamicDims;

compiler/src/iree/compiler/DispatchCreation/test/fuse_encoding_ops_into_dispatch_regions.mlir

@@ -187,3 +187,30 @@ util.func public @multi_encoding_fusion_dynamic(%arg0: tensor<?x?x?xf32>, %d0: i
 // CHECK:         flow.return %[[SET_ENCODING]] :
 // CHECK:       }
 // CHECK:       util.return %[[DISPATCH]], %[[DISPATCH]]
+
+// -----
+
+#encoding = #iree_encoding.testing<>
+util.func public @reshape_fusion(%arg0: tensor<32x32xf32>) -> tensor<16x64xf32, #encoding> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<32x32xf32>
+  %1 = flow.dispatch.region -> (tensor<32x32xf32>) {
+    %3 = linalg.add ins(%arg0, %arg0 : tensor<32x32xf32>, tensor<32x32xf32>)
+        outs(%0 : tensor<32x32xf32>) -> tensor<32x32xf32>
+    flow.return %3 : tensor<32x32xf32>
+  }
+  %collapsed = tensor.collapse_shape %1 [[0, 1]] : tensor<32x32xf32> into tensor<1024xf32>
+  %expanded = tensor.expand_shape %collapsed [[0, 1]] output_shape [16, 64] : tensor<1024xf32> into tensor<16x64xf32>
+  %2 = iree_encoding.set_encoding %expanded : tensor<16x64xf32> -> tensor<16x64xf32, #encoding>
+  util.return %2 : tensor<16x64xf32, #encoding>
+}
+// CHECK:       #[[$ENCODING:.+]] = #iree_encoding.testing<>
+// CHECK-LABEL: @reshape_fusion
+// CHECK:       %[[DISPATCH0:.+]] = flow.dispatch.region -> (tensor<16x64xf32, #[[$ENCODING]]>)
+// CHECK:         linalg.add
+// CHECK:         tensor.collapse_shape
+// CHECK:         tensor.expand_shape
+// CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding
+// CHECK:         flow.return %[[SET_ENCODING]] :
+// CHECK:       }
+// CHECK:       util.return %[[DISPATCH0]] : tensor<16x64xf32, #[[$ENCODING]]>