diff --git a/quadrants/runtime/amdgpu/jit_amdgpu.cpp b/quadrants/runtime/amdgpu/jit_amdgpu.cpp index 7a2dccf783..479d404601 100644 --- a/quadrants/runtime/amdgpu/jit_amdgpu.cpp +++ b/quadrants/runtime/amdgpu/jit_amdgpu.cpp @@ -2,7 +2,12 @@ #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/Analysis/LoopInfo.h" +#include "llvm/Transforms/Scalar/LoopPassManager.h" #include "llvm/Transforms/Utils/Cloning.h" #include @@ -11,6 +16,68 @@ 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(); + } +}; +} // 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); @@ -142,6 +209,14 @@ std::string JITSessionAMDGPU::compile_module_to_hsaco(std::unique_ptr