Skip to content

Commit e1e6df5

Browse files
[SYCL][Fusion] JIT compiler kernel fusion passes (#7661)
This is the fourth patch in a series of patches to add an implementation of the [kernel fusion extension](#7098). We have split the implementation into multiple patches to make them more easy to review. This patch adds the LLVM passes that perform the kernel fusion and related optimizations: * A pass creating the function definition for the fused kernel from the input kernel definitions. * A pass performing internalization of dataflow internal to the fused kernel into either private or local memory. The type of memory to use is currently specified by the user in the runtime. * A pass propagating values for scalars and by-val aggregates from the SYCL runtime to the fused kernel as constants. The information is propagated from the SYCL runtime to the passes via LLVM metadata inserted by the JIT compiler frontend. After and between the fusion passes, some standard LLVM optimization and transformation passes are executed to enable passes and optimize the fused kernel. Signed-off-by: Lukas Sommer <[email protected]> Co-authored-by: Victor Perez <[email protected]>
1 parent 1a1fd8d commit e1e6df5

38 files changed

+4375
-4
lines changed

sycl-fusion/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,5 @@ set(SYCL_JIT_BASE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
1010
set(LLVM_SPIRV_INCLUDE_DIRS "${LLVM_MAIN_SRC_DIR}/../llvm-spirv/include")
1111

1212
add_subdirectory(jit-compiler)
13+
add_subdirectory(passes)
14+
add_subdirectory(test)

sycl-fusion/jit-compiler/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ add_llvm_library(sycl-fusion
44
lib/KernelFusion.cpp
55
lib/JITContext.cpp
66
lib/translation/SPIRVLLVMTranslation.cpp
7+
lib/fusion/FusionPipeline.cpp
78
lib/fusion/FusionHelper.cpp
89
lib/fusion/ModuleHelper.cpp
910
lib/helper/ConfigHelper.cpp
@@ -34,6 +35,7 @@ find_package(Threads REQUIRED)
3435
target_link_libraries(sycl-fusion
3536
PRIVATE
3637
LLVMSPIRVLib
38+
SYCLKernelFusionPasses
3739
${CMAKE_THREAD_LIBS_INIT}
3840
)
3941

sycl-fusion/jit-compiler/lib/KernelFusion.cpp

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include "KernelIO.h"
1212
#include "Options.h"
1313
#include "fusion/FusionHelper.h"
14+
#include "fusion/FusionPipeline.h"
1415
#include "helper/ConfigHelper.h"
1516
#include "helper/ErrorHandling.h"
1617
#include "translation/SPIRVLLVMTranslation.h"
@@ -80,9 +81,10 @@ FusionResult KernelFusion::fuseKernels(
8081
}
8182
std::unique_ptr<llvm::Module> NewMod = std::move(*NewModOrError);
8283

83-
// TODO: Invoke the actual fusion via LLVM pass manager.
84-
// This will be added in a later PR.
85-
auto NewModInfo = std::make_unique<SYCLModuleInfo>();
84+
// Invoke the actual fusion via LLVM pass manager.
85+
std::unique_ptr<SYCLModuleInfo> NewModInfo =
86+
fusion::FusionPipeline::runFusionPasses(*NewMod, ModuleInfo,
87+
BarriersFlags);
8688

8789
// Get the updated kernel info for the fused kernel and add the information to
8890
// the existing KernelInfo.

sycl-fusion/jit-compiler/lib/fusion/FusionHelper.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
//==--------- FusionHelper.h - helpers to insert fused kernel stub ---------==//
1+
//==--------- FusionHelper.h - Helpers to insert fused kernel stub ---------==//
22
//
33
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
44
// See https://llvm.org/LICENSE.txt for license information.
Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
//==-------------------------- FusionPipeline.cpp --------------------------==//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "FusionPipeline.h"
10+
11+
#include "debug/PassDebug.h"
12+
#include "helper/ConfigHelper.h"
13+
#include "internalization/Internalization.h"
14+
#include "kernel-fusion/SYCLKernelFusion.h"
15+
#include "kernel-info/SYCLKernelInfo.h"
16+
#include "syclcp/SYCLCP.h"
17+
18+
#include "llvm/IR/PassManager.h"
19+
#include "llvm/Transforms/Scalar/IndVarSimplify.h"
20+
#include "llvm/Transforms/Scalar/InferAddressSpaces.h"
21+
#include "llvm/Transforms/Scalar/LoopUnrollPass.h"
22+
#ifndef NDEBUG
23+
#include "llvm/IR/Verifier.h"
24+
#endif // NDEBUG
25+
#include "llvm/Passes/PassBuilder.h"
26+
#include "llvm/Transforms/InstCombine/InstCombine.h"
27+
#include "llvm/Transforms/Scalar/ADCE.h"
28+
#include "llvm/Transforms/Scalar/EarlyCSE.h"
29+
#include "llvm/Transforms/Scalar/SCCP.h"
30+
#include "llvm/Transforms/Scalar/SROA.h"
31+
#include "llvm/Transforms/Scalar/SimplifyCFG.h"
32+
33+
using namespace llvm;
34+
using namespace jit_compiler;
35+
using namespace jit_compiler::fusion;
36+
37+
std::unique_ptr<SYCLModuleInfo>
38+
FusionPipeline::runFusionPasses(Module &Mod, SYCLModuleInfo &InputInfo,
39+
int BarriersFlags) {
40+
// Perform the actual kernel fusion, i.e., generate a kernel function for the
41+
// fused kernel from the kernel functions of the input kernels. This is done
42+
// by the SYCLKernelFusion LLVM pass, which is run here through a custom LLVM
43+
// pass pipeline. In order to perform internalization, we run the
44+
// SYCLInternalizer pass.
45+
46+
bool DebugEnabled = ConfigHelper::get<option::JITEnableVerbose>();
47+
if (DebugEnabled) {
48+
// Enabled debug output from the fusion passes.
49+
jit_compiler::PassDebug = true;
50+
}
51+
52+
// Initialize the analysis managers with all the registered analyses.
53+
PassBuilder PB;
54+
LoopAnalysisManager LAM;
55+
FunctionAnalysisManager FAM;
56+
CGSCCAnalysisManager CGAM;
57+
ModuleAnalysisManager MAM;
58+
PB.registerModuleAnalyses(MAM);
59+
PB.registerCGSCCAnalyses(CGAM);
60+
PB.registerFunctionAnalyses(FAM);
61+
PB.registerLoopAnalyses(LAM);
62+
PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
63+
64+
// Make the existing SYCLModuleInfo available to the pass pipeline via the
65+
// corresponding analysis pass.
66+
MAM.registerPass([&]() {
67+
auto ModInfo = std::make_unique<SYCLModuleInfo>(InputInfo);
68+
return SYCLModuleInfoAnalysis{std::move(ModInfo)};
69+
});
70+
ModulePassManager MPM;
71+
// Run the fusion pass on the LLVM IR module.
72+
MPM.addPass(SYCLKernelFusion{BarriersFlags});
73+
{
74+
FunctionPassManager FPM;
75+
// Run loop unrolling and SROA to split the kernel functor struct into its
76+
// scalar parts, to avoid problems with address-spaces and enable
77+
// internalization.
78+
FPM.addPass(createFunctionToLoopPassAdaptor(IndVarSimplifyPass{}));
79+
LoopUnrollOptions UnrollOptions;
80+
FPM.addPass(LoopUnrollPass{UnrollOptions});
81+
FPM.addPass(SROAPass{});
82+
// Run the InferAddressSpace pass to remove as many address-space casts
83+
// to/from generic address-space as possible, because these hinder
84+
// internalization.
85+
// FIXME: TTI should tell the pass which address space to use.
86+
// Ideally, the static compiler should have performed that job.
87+
constexpr unsigned FlatAddressSpace = 4;
88+
FPM.addPass(InferAddressSpacesPass(FlatAddressSpace));
89+
MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
90+
}
91+
// Run dataflow internalization and runtime constant propagation.
92+
MPM.addPass(SYCLInternalizer{});
93+
MPM.addPass(SYCLCP{});
94+
// Run additional optimization passes after completing fusion.
95+
{
96+
FunctionPassManager FPM;
97+
FPM.addPass(SROAPass{});
98+
FPM.addPass(SCCPPass{});
99+
FPM.addPass(InstCombinePass{});
100+
FPM.addPass(SimplifyCFGPass{});
101+
FPM.addPass(SROAPass{});
102+
FPM.addPass(InstCombinePass{});
103+
FPM.addPass(SimplifyCFGPass{});
104+
FPM.addPass(ADCEPass{});
105+
FPM.addPass(EarlyCSEPass{/*UseMemorySSA*/ true});
106+
MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
107+
}
108+
MPM.run(Mod, MAM);
109+
110+
if (DebugEnabled) {
111+
// Restore debug option
112+
jit_compiler::PassDebug = false;
113+
}
114+
115+
assert(!verifyModule(Mod, &errs()) && "Invalid LLVM IR generated");
116+
117+
auto NewModInfo = MAM.getResult<SYCLModuleInfoAnalysis>(Mod);
118+
assert(NewModInfo.ModuleInfo && "Failed to retrieve SYCL module info");
119+
120+
return std::make_unique<SYCLModuleInfo>(std::move(*NewModInfo.ModuleInfo));
121+
}
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
//==--- FusionPipeline - LLVM pass pipeline definition for kernel fusion ---==//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef SYCL_FUSION_JIT_COMPILER_FUSION_FUSIONPIPELINE_H
10+
#define SYCL_FUSION_JIT_COMPILER_FUSION_FUSIONPIPELINE_H
11+
12+
#include "Kernel.h"
13+
#include "llvm/IR/Module.h"
14+
15+
namespace jit_compiler {
16+
namespace fusion {
17+
18+
class FusionPipeline {
19+
public:
20+
///
21+
/// Run the necessary passes in a custom pass pipeline to perform kernel
22+
/// fusion on the given module. The module should contain the stub functions
23+
/// and fusion metadata. The given SYCLModuleInfo must contain information
24+
/// about all input kernels. The returned SYCLModuleInfo will additionally
25+
/// contain an entry for the fused kernel.
26+
static std::unique_ptr<SYCLModuleInfo>
27+
runFusionPasses(llvm::Module &Mod, SYCLModuleInfo &InputInfo,
28+
int BarriersFlags);
29+
};
30+
} // namespace fusion
31+
} // namespace jit_compiler
32+
33+
#endif // SYCL_FUSION_JIT_COMPILER_FUSION_FUSIONPIPELINE_H

sycl-fusion/passes/CMakeLists.txt

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Module library for usage as library/pass-plugin with LLVM opt.
2+
add_llvm_library(SYCLKernelFusion SHARED
3+
SYCLFusionPasses.cpp
4+
kernel-fusion/SYCLKernelFusion.cpp
5+
kernel-info/SYCLKernelInfo.cpp
6+
internalization/Internalization.cpp
7+
syclcp/SYCLCP.cpp
8+
cleanup/Cleanup.cpp
9+
debug/PassDebug.cpp
10+
)
11+
12+
target_include_directories(SYCLKernelFusion
13+
PUBLIC
14+
${CMAKE_CURRENT_SOURCE_DIR}
15+
PRIVATE
16+
${SYCL_JIT_BASE_DIR}/common/include
17+
)
18+
19+
# Static library for linking with the jit_compiler
20+
add_llvm_library(SYCLKernelFusionPasses
21+
SYCLFusionPasses.cpp
22+
kernel-fusion/SYCLKernelFusion.cpp
23+
kernel-info/SYCLKernelInfo.cpp
24+
internalization/Internalization.cpp
25+
syclcp/SYCLCP.cpp
26+
cleanup/Cleanup.cpp
27+
debug/PassDebug.cpp
28+
29+
LINK_COMPONENTS
30+
Core
31+
Support
32+
TransformUtils
33+
Passes
34+
)
35+
36+
target_include_directories(SYCLKernelFusionPasses
37+
PUBLIC
38+
$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
39+
PRIVATE
40+
${SYCL_JIT_BASE_DIR}/common/include
41+
)
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
//==------------------------ SYCLFusionPasses.cpp --------------------------==//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "llvm/Passes/PassBuilder.h"
10+
#include "llvm/Passes/PassPlugin.h"
11+
12+
#include "internalization/Internalization.h"
13+
#include "kernel-fusion/SYCLKernelFusion.h"
14+
#include "kernel-info/SYCLKernelInfo.h"
15+
#include "syclcp/SYCLCP.h"
16+
17+
using namespace llvm;
18+
19+
cl::opt<bool>
20+
NoBarriers("sycl-kernel-fusion-no-barriers",
21+
cl::desc("Disable barrier insertion for SYCL kernel fusion."));
22+
23+
llvm::PassPluginLibraryInfo getSYCLKernelFusionPluginInfo() {
24+
return {
25+
LLVM_PLUGIN_API_VERSION, "SYCL-Module-Info", LLVM_VERSION_STRING,
26+
[](PassBuilder &PB) {
27+
PB.registerPipelineParsingCallback(
28+
[](StringRef Name, ModulePassManager &MPM,
29+
ArrayRef<PassBuilder::PipelineElement>) {
30+
if (Name == "sycl-kernel-fusion") {
31+
int BarrierFlag =
32+
(NoBarriers) ? -1 : SYCLKernelFusion::DefaultBarriersFlags;
33+
MPM.addPass(SYCLKernelFusion(BarrierFlag));
34+
return true;
35+
}
36+
if (Name == "sycl-internalization") {
37+
MPM.addPass(SYCLInternalizer());
38+
return true;
39+
}
40+
if (Name == "sycl-cp") {
41+
MPM.addPass(SYCLCP());
42+
return true;
43+
}
44+
if (Name == "print-sycl-module-info") {
45+
MPM.addPass(SYCLModuleInfoPrinter());
46+
return true;
47+
}
48+
return false;
49+
});
50+
PB.registerAnalysisRegistrationCallback([](ModuleAnalysisManager &MAM) {
51+
MAM.registerPass([]() { return SYCLModuleInfoAnalysis{}; });
52+
});
53+
}};
54+
}
55+
56+
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
57+
llvmGetPassPluginInfo() {
58+
return getSYCLKernelFusionPluginInfo();
59+
}

0 commit comments

Comments
 (0)