[SPIRV] Run pad alloc pass in spirv (#9326)
Re-use the pass already ran by LLVMGPU pipelines.
diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD b/compiler/src/iree/compiler/Codegen/Common/BUILD
index 768b34f..d748ea7 100644
--- a/compiler/src/iree/compiler/Codegen/Common/BUILD
+++ b/compiler/src/iree/compiler/Codegen/Common/BUILD
@@ -51,6 +51,7 @@
"LinalgBufferizePass.cpp",
"MemrefCopyToLinalg.cpp",
"OptimizeVectorTransferPass.cpp",
+ "PadDynamicAlloc.cpp",
"PolynomialApproximationPass.cpp",
"RemoveTrivialLoops.cpp",
"RewriteLinalgDestructiveUpdatesPass.cpp",
diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
index 025fe35..55286cc 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
@@ -43,6 +43,7 @@
"LinalgBufferizePass.cpp"
"MemrefCopyToLinalg.cpp"
"OptimizeVectorTransferPass.cpp"
+ "PadDynamicAlloc.cpp"
"PolynomialApproximationPass.cpp"
"RemoveTrivialLoops.cpp"
"RewriteLinalgDestructiveUpdatesPass.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUPadDynamicAlloc.cpp b/compiler/src/iree/compiler/Codegen/Common/PadDynamicAlloc.cpp
similarity index 81%
rename from compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUPadDynamicAlloc.cpp
rename to compiler/src/iree/compiler/Codegen/Common/PadDynamicAlloc.cpp
index 27a93f6..38af3af 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUPadDynamicAlloc.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/PadDynamicAlloc.cpp
@@ -7,7 +7,6 @@
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Passes.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
-#include "mlir/Dialect/GPU/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
namespace mlir {
@@ -49,18 +48,13 @@
namespace {
-struct LLVMGPUPadDynamicAllocPass
- : public LLVMGPUPadDynamicAllocBase<LLVMGPUPadDynamicAllocPass> {
+struct PadDynamicAllocPass : public PadDynamicAllocBase<PadDynamicAllocPass> {
void runOnOperation() override {
auto funcOp = getOperation();
SmallVector<memref::AllocOp> sharedMemAllocs;
// Collect all the alloc operations.
- funcOp.walk([&](memref::AllocOp allocOp) {
- if (allocOp.getType().getMemorySpaceAsInt() ==
- gpu::GPUDialect::getWorkgroupAddressSpace()) {
- sharedMemAllocs.push_back(allocOp);
- }
- });
+ funcOp.walk(
+ [&](memref::AllocOp allocOp) { sharedMemAllocs.push_back(allocOp); });
for (memref::AllocOp alloc : sharedMemAllocs) {
if (failed(padAlloc(alloc))) return signalPassFailure();
}
@@ -68,8 +62,8 @@
};
} // namespace
-std::unique_ptr<OperationPass<func::FuncOp>> createLLVMGPUPadDynamicAlloc() {
- return std::make_unique<LLVMGPUPadDynamicAllocPass>();
+std::unique_ptr<OperationPass<func::FuncOp>> createPadDynamicAlloc() {
+ return std::make_unique<PadDynamicAllocPass>();
}
} // namespace iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD b/compiler/src/iree/compiler/Codegen/Common/test/BUILD
index f44b137..5a93dcc 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD
+++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD
@@ -32,6 +32,7 @@
"insert_distribution_info.mlir",
"iree_comprehensive_bufferize.mlir",
"linalg_bufferize.mlir",
+ "pad_dynamic_alloc.mlir",
"remove_dead_allocs.mlir",
"remove_trivial_loops.mlir",
"rewrite_linalg_destructive_updates.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
index b71a981..803f5b0 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
@@ -27,6 +27,7 @@
"insert_distribution_info.mlir"
"iree_comprehensive_bufferize.mlir"
"linalg_bufferize.mlir"
+ "pad_dynamic_alloc.mlir"
"remove_dead_allocs.mlir"
"remove_trivial_loops.mlir"
"rewrite_linalg_destructive_updates.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/pad_dynamic_alloc.mlir b/compiler/src/iree/compiler/Codegen/Common/test/pad_dynamic_alloc.mlir
similarity index 91%
rename from compiler/src/iree/compiler/Codegen/LLVMGPU/test/pad_dynamic_alloc.mlir
rename to compiler/src/iree/compiler/Codegen/Common/test/pad_dynamic_alloc.mlir
index 3181b7e..efb223e 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/pad_dynamic_alloc.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/pad_dynamic_alloc.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -iree-llvmgpu-pad-dynamic-alloc %s | FileCheck %s
+// RUN: iree-opt -iree-codegen-pad-dynamic-alloc %s | FileCheck %s
// CHECK-LABEL: dynamic_alloc
func.func @dynamic_alloc(%id : index) {
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD b/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD
index 7a7d979..02f5f54 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD
@@ -21,7 +21,6 @@
"KernelConfig.cpp",
"LLVMGPULowerExecutableTarget.cpp",
"LLVMGPUMultiBuffering.cpp",
- "LLVMGPUPadDynamicAlloc.cpp",
"LLVMGPUReduceBankConflicts.cpp",
"LLVMGPUTensorCoreVectorization.cpp",
"LLVMGPUTileAndDistribute.cpp",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt
index a9ec926..0d5d229 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt
@@ -23,7 +23,6 @@
"KernelConfig.cpp"
"LLVMGPULowerExecutableTarget.cpp"
"LLVMGPUMultiBuffering.cpp"
- "LLVMGPUPadDynamicAlloc.cpp"
"LLVMGPUReduceBankConflicts.cpp"
"LLVMGPUTensorCoreVectorization.cpp"
"LLVMGPUTileAndDistribute.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index 665ddd8..093b2ae 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -189,7 +189,7 @@
// Pad allocations with dynamic dimension before lowering of SCF and affine
// but after linalg lowering.
- pm.addNestedPass<func::FuncOp>(createLLVMGPUPadDynamicAlloc());
+ pm.addNestedPass<func::FuncOp>(createPadDynamicAlloc());
pm.addPass(createLowerAffinePass());
pm.addPass(createCanonicalizerPass());
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD
index 182722f..f3e1b85 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD
@@ -24,7 +24,6 @@
"distribute_to_thread.mlir",
"gpu_set_num_workgroups.mlir",
"nvvm_pipeline_test.mlir",
- "pad_dynamic_alloc.mlir",
"reduce_bank_conflicts.mlir",
"rocdl_pipeline_test.mlir",
"illegal_configuration.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
index a17c533..c96a726 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
@@ -21,7 +21,6 @@
"illegal_configuration.mlir"
"legalize.mlir"
"nvvm_pipeline_test.mlir"
- "pad_dynamic_alloc.mlir"
"reduce_bank_conflicts.mlir"
"rocdl_pipeline_test.mlir"
"tensorcore_vectorization.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Passes.h b/compiler/src/iree/compiler/Codegen/Passes.h
index a92efbc..301a4bd 100644
--- a/compiler/src/iree/compiler/Codegen/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/Passes.h
@@ -167,6 +167,9 @@
std::unique_ptr<OperationPass<func::FuncOp>> createWorkGroupSwizzle(
unsigned swizzleLogTile = 0);
+/// Pad dynamic alloc op to convert them into static one.
+std::unique_ptr<OperationPass<func::FuncOp>> createPadDynamicAlloc();
+
//----------------------------------------------------------------------------//
// Common codegen patterns.
//----------------------------------------------------------------------------//
@@ -372,9 +375,6 @@
/// Converts vector ops to gpu dialect.
std::unique_ptr<OperationPass<func::FuncOp>> createLLVMGPUVectorToGPU();
-/// Pad dynamic alloc op to convert them into static one.
-std::unique_ptr<OperationPass<func::FuncOp>> createLLVMGPUPadDynamicAlloc();
-
//------------------------------------------------------------------------------
// SPIR-V Passes
//------------------------------------------------------------------------------
diff --git a/compiler/src/iree/compiler/Codegen/Passes.td b/compiler/src/iree/compiler/Codegen/Passes.td
index edf2e1a..6ffb0cd 100644
--- a/compiler/src/iree/compiler/Codegen/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Passes.td
@@ -185,6 +185,12 @@
];
}
+def PadDynamicAlloc :
+ Pass<"iree-codegen-pad-dynamic-alloc", "func::FuncOp"> {
+ let summary = "Pass to pad dynamic alloc into static one.";
+ let constructor = "mlir::iree_compiler::createPadDynamicAlloc()";
+}
+
//------------------------------------------------------------------------------
// LLVMCPU
//------------------------------------------------------------------------------
@@ -322,12 +328,6 @@
let constructor = "mlir::iree_compiler::createLLVMGPUVectorToGPU()";
}
-def LLVMGPUPadDynamicAlloc :
- Pass<"iree-llvmgpu-pad-dynamic-alloc", "func::FuncOp"> {
- let summary = "Pass to pad dynamic alloc into static one.";
- let constructor = "mlir::iree_compiler::createLLVMGPUPadDynamicAlloc()";
-}
-
//------------------------------------------------------------------------------
// SPIR-V
//------------------------------------------------------------------------------
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
index 06d89d6..a13331f 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
@@ -116,6 +116,8 @@
// math dialect elementry functions -> polynomial form.
pm.addNestedPass<func::FuncOp>(createPolynomialApproximationPass());
+ pm.addNestedPass<func::FuncOp>(createPadDynamicAlloc());
+
// Fold load/store from/to subview ops into the original memref when possible.
// In SPIR-V we don't use memref descriptor so it's not possible to handle
// subview ops.
