diff --git a/quadrants/runtime/amdgpu/jit_amdgpu.cpp b/quadrants/runtime/amdgpu/jit_amdgpu.cpp index 7a2dccf783..28ab970f4b 100644 --- a/quadrants/runtime/amdgpu/jit_amdgpu.cpp +++ b/quadrants/runtime/amdgpu/jit_amdgpu.cpp @@ -2,7 +2,13 @@ #include "quadrants/runtime/llvm/llvm_context.h" #include "quadrants/runtime/llvm/llvm_context_pass.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/Metadata.h" #include "llvm/IR/Module.h" +#include "llvm/IR/PassManager.h" +#include "llvm/Passes/PassBuilder.h" +#include "llvm/Analysis/LoopInfo.h" +#include "llvm/Transforms/Scalar/LoopPassManager.h" #include "llvm/Transforms/Utils/Cloning.h" #include @@ -11,6 +17,72 @@ namespace quadrants { namespace lang { #if defined(QD_WITH_AMDGPU) + +namespace { +// New-PM loop pass: annotates innermost loops with llvm.loop.interleave.count IR metadata so the loop vectorizer uses +// the requested interleave factor. Scoped to the compilation pipeline where it is registered; does not touch any +// process-wide LLVM command-line state. Kept private to this AMDGPU JIT translation unit to avoid expanding the shared +// LLVM header surface area. +struct AMDGPUSetLoopInterleavePass : public llvm::PassInfoMixin { + unsigned count_; + explicit AMDGPUSetLoopInterleavePass(unsigned count) : count_(count) { + } + + llvm::PreservedAnalyses run(llvm::Loop &L, + llvm::LoopAnalysisManager &, + llvm::LoopStandardAnalysisResults &, + llvm::LPMUpdater &) { + // Only annotate innermost loops + if (!L.getSubLoops().empty()) + return llvm::PreservedAnalyses::all(); + + llvm::LLVMContext &ctx = L.getHeader()->getContext(); + llvm::MDNode *existing_id = L.getLoopID(); + + // Skip if interleave count is already specified + if (existing_id) { + for (unsigned i = 1, e = existing_id->getNumOperands(); i < e; ++i) { + if (auto *node = llvm::dyn_cast(existing_id->getOperand(i))) { + if (node->getNumOperands() > 0) { + if (auto *s = llvm::dyn_cast(node->getOperand(0))) { + if (s->getString() == "llvm.loop.interleave.count") + return llvm::PreservedAnalyses::all(); + } + } + } + } + } + + // Build new interleave hint node + llvm::MDNode *hint = llvm::MDNode::get( + ctx, {llvm::MDString::get(ctx, "llvm.loop.interleave.count"), + llvm::ConstantAsMetadata::get(llvm::ConstantInt::get(llvm::Type::getInt32Ty(ctx), count_))}); + + // Assemble updated loop metadata: [self-ref, ...existing..., hint] + llvm::SmallVector ops; + ops.push_back(nullptr); // placeholder for self-reference + if (existing_id) { + for (unsigned i = 1, e = existing_id->getNumOperands(); i < e; ++i) + ops.push_back(existing_id->getOperand(i).get()); + } + ops.push_back(hint); + + llvm::MDNode *new_id = llvm::MDNode::get(ctx, ops); + new_id->replaceOperandWith(0, new_id); // fix self-reference + L.setLoopID(new_id); + + return llvm::PreservedAnalyses::none(); + } +}; + +// Registers AMDGPUSetLoopInterleavePass on the given PassBuilder. Shared by the HSACO code-generation pipeline and the +// print_kernel_amdgcn dump pipeline so the emitted .gcn reflects the same interleave metadata as the linked object. +void register_loop_interleave(llvm::PassBuilder &pb) { + pb.registerLoopOptimizerEndEPCallback( + [](llvm::LoopPassManager &lpm, llvm::OptimizationLevel) { lpm.addPass(AMDGPUSetLoopInterleavePass(8)); }); +} +} // namespace + JITModule *JITSessionAMDGPU ::add_module(std::unique_ptr M, int max_reg) { auto hsaco = compile_module_to_hsaco(M); QD_TRACE("hsaco size: {:.2f}KB", hsaco.size() / 1024.0); @@ -113,6 +185,9 @@ std::string JITSessionAMDGPU::compile_module_to_hsaco(std::unique_ptr