Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions quadrants/runtime/amdgpu/jit_amdgpu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 <fstream>
Expand All @@ -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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

comments should be 120c wrapped, not 80c wrapped

// 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<AMDGPUSetLoopInterleavePass> {
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<llvm::MDNode>(existing_id->getOperand(i))) {
if (node->getNumOperands() > 0) {
if (auto *s = llvm::dyn_cast<llvm::MDString>(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<llvm::Metadata *, 4> 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<llvm::Module> M, int max_reg) {
auto hsaco = compile_module_to_hsaco(M);
QD_TRACE("hsaco size: {:.2f}KB", hsaco.size() / 1024.0);
Expand Down Expand Up @@ -142,6 +209,14 @@ std::string JITSessionAMDGPU::compile_module_to_hsaco(std::unique_ptr<llvm::Modu
pb.registerLoopAnalyses(lam);
pb.crossRegisterProxies(lam, fam, cgam, mam);

// Annotate innermost loops with interleave-count=8 via IR metadata so the

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

120c wrap

// LLVM loop vectorizer uses it. This is scoped to this compilation pipeline
// and avoids mutating process-wide command-line state.
pb.registerLoopOptimizerEndEPCallback(
[](llvm::LoopPassManager &lpm, llvm::OptimizationLevel) {
lpm.addPass(AMDGPUSetLoopInterleavePass(8));
});
Comment on lines +215 to +218

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep the GCN dump pipeline in sync

When print_kernel_amdgcn is enabled, the earlier dump path builds and runs a separate PassBuilder for module_clone without this callback, while the HSACO path below does get the interleave metadata. That means the emitted quadrants_kernel_amdgcn_*.gcn no longer represents the code being linked for this perf change, which can mislead AMDGPU tuning/debugging; register the same callback on the clone pipeline or share a helper.

Useful? React with 👍 / 👎.


llvm::ModulePassManager mpm = pb.buildPerModuleDefaultPipeline(llvm::OptimizationLevel::O3);

// Run the new optimization pipeline
Expand Down
3 changes: 3 additions & 0 deletions quadrants/runtime/llvm/llvm_context_pass.h
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,12 @@
#include "llvm/Transforms/IPO.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Metadata.h"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we still need these header changes?

#include "llvm/Support/SourceMgr.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Transforms/Scalar/LoopPassManager.h"
Comment on lines +15 to +16

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3 Badge Keep AMDGPU-only includes out of the shared header

The repo AGENTS.md asks reviewers to flag unnecessary contact area, and these loop-pass headers are only needed by the private AMDGPU pass in jit_amdgpu.cpp where they are already included. Since llvm_context_pass.h is included by common/CPU/CUDA code too, adding these dependencies expands the rebuild and dependency surface for unrelated backends; keep them local to the AMDGPU translation unit.

Useful? React with 👍 / 👎.


#if defined(QD_WITH_AMDGPU)
#include "quadrants/rhi/amdgpu/amdgpu_context.h"
Expand Down
Loading