Change vector interleaving to 8 (#8)

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

quadrants/runtime/amdgpu/jit_amdgpu.cpp

@@ -142,6 +142,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
+  // 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));
+      });
+
   llvm::ModulePassManager mpm = pb.buildPerModuleDefaultPipeline(llvm::OptimizationLevel::O3);
 
   // Run the new optimization pipeline

quadrants/runtime/llvm/llvm_context_pass.h

@@ -8,9 +8,12 @@
 #include "llvm/Transforms/IPO.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Metadata.h"
 #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"
 
 #if defined(QD_WITH_AMDGPU)
 #include "quadrants/rhi/amdgpu/amdgpu_context.h"
@@ -295,5 +298,63 @@ struct AMDGPUConvertFuncParamAddressSpacePass : public ModulePass {
 
 #endif
 
+// 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.
+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 lang
 }  // namespace quadrants