[DispatchCreation][DT] Only fuse encodings with element-wise (#20143)

This PR restricts the fusion of set_encoding ops to only dispatches with
element-wise linalg ops. We don't have good codegen for fusions with
reduction ops on GPU yet, so this will allow using the late
materialization data tiling path on GPU while the codegen improves.

---------

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

Author: Max191

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

@@ -11,6 +11,7 @@
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
@@ -25,10 +26,37 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 namespace {
 
-// Return true if the op is fusable with a SetEncodingOp consumer.
-// For now, just check if it is a LinalgOp.
+// Return true if the op is fusable with a SetEncodingOp consumer. For now,
+// the op's containing dispatch region must not contain any ops other than
+// element-wise linalg ops and some tensor ops. This is quite conservative,
+// and could be extended to more ops when we are confident that the codegen
+// backends can support it.
 static bool isFusableWithSetEncoding(Operation *op) {
-  return isa<linalg::LinalgOp>(op);
+  auto parentRegion = op->getParentOfType<IREE::Flow::DispatchRegionOp>();
+  // Make sure the dispatch region has only one block.
+  if (!llvm::hasSingleElement(parentRegion.getBody())) {
+    return false;
+  }
+  // Check that there are no ops other than reshapes and element-wise linalg
+  // ops in the dispatch region.
+  Block &regionBlock = parentRegion.getBody().getBlocks().front();
+  for (Operation &op : regionBlock.getOperations()) {
+    if (llvm::none_of(op.getResultTypes(), llvm::IsaPred<ShapedType>)) {
+      continue;
+    }
+    if (isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp, tensor::EmptyOp>(
+            op)) {
+      continue;
+    }
+    auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+    if (!linalgOp) {
+      return false;
+    }
+    if (linalgOp.getNumReductionLoops() != 0) {
+      return false;
+    }
+  }
+  return true;
 }
 
 struct FuseEncodingOpsIntoDispatchRegionsPass

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

@@ -69,12 +69,15 @@ module {
 // CHECK-DAG:   #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 // CHECK-DAG:   #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type =  matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
 // CHECK-LABEL: @reduction_fusion
-// CHECK:       %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #[[$ENCODING]]
+// CHECK:       %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32>)
 // CHECK:         %[[REDUCTION:.+]] = linalg.generic
-// CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[REDUCTION]]
+// CHECK:         flow.return %[[REDUCTION]] :
+// CHECK:       }
+// CHECK:       %[[DISPATCH_SE:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #[[$ENCODING]]>)
+// CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[DISPATCH]]
 // CHECK:         flow.return %[[SET_ENCODING]] :
 // CHECK:       }
-// CHECK:       util.return %[[DISPATCH]] : tensor<2x11008x128xf32, #[[$ENCODING]]>
+// CHECK:       util.return %[[DISPATCH_SE]] : tensor<2x11008x128xf32, #[[$ENCODING]]>
 
 // -----