DispatchLinalgOnTensorsViaRegions: Lower ExtractSliceOps to FlowTenso… (#11084)
…rSliceOps
This is to ensure that DispatchLinalgOnTensorsViaRegions produces the
same number of dispatches as DispatchLinalgOnTensors.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp b/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp
index 0d78b79..16b807c 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp
@@ -135,6 +135,9 @@
LogicalResult convertExtractSliceOpToFlowSliceOp(
RewriterBase &rewriter, tensor::ExtractSliceOp sliceOp) {
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(sliceOp);
+
if (sliceOp->getParentOfType<Flow::DispatchWorkgroupsOp>()) {
return failure();
}
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensorsViaRegionOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensorsViaRegionOps.cpp
index f77661d..39e093a 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensorsViaRegionOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensorsViaRegionOps.cpp
@@ -720,6 +720,71 @@
return wrapInWorkgroupsOp(rewriter, rootOps, generateWorkloadRegion);
}
+/// Return `true` if the given op is contained in DispatchWorkgroupsOp or in a
+/// DispatchRegionOp.
+static bool isInDispatchRegion(Operation *op) {
+ return op->getParentOfType<Flow::DispatchWorkgroupsOp>() ||
+ op->getParentOfType<Flow::DispatchRegionOp>();
+}
+
+/// Rewrite top-level InsertSliceOps to FlowUpdateOps or wrap them in a
+/// dispatch region.
+LogicalResult convertInsertSliceOps(
+ TensorDimTrackingRewriter &rewriter, mlir::FunctionOpInterface funcOp,
+ SmallVector<Flow::DispatchWorkgroupsOp> &workgroupsOps,
+ bool generateWorkloadRegion) {
+ // Find eligible InsertSliceOps.
+ SmallVector<tensor::InsertSliceOp> insertSliceOps;
+ funcOp.walk([&](tensor::InsertSliceOp op) {
+ if (!isInDispatchRegion(op)) insertSliceOps.push_back(op);
+ });
+
+ // Rewrite InsertSliceOps to FlowUpdateOps.
+ SmallVector<Operation *> remainingInsertSliceOps;
+ for (tensor::InsertSliceOp insertSliceOp : insertSliceOps)
+ if (failed(
+ Flow::convertInsertSliceOpToFlowUpdateOp(rewriter, insertSliceOp)))
+ remainingInsertSliceOps.push_back(insertSliceOp);
+
+ // Create a DispatchWorkgroupsOp for every remaining InsertSliceOp.
+ FailureOr<SmallVector<Flow::DispatchWorkgroupsOp>> newWorkgroupsOps =
+ wrapInWorkgroupsOp(rewriter, remainingInsertSliceOps,
+ generateWorkloadRegion);
+ if (failed(newWorkgroupsOps)) return failure();
+ workgroupsOps.append(newWorkgroupsOps->begin(), newWorkgroupsOps->end());
+
+ return success();
+}
+
+/// Rewrite top-level ExtractSliceOps to FlowSliceOps or wrap them in a
+/// dispatch region.
+LogicalResult convertExtractSliceOps(
+ TensorDimTrackingRewriter &rewriter, mlir::FunctionOpInterface funcOp,
+ SmallVector<Flow::DispatchWorkgroupsOp> &workgroupsOps,
+ bool generateWorkloadRegion) {
+ // Find eligible ExtractSliceOps.
+ SmallVector<tensor::ExtractSliceOp> extractSliceOps;
+ funcOp.walk([&](tensor::ExtractSliceOp op) {
+ if (!isInDispatchRegion(op)) extractSliceOps.push_back(op);
+ });
+
+ // Rewrite ExtractSliceOps to FlowSliceOps.
+ SmallVector<Operation *> remainingExtractSliceOps;
+ for (tensor::ExtractSliceOp extractSliceOp : extractSliceOps)
+ if (failed(
+ Flow::convertExtractSliceOpToFlowSliceOp(rewriter, extractSliceOp)))
+ remainingExtractSliceOps.push_back(extractSliceOp);
+
+ // Create a DispatchWorkgroupsOp for every remaining ExtractSliceOp.
+ FailureOr<SmallVector<Flow::DispatchWorkgroupsOp>> newWorkgroupsOps =
+ wrapInWorkgroupsOp(rewriter, remainingExtractSliceOps,
+ generateWorkloadRegion);
+ if (failed(newWorkgroupsOps)) return failure();
+ workgroupsOps.append(newWorkgroupsOps->begin(), newWorkgroupsOps->end());
+
+ return success();
+}
+
namespace {
/// Pass declaration.
struct DispatchLinalgOnTensorsViaRegionOpsPass
@@ -745,7 +810,7 @@
} // namespace
void DispatchLinalgOnTensorsViaRegionOpsPass::runOnOperation() {
- auto funcOp = getOperation();
+ mlir::FunctionOpInterface funcOp = getOperation();
MLIRContext *context = &getContext();
DominanceInfo const &dominanceInfo = getAnalysis<DominanceInfo>();
@@ -763,30 +828,20 @@
llvm::dbgs() << "\n\n";
});
- // Step 2a: Rewrite InsertSliceOps to TensorUpdateOps.
- SmallVector<tensor::InsertSliceOp> insertSliceOps;
- SmallVector<Operation *> remainingInsertSliceOps;
- funcOp.walk([&](tensor::InsertSliceOp op) {
- if (!op->getParentOfType<Flow::DispatchRegionOp>())
- insertSliceOps.push_back(op);
- });
- for (tensor::InsertSliceOp insertSliceOp : insertSliceOps)
- if (failed(
- Flow::convertInsertSliceOpToFlowUpdateOp(rewriter, insertSliceOp)))
- remainingInsertSliceOps.push_back(insertSliceOp);
+ // Step 2: Rewrite InsertSliceOps to FlowUpdateOps.
+ if (failed(convertInsertSliceOps(rewriter, funcOp, workgroupsOps,
+ generateWorkloadRegion)))
+ return signalPassFailure();
- // Step 2b: Create a DispatchWorkgroupsOp for every remaining InsertSliceOp.
+ // Step 3: Rewrite ExtractSliceOps to FlowUpdateOps.
+ if (failed(convertExtractSliceOps(rewriter, funcOp, workgroupsOps,
+ generateWorkloadRegion)))
+ return signalPassFailure();
+
+ // Step 4: Create a DispatchWorkgroupsOp for certain other ops.
FailureOr<SmallVector<Flow::DispatchWorkgroupsOp>> newWorkgroupsOps =
- wrapInWorkgroupsOp(rewriter, remainingInsertSliceOps,
- generateWorkloadRegion);
- if (failed(newWorkgroupsOps)) return signalPassFailure();
- workgroupsOps.append(newWorkgroupsOps->begin(), newWorkgroupsOps->end());
-
- // Step 3: Create a DispatchWorkgroupsOp for certain other ops.
- newWorkgroupsOps =
- wrapInWorkgroupsOp<tensor::ExtractSliceOp, LinalgExt::SetEncodingOp,
- LinalgExt::UnsetEncodingOp>(rewriter, funcOp,
- generateWorkloadRegion);
+ wrapInWorkgroupsOp<LinalgExt::SetEncodingOp, LinalgExt::UnsetEncodingOp>(
+ rewriter, funcOp, generateWorkloadRegion);
if (failed(newWorkgroupsOps)) return signalPassFailure();
workgroupsOps.append(newWorkgroupsOps->begin(), newWorkgroupsOps->end());
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
index 2a3d3cf..54e6a33 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
@@ -2054,3 +2054,31 @@
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[FILL]] :
// CHECK: flow.dispatch.tensor.store %[[GEMM]], %[[RESULT]]
+
+// -----
+
+func.func @extract_slice1(%arg0 : tensor<5x24x48xf32>) -> tensor<4xf32> {
+ %0 = tensor.extract_slice %arg0[2, 3, 4] [1, 1, 4] [1, 1, 1]
+ : tensor<5x24x48xf32> to tensor<4xf32>
+ return %0 : tensor<4xf32>
+}
+
+// CHECK-LABEL: func.func @extract_slice1(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<5x24x48xf32>)
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK: %[[SLICE:.+]] = flow.tensor.slice %[[ARG0]][%[[C2]], %[[C3]], %[[C4]] for %[[C1]], %[[C1]], %[[C4]]]
+// CHECK: %[[RESULT:.+]] = flow.tensor.reshape %[[SLICE]]
+// CHECK: return %[[RESULT]]
+
+// CHECK-VIA-REGIONS-LABEL: func.func @extract_slice1(
+// CHECK-VIA-REGIONS-SAME: %[[ARG0:.+]]: tensor<5x24x48xf32>)
+// CHECK-VIA-REGIONS-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-VIA-REGIONS-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK-VIA-REGIONS-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-VIA-REGIONS-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-VIA-REGIONS: %[[SLICE:.+]] = flow.tensor.slice %[[ARG0]][%[[C2]], %[[C3]], %[[C4]] for %[[C1]], %[[C1]], %[[C4]]]
+// CHECK-VIA-REGIONS: %[[RESULT:.+]] = flow.tensor.reshape %[[SLICE]]
+// CHECK-VIA-REGIONS: return %[[RESULT]]
