[Codegen][GPU] Add fusion barrier after result promotion (#21709)
Currently result promotion relies on inserting two copies, one to
workgroup memory that is supposed to fuse with its producer, and one
with `derived_thread_config` that controls the distribution of the
consumer. This however has nothing forcing the intended fusion behavior,
and in all exercised cases today relies on reshapes missing fusion
patterns.
This patch introduces a fusion barrier op that sits between the two
copies to prevent fusion and gets dropped when bufferizing.
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp
index 0809e83..d953725 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp
@@ -6,6 +6,7 @@
#include "iree/compiler/Codegen/Common/GPU/Passes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
@@ -95,8 +96,13 @@
alloc.setMemorySpaceAttr(addressSpace);
auto copy =
linalg::CopyOp::create(rewriter, loc, valToMakeShared, alloc.getResult());
-
Value replacement = copy.getResult(0);
+
+ // Insert a fusion barrier to prevent the copy to workgroup memory from fusing
+ // into the promoted copy.
+ replacement = IREE::Codegen::FusionBarrierOp::create(
+ rewriter, replacement.getLoc(), replacement.getType(), replacement);
+
// If in extract slice is present we make it consume the new copy.
if (extractSliceOp) {
extractSliceOp.getSourceMutable().assign(replacement);
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
index e74a7ce..87c4402 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
@@ -261,6 +261,7 @@
"::mlir::bufferization::BufferizationDialect",
"::mlir::gpu::GPUDialect",
"::mlir::linalg::LinalgDialect",
+ "::mlir::iree_compiler::IREE::Codegen::IREECodegenDialect",
"::mlir::iree_compiler::IREE::GPU::IREEGPUDialect"
];
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_promote_matmul_operands.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_promote_matmul_operands.mlir
index d8b7959..cb209b3 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_promote_matmul_operands.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_promote_matmul_operands.mlir
@@ -125,9 +125,10 @@
// CHECK: %[[COPY1:.+]] = linalg.copy
// CHECK-SAME: ins(%[[MATMUL]] : tensor<?x?xf32>) outs(%[[ALLOC]] : tensor<?x?xf32>)
// CHECK-SAME: -> tensor<?x?xf32>
+// CHECK: %[[BARRIER:.+]] = iree_codegen.fusion_barrier %[[COPY1]]
// CHECK: %[[COPY2:.+]] = linalg.copy
// CHECK-SAME: {lowering_config = #iree_gpu.derived_thread_config}
-// CHECK-SAME: ins(%[[COPY1]] : tensor<?x?xf32>)
+// CHECK-SAME: ins(%[[BARRIER]] : tensor<?x?xf32>)
// CHECK: return %[[COPY2]] : tensor<?x?xf32>
// -----
@@ -152,7 +153,8 @@
// CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor
// CHECK: %[[COPY1:.+]] = linalg.copy
// CHECK-SAME: ins(%[[MATMUL]] : tensor<?x?xf32>) outs(%[[ALLOC]] : tensor<?x?xf32>)
-// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[COPY1]]
+// CHECK: %[[BARRIER:.+]] = iree_codegen.fusion_barrier %[[COPY1]]
+// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[BARRIER]]
// CHECK: %[[COPY2:.+]] = linalg.copy
// CHECK-SAME: {lowering_config = #iree_gpu.derived_thread_config}
// CHECK-SAME: ins(%[[EXTRACT]] : tensor<?x?xf32>)
diff --git a/compiler/src/iree/compiler/Codegen/Common/IREEComprehensiveBufferizePass.cpp b/compiler/src/iree/compiler/Codegen/Common/IREEComprehensiveBufferizePass.cpp
index 6187a9b..c6acdd0 100644
--- a/compiler/src/iree/compiler/Codegen/Common/IREEComprehensiveBufferizePass.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/IREEComprehensiveBufferizePass.cpp
@@ -12,6 +12,7 @@
#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Codegen/Common/Transforms.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Interfaces/BufferizationInterfaces.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
@@ -238,6 +239,14 @@
/// Run comprehensive bufferize.
void IREEComprehensiveBufferizePass::runOnOperation() {
mlir::FunctionOpInterface funcOp = getOperation();
+
+ // First drop all fusion barriers. Fusion is done at this point and these can
+ // be safely removed.
+ IRRewriter rewriter(funcOp.getContext());
+ funcOp.walk([&](IREE::Codegen::FusionBarrierOp barrier) {
+ rewriter.replaceOp(barrier, barrier.getSource());
+ });
+
IREEOneShotBufferizationOptions options = getBufferizationOptions();
options.testAnalysisOnly = testAnalysisOnly;
options.printConflicts = printConflicts;
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir b/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir
index 6f22330..fb9d0ef 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir
@@ -3207,3 +3207,16 @@
// CHECK: %[[CST:.+]] = arith.constant dense<0> : tensor<6xi32>
// CHECK: %[[MEMREF:.+]] = bufferization.to_buffer %[[CST]] read_only
// CHECK: memref.copy %[[MEMREF]]
+
+// -----
+
+func.func @drop_fusion_barrier() -> memref<6xf32> {
+ %alloc = bufferization.alloc_tensor() : tensor<6xf32>
+ %0 = iree_codegen.fusion_barrier %alloc : tensor<6xf32>
+ %memref = bufferization.to_buffer %0 : tensor<6xf32> to memref<6xf32>
+ return %memref : memref<6xf32>
+}
+
+// CHECK-LABEL: func.func @drop_fusion_barrier
+// CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<6xf32>
+// CHECK: return %[[ALLOC]]
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.td b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.td
index efd692f..1cf2d80 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.td
@@ -112,6 +112,22 @@
}
//===----------------------------------------------------------------------===//
+// FusionBarrierOp
+//===----------------------------------------------------------------------===//
+
+def IREECodegen_FusionBarrierOp :
+ Op<IREECodegen_Dialect, "fusion_barrier", [
+ SameOperandsAndResultType, Pure]> {
+ let summary = [{Prevents fusion through a tensor value}];
+
+ let arguments = (ins AnyRankedTensor:$source);
+ let results = (outs AnyRankedTensor:$result);
+ let assemblyFormat = [{
+ attr-dict $source `:` type($result)
+ }];
+}
+
+//===----------------------------------------------------------------------===//
// SwizzleHintOp
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/roundtrip.mlir b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/roundtrip.mlir
index f28f9d8..b3fca67 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/roundtrip.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/roundtrip.mlir
@@ -72,3 +72,13 @@
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]:
// CHECK: iree_codegen.store_to_buffer %[[ARG0]], %[[ARG1]]
// CHECK-SAME: : tensor<?x?xf32> into memref<?x?xf32, strided<[?, 1], offset: ?>>
+
+// -----
+
+func.func @fusion_barrier(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+ %0 = iree_codegen.fusion_barrier %arg0 : tensor<?xf32>
+ return %0 : tensor<?xf32>
+}
+// CHECK-LABEL: func.func @fusion_barrier(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?xf32>
+// CHECK: iree_codegen.fusion_barrier %[[ARG0]] : tensor<?xf32>
