[DT] Fuse encoding ops more aggressively for multi-use, gather, and slices ops. (iree-org#21830)

The fusion constraint of multi-use dispatch is only required by
SetEncoding pass, because it has to move consumer dispatches around. It
is not required by encoding fusion, because it is just moving a
SetEncoding op into its producer dispatch.

The revision also allows the fusion when the dispatch region contains
tensor.extract_slice op and iree_linalg_ext.gather ops. It reduces the
number of dispatches to 644 in llama fp8 model, the same as without data
tiling. The latency drops 25ms, from 378ms to 353ms.

| | No Data Tiling | Data Tiling w/o the revision | Data Tiling w/ the
revision |
| ------------- | ------------- | ------------- | ------------- |
| Benchmark latency | 243ms  |  378ms | 353ms |
| Memory usage (HIP unpooled)  | 15.9GB  |  31.14GB | 31.11GB |
| Number of dispatches | 644  | 741 | 644 |

| | No Data Tiling (ms) | Data Tiling w/o the revision | Data Tiling w/
the revision |
| ------------- | ------------- | ------------- | ------------- |
| dispatch_15_attention_4x8x4xDx128xf8  | 62.29 |  55.35 | 59.21 |
| dispatch_20_matmul_like_Dx14336x4096_f8xf8xf32 | 40.13 | 89.14 |
93.72|
| dispatch_19_matmul_like_Dx14336x4096_f8xf8xf32 | 28.01 | 44.78 | 44.59
|
| dispatch_21_matmul_like_Dx4096x14336_f8xf8xf32 | 27.25 | 40.18 | 39.99
|
| dispatch_643_matmul_like_Dx128256x4096_f16xf16xf32 | 17.1 | 29.76 |
29.21 |
| dispatch_16_matmul_like_Dx4096x4096_f8xf8xf32 | 8.83 | 17.92 | 17.91 |
| dispatch_23_matmul_like_Dx4096x4096_f8xf8xf32 | 9.27 | 16.69 | 16.59 |
| encoding_10_encode_Dx4096xf8_to_Dx4096xf8 | - |  32.15 | - |
| encoding_6_encode_Dx14336xf32_to_Dx14336xf32 | - |  0.318 | - |

---------

Signed-off-by: hanhanW <hanhan0912@gmail.com>

Author: hanhanW

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

@@ -31,12 +31,13 @@ namespace mlir::iree_compiler::DispatchCreation {
 namespace {
 
 /// Return true if the op is fusable with a SetEncodingOp consumer. The op's
-/// containing dispatch must contain only reshape ops, encoding ops, linalg ops,
-/// and attention ops. Non ShapedType ops (like arith ops, dim ops, etc.) are
-/// also allowed.
-/// TODO(#20179): It should be done by interface methods.
-static bool isFusableWithSetEncoding(Operation *op) {
-  auto parentRegion = op->getParentOfType<IREE::Flow::DispatchRegionOp>();
+/// containing dispatch must contain only:
+///   - Reshape ops, encoding ops, linalg ops, gather ops, and attention ops.
+///   - Non ShapedType ops, e.g., like arith ops, dim ops, etc.
+///   - tensor::ExtractSliceOp is allowed as they can be folded into dispatch
+///     tensor load ops.
+static bool isFusableWithSetEncoding(Operation *target) {
+  auto parentRegion = target->getParentOfType<IREE::Flow::DispatchRegionOp>();
   // Make sure the dispatch region has only one block.
   if (!llvm::hasSingleElement(parentRegion.getBody())) {
     return false;
@@ -49,8 +50,9 @@ static bool isFusableWithSetEncoding(Operation *op) {
       continue;
     }
     if (!isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp, tensor::EmptyOp,
-             IREE::Encoding::SetEncodingOp, IREE::Encoding::UnsetEncodingOp,
-             linalg::LinalgOp, IREE::LinalgExt::AttentionOp>(op)) {
+             tensor::ExtractSliceOp, IREE::Encoding::SetEncodingOp,
+             IREE::Encoding::UnsetEncodingOp, linalg::LinalgOp,
+             IREE::LinalgExt::AttentionOp, IREE::LinalgExt::GatherOp>(op)) {
       return false;
     }
   }

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

@@ -145,13 +145,8 @@ getProducerDispatchValueAndOpChain(Value operand) {
 
   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) {
+      !llvm::hasSingleElement(producerDispatch.getBody())) {
     return std::nullopt;
   }
   if (!llvm::hasSingleElement(producerValue.getUses())) {

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

@@ -28,8 +28,8 @@ bool areFusableAsElementwiseOps(MLIRContext *context, OpOperand *operand,
 
 /// 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.
+/// dispatch. Returns std::nullopt if the dispatch can not be found in the
+/// chain or any op in the chain is not a reshape-like op.
 std::optional<std::pair<OpResult, SmallVector<Operation *>>>
 getProducerDispatchValueAndOpChain(Value operand);
 

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

@@ -392,16 +392,24 @@ static SmallVector<unsigned> getOperandsToPad(Operation *op) {
     OpOperand &operand = op->getOpOperand(operandNum);
     std::optional<std::pair<OpResult, SmallVector<Operation *>>>
         dispatchAndOpChain = getProducerDispatchValueAndOpChain(operand.get());
-    if (dispatchAndOpChain.has_value()) {
-      auto producerDispatch = cast<IREE::Flow::DispatchRegionOp>(
-          dispatchAndOpChain->first.getOwner());
-      WalkResult res =
-          producerDispatch->walk([&](IREE::LinalgExt::AttentionOp op) {
-            return WalkResult::interrupt();
-          });
-      if (res.wasInterrupted()) {
-        return {};
-      }
+    if (!dispatchAndOpChain.has_value()) {
+      continue;
+    }
+    auto producerDispatch = cast<IREE::Flow::DispatchRegionOp>(
+        dispatchAndOpChain->first.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->getNumResults() != 1) {
+      continue;
+    }
+    WalkResult res =
+        producerDispatch->walk([&](IREE::LinalgExt::AttentionOp op) {
+          return WalkResult::interrupt();
+        });
+    if (res.wasInterrupted()) {
+      return {};
     }
   }
 

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

@@ -280,6 +280,24 @@ util.func public @encoding_fusion(%arg0: tensor<128xf32, #encoding0>) -> tensor<
 
 // -----
 
+#encoding = #iree_encoding.testing<>
+util.func public @extract_slice_fusion(%arg0: tensor<192x1024x64xf32>) -> tensor<96x512x32xf32, #encoding> {
+  %0 = flow.dispatch.region -> (tensor<96x512x32xf32>) {
+    %1 = tensor.extract_slice %arg0[0, 0, 0] [96, 512, 32] [1, 1, 1] : tensor<192x1024x64xf32> to tensor<96x512x32xf32>
+    flow.return %1 : tensor<96x512x32xf32>
+  }
+  %2 = iree_encoding.set_encoding %0 : tensor<96x512x32xf32> -> tensor<96x512x32xf32, #encoding>
+  util.return %2 : tensor<96x512x32xf32, #encoding>
+}
+// CHECK-LABEL: @extract_slice_fusion
+// CHECK:       %[[DISPATCH0:.+]] = flow.dispatch.region
+// CHECK:         tensor.extract_slice
+// CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding
+// CHECK:         flow.return %[[SET_ENCODING]] :
+// CHECK:       util.return %[[DISPATCH0]]
+
+// -----
+
 #encoding = #iree_encoding.testing<>
 util.func public @attention_fusion(
     %query: tensor<192x1024x64xf32>, %key: tensor<192x1024x64xf32>,
@@ -309,3 +327,55 @@ util.func public @attention_fusion(
 // CHECK:         flow.return %[[SET_ENCODING]] :
 // CHECK:       }
 // CHECK:       util.return %[[DISPATCH0]]
+
+// -----
+
+#encoding = #iree_encoding.testing<>
+util.func public @gather_fusion(
+    %source: tensor<10x10xi32>, %indices: tensor<1xi32>) -> tensor<1x10xi32, #encoding> {
+  %empty = tensor.empty() : tensor<1x10xi32>
+  %1 = flow.dispatch.region -> (tensor<1x10xi32>) {
+    %3 = iree_linalg_ext.gather dimension_map = [0]
+           ins(%source, %indices : tensor<10x10xi32>, tensor<1xi32>)
+           outs(%empty : tensor<1x10xi32>) -> tensor<1x10xi32>
+    flow.return %3 : tensor<1x10xi32>
+  }
+  %2 = iree_encoding.set_encoding %1 : tensor<1x10xi32> -> tensor<1x10xi32, #encoding>
+  util.return %2 : tensor<1x10xi32, #encoding>
+}
+// CHECK-LABEL: @gather_fusion
+// CHECK:       %[[DISPATCH0:.+]] = flow.dispatch.region
+// CHECK:         iree_linalg_ext.gather
+// CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding
+// CHECK:         flow.return %[[SET_ENCODING]] :
+// CHECK:       util.return %[[DISPATCH0]]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#encoding = #iree_encoding.testing<>
+util.func public @multi_result_fusion(%arg0: tensor<123x456xf32>) -> (tensor<123x456xf32>, tensor<123x456xf32, #encoding>) {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<123x456xf32>
+  %1:2 = flow.dispatch.region -> (tensor<123x456xf32>, tensor<123x456xf32>) {
+    %3:2 = linalg.generic {
+        indexing_maps = [#map, #map, #map, #map],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%arg0, %arg0 : tensor<123x456xf32>, tensor<123x456xf32>)
+        outs(%0, %0 : tensor<123x456xf32>, tensor<123x456xf32>) {
+    ^bb0(%in: f32, %in_0: f32, %out: f32, %out2: f32):
+      %4 = arith.addf %in, %in_0 : f32
+      %5 = arith.mulf %in, %in_0 : f32
+      linalg.yield %4, %5 : f32, f32
+    } -> (tensor<123x456xf32>, tensor<123x456xf32>)
+    flow.return %3#0, %3#1 : tensor<123x456xf32>, tensor<123x456xf32>
+  }
+  %2 = iree_encoding.set_encoding %1#1 : tensor<123x456xf32> -> tensor<123x456xf32, #encoding>
+  util.return %1#0, %2 : tensor<123x456xf32>, tensor<123x456xf32, #encoding>
+}
+// CHECK-LABEL: @multi_result_fusion
+// CHECK:       %[[DISPATCH0:.+]]:2 = flow.dispatch.region
+// CHECK:         %[[ELEM:.+]]:2 = linalg.generic
+// CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[ELEM]]#1
+// CHECK:         flow.return %[[ELEM]]#0, %[[SET_ENCODING]]
+// CHECK:       util.return %[[DISPATCH0]]#0, %[[DISPATCH0]]#1