[DispatchCreation] hoist entire chains of set_encoding ops (#24904)
Capture the entire producer chains of `set_encoding` ops, and hoist them
if there is no unhoistable op in the source backward slice. The metadata
backward slice, i.e. the backward slice from the `encoding_dims` can
always be hoisted.
Fixes #24876.
Assisted-by: Claude Code
---------
Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
index 9c3b674..81312b9 100644
--- a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
@@ -5,7 +5,6 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include <cassert>
-#include <queue>
#include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
@@ -22,6 +21,7 @@
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
+#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
@@ -329,36 +329,38 @@
return;
}
+ // Every op within the dispatch that the SetEncodingOp transitively depends
+ // on (data and metadata) has to be hoisted along with it. While computing
+ // the slice, record the ops on the data path, i.e. the transitive defs of
+ // the source. The source operand is traversed before the encoding dims.
+ Operation *sourceOp = setEncodingOp.getSource().getDefiningOp();
+ DenseSet<Operation *> dataPathOps;
+ BackwardSliceOptions sliceOptions;
bool isHoistable = true;
- SetVector<Operation *> opsWithinDispatch, seen;
- std::queue<Operation *> worklist;
- worklist.push(setEncodingOp);
- seen.insert(setEncodingOp);
- while (!worklist.empty()) {
- Operation *op = worklist.front();
- worklist.pop();
- opsWithinDispatch.insert(op);
- // For SetEncodingOp, the encoding_dims operands are metadata and should
- // not block hoistability. We still add their defining ops to the worklist
- // so they get hoisted, but only the source operand affects hoistability.
- auto setEnc = dyn_cast<IREE::Encoding::SetEncodingOp>(op);
- for (Value input : op->getOperands()) {
- auto inputOp = input.getDefiningOp();
- if (inputOp &&
- inputOp->getParentOfType<IREE::Flow::DispatchRegionOp>() &&
- !seen.contains(inputOp)) {
- // For SetEncodingOp, only the source operand affects hoistability.
- bool affectsHoistability = !setEnc || input == setEnc.getSource();
- if (affectsHoistability && !isHoistableOp(inputOp)) {
- isHoistable = false;
- }
- worklist.push(inputOp);
- seen.insert(inputOp);
- }
+ sliceOptions.inclusive = true;
+ sliceOptions.omitBlockArguments = true;
+ sliceOptions.filter = [&dataPathOps, &isHoistable,
+ &sourceOp](Operation *op) {
+ if (!op->getParentOfType<IREE::Flow::DispatchRegionOp>()) {
+ return false;
}
- }
+ if (op == sourceOp || llvm::any_of(op->getUsers(), [&](Operation *user) {
+ return dataPathOps.contains(user);
+ })) {
+ dataPathOps.insert(op);
+ // Only the ops on the data path gate hoistability. Ops that are only
+ // reachable through `encoding_dims` are metadata and come along for
+ // free.
+ isHoistable &= isHoistableOp(op);
+ }
+ return true;
+ };
+ SetVector<Operation *> opsWithinDispatch;
+ (void)getBackwardSlice(setEncodingOp.getOperation(), &opsWithinDispatch,
+ sliceOptions);
+
if (isHoistable) {
- candidates.push_back(llvm::to_vector(llvm::reverse(opsWithinDispatch)));
+ candidates.push_back(llvm::to_vector(opsWithinDispatch));
return;
}
@@ -371,7 +373,7 @@
}
if (policy.getDecision(constInfo)->getOutcome() ==
IREE::Util::ConstExprHoistingPolicy::ENABLE_HOIST) {
- candidates.push_back(llvm::to_vector(llvm::reverse(opsWithinDispatch)));
+ candidates.push_back(llvm::to_vector(opsWithinDispatch));
return;
}
LDBG() << "Non-hoistable op: " << setEncodingOp;
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
index 88d2381..e205928 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(iree-dispatch-creation-hoist-encoding-ops,cse)" --split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(iree-dispatch-creation-hoist-encoding-ops,cse)" --split-input-file --allow-unregistered-dialect %s | FileCheck %s
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
@@ -785,7 +785,7 @@
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK-DAG: %[[C11008:.+]] = arith.constant 11008 : index
-// CHECK: %[[DIM:.+]] = tensor.dim %{{.*}}, %[[C0]]
+// CHECK-DAG: %[[DIM:.+]] = tensor.dim %{{.*}}, %[[C0]]
// CHECK-DAG: iree_encoding.set_encoding %{{.*}} encoding_dims{%[[DIM]], %[[C11008]], %[[C128]]}
// CHECK-DAG: iree_encoding.set_encoding %{{.*}} encoding_dims{%[[DIM]], %[[C11008]], %[[C128]]}
// CHECK-DAG: iree_encoding.set_encoding %{{.*}} encoding_dims{%[[DIM]], %[[C11008]], %[[C128]]}
@@ -1115,3 +1115,81 @@
// CHECK: flow.return %[[DEQUANT]]
// CHECK: }
// CHECK: util.return %[[DISPATCH]]
+
+// -----
+
+// `encoding_dims` are produced inside the dispatch by tensor.dim / affine.apply
+// of values defined outside the dispatch. These ops only feed metadata, so they
+// must not block hoisting of the set_encoding ops.
+
+#map_md = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_md1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_md2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map_md3 = affine_map<()[s0] -> (s0 * 8)>
+#enc_md_lhs = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_md, #map_md1, #map_md2], iteration_sizes = [?, 256, ?]>
+#enc_md_rhs = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_md, #map_md1, #map_md2], iteration_sizes = [?, 256, ?]>
+#enc_md_out = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_md, #map_md1, #map_md2], iteration_sizes = [?, 256, ?]>
+util.func public @hoist_matmul_with_metadata_encoding_dims(
+ %reshape: tensor<8x?x?xf32>, %collapsed: tensor<?x?xf32>,
+ %rhs: tensor<?x256xf32>, %out_sz: index) -> tensor<?x256xf32> {
+ %0 = flow.dispatch.region -> (tensor<?x256xf32>{%out_sz}) {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %dim1 = tensor.dim %reshape, %c1 : tensor<8x?x?xf32>
+ %m = affine.apply #map_md3()[%dim1]
+ %k = tensor.dim %reshape, %c2 : tensor<8x?x?xf32>
+ %lhs_enc = iree_encoding.set_encoding %collapsed encoding_dims{%m, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #enc_md_lhs>
+ %rhs_enc = iree_encoding.set_encoding %rhs encoding_dims{%m, %k} : tensor<?x256xf32> -> tensor<?x256xf32, #enc_md_rhs>
+ %init = tensor.empty(%out_sz) : tensor<?x256xf32, #enc_md_out>
+ %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x256xf32, #enc_md_out>) -> tensor<?x256xf32, #enc_md_out>
+ %mm = linalg.matmul ins(%lhs_enc, %rhs_enc : tensor<?x?xf32, #enc_md_lhs>, tensor<?x256xf32, #enc_md_rhs>) outs(%fill : tensor<?x256xf32, #enc_md_out>) -> tensor<?x256xf32, #enc_md_out>
+ %unset = iree_encoding.unset_encoding %mm encoding_dims{%m, %k} : tensor<?x256xf32, #enc_md_out> -> tensor<?x256xf32>{%out_sz}
+ flow.return %unset : tensor<?x256xf32>
+ }
+ util.return %0 : tensor<?x256xf32>
+}
+// CHECK-LABEL: @hoist_matmul_with_metadata_encoding_dims
+// CHECK-SAME: %[[RESHAPE:.+]]: tensor<8x?x?xf32>, %[[COLLAPSED:.+]]: tensor<?x?xf32>, %[[RHS:.+]]: tensor<?x256xf32>
+// The tensor.dim / affine.apply that produce the encoding_dims hoist out
+// alongside the set_encoding ops.
+// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[RESHAPE]], %c1
+// CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[RESHAPE]], %c2
+// CHECK-DAG: %[[M:.+]] = affine.apply {{.*}}[%[[DIM1]]]
+// CHECK-DAG: iree_encoding.set_encoding %[[COLLAPSED]] encoding_dims{%[[M]], %[[DIM2]]}
+// CHECK-DAG: iree_encoding.set_encoding %[[RHS]] encoding_dims{%[[M]], %[[DIM2]]}
+// CHECK: flow.dispatch.region
+// CHECK: linalg.matmul
+// CHECK: iree_encoding.unset_encoding
+// CHECK: flow.return
+
+// -----
+
+// Mixed-chain rejection: `some_unhoistable_op` is shared between the metadata path
+// (it is an encoding_dim of the SetEncodingOp) and the data path (it is a
+// dynamic size of the tensor.extract_slice that becomes the SetEncodingOp's
+// source).
+
+#map_rej = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_rej1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_rej2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#enc_rej = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_rej, #map_rej1, #map_rej2], iteration_sizes = [?, 256, 256]>
+util.func public @no_hoist_when_data_path_shares_metadata_op(
+ %src: tensor<?x256xf32>, %m: index) -> tensor<?x256xf32> {
+ %0 = flow.dispatch.region -> (tensor<?x256xf32>{%m}) {
+ %c0 = arith.constant 0 : index
+ %d = "some_unhoistable_op"(%src) : (tensor<?x256xf32>) -> index
+ %slice = tensor.extract_slice %src[0, 0] [%d, 256] [1, 1] : tensor<?x256xf32> to tensor<?x256xf32>
+ %enc = iree_encoding.set_encoding %slice encoding_dims{%d} : tensor<?x256xf32> -> tensor<?x256xf32, #enc_rej>
+ %unset = iree_encoding.unset_encoding %enc encoding_dims{%d} : tensor<?x256xf32, #enc_rej> -> tensor<?x256xf32>{%m}
+ flow.return %unset : tensor<?x256xf32>
+ }
+ util.return %0 : tensor<?x256xf32>
+}
+// CHECK-LABEL: @no_hoist_when_data_path_shares_metadata_op
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[DIM:.+]] = "some_unhoistable_op"
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice
+// CHECK: %[[SET:.+]] = iree_encoding.set_encoding %[[SLICE]] encoding_dims{%[[DIM]]}
+// CHECK: iree_encoding.unset_encoding %[[SET]] encoding_dims{%[[DIM]]}
+// CHECK: flow.return
