[feat] Add original argument order mapping to kernel metadata

Polygeist reorders kernel arguments (scalars before pointers) when lowering
to GPU dialect. This commit extracts the original argument order from the
host wrapper function and includes it in the kernel metadata JSON.

The mapping allows the host stub generator to create launchers that accept
arguments in original order (matching hipLaunchKernelGGL) and reorder them
when packing for the device kernel.

Also fixes extractBaseKernelName to handle kernel names that contain "_kernel"
(e.g., vecadd_kernel_kernel94... now correctly extracts vecadd_kernel).

Author: ymweiss

lib/polygeist/Passes/ConvertGPUToVortex.cpp

@@ -54,15 +54,30 @@ constexpr uint32_t VX_CSR_LOCAL_MEM_BASE = 0xFC3;
 
 /// Extract base kernel name by removing Polygeist variant suffix
 /// Example: _Z12launch_basicPiS_ji_kernel94565344022848 -> _Z12launch_basicPiS_ji
+/// Example: __polygeist_launch_vecadd_kernel_kernel94... -> __polygeist_launch_vecadd_kernel
 static StringRef extractBaseKernelName(StringRef mangledName) {
-  size_t pos = mangledName.find("_kernel");
-  if (pos != StringRef::npos) {
-    // Find where the numeric suffix starts after "_kernel"
+  // Search from the end for "_kernel" followed by digits
+  // This handles cases like "vecadd_kernel_kernel94..." where the kernel name
+  // itself contains "_kernel"
+  size_t searchStart = 0;
+  size_t lastValidPos = StringRef::npos;
+
+  while (true) {
+    size_t pos = mangledName.find("_kernel", searchStart);
+    if (pos == StringRef::npos)
+      break;
+
     size_t suffixStart = pos + 7; // Length of "_kernel"
     if (suffixStart < mangledName.size() &&
         std::isdigit(mangledName[suffixStart])) {
-      return mangledName.substr(0, pos);
+      // Found "_kernel" followed by digit - this is a potential suffix
+      lastValidPos = pos;
     }
+    searchStart = pos + 1;
+  }
+
+  if (lastValidPos != StringRef::npos) {
+    return mangledName.substr(0, lastValidPos);
   }
   return mangledName;
 }
@@ -812,7 +827,8 @@ static std::string generateKernelArgsHeader(const KernelMetadata &meta) {
 }
 
 /// Generate JSON string for kernel metadata (for runtime dynamic loading)
-static std::string generateMetadataJSON(const KernelMetadata &meta) {
+static std::string generateMetadataJSON(const KernelMetadata &meta,
+                                         const std::vector<unsigned> &originalOrder = {}) {
   std::ostringstream json;
   json << "{\n";
   json << "  \"kernel_name\": \"" << meta.kernelName << "\",\n";
@@ -834,6 +850,19 @@ static std::string generateMetadataJSON(const KernelMetadata &meta) {
 
   json << "  ],\n";
   json << "  \"total_args_size\": " << meta.totalArgsSize << ",\n";
+
+  // Include original argument order mapping if available
+  // This maps from original (hipLaunchKernelGGL) order to device order
+  if (!originalOrder.empty()) {
+    json << "  \"original_arg_order\": [";
+    for (size_t i = 0; i < originalOrder.size(); ++i) {
+      json << originalOrder[i];
+      if (i < originalOrder.size() - 1)
+        json << ", ";
+    }
+    json << "],\n";
+  }
+
   json << "  \"architecture\": \"rv32\"\n";
   json << "}\n";
 
@@ -843,8 +872,10 @@ static std::string generateMetadataJSON(const KernelMetadata &meta) {
 /// Extract metadata from a GPU function and write metadata files
 /// Generates both .meta.json (for runtime) and _args.h (for compile-time)
 /// If outputDir is empty, uses current working directory
+/// Uses pre-built originalArgIsPointer map for computing argument order mapping
 static void emitKernelMetadata(gpu::GPUFuncOp funcOp,
-                                StringRef outputDir) {
+                                StringRef outputDir,
+                                const llvm::StringMap<std::vector<bool>> &originalArgIsPointer) {
   if (!funcOp.isKernel())
     return;
 
@@ -897,6 +928,42 @@ static void emitKernelMetadata(gpu::GPUFuncOp funcOp,
 
   meta.totalArgsSize = offset;
 
+  // Look up pre-computed original argument types from host wrapper
+  // Base name should match the host wrapper function name
+  std::vector<unsigned> originalOrder;
+
+  auto it = originalArgIsPointer.find(baseName);
+  if (it != originalArgIsPointer.end()) {
+    const std::vector<bool> &hostIsPointer = it->second;
+
+    if (hostIsPointer.size() == meta.arguments.size()) {
+      // Build mapping from original order to device order
+      // Device order: scalars first, then pointers (preserving relative order)
+      // Original order: as declared in kernel signature
+
+      // Count scalars in host (original) order
+      unsigned numScalars = 0;
+      for (bool isPtr : hostIsPointer) {
+        if (!isPtr) numScalars++;
+      }
+
+      // Build the mapping: original_arg_order[device_idx] = original_idx
+      originalOrder.resize(hostIsPointer.size());
+      unsigned deviceScalarIdx = 0;
+      unsigned devicePtrIdx = numScalars;
+
+      for (unsigned origIdx = 0; origIdx < hostIsPointer.size(); ++origIdx) {
+        if (!hostIsPointer[origIdx]) {
+          // Scalar - goes to front of device args
+          originalOrder[deviceScalarIdx++] = origIdx;
+        } else {
+          // Pointer - goes to back of device args
+          originalOrder[devicePtrIdx++] = origIdx;
+        }
+      }
+    }
+  }
+
   // Determine output directory
   SmallString<256> outDir;
   if (outputDir.empty()) {
@@ -905,7 +972,7 @@ static void emitKernelMetadata(gpu::GPUFuncOp funcOp,
     outDir = outputDir;
   }
 
-  // Write JSON metadata file
+  // Write JSON metadata file (with original order mapping if available)
   {
     SmallString<256> jsonPath(outDir);
     llvm::sys::path::append(jsonPath, meta.kernelName + ".meta.json");
@@ -916,7 +983,7 @@ static void emitKernelMetadata(gpu::GPUFuncOp funcOp,
       llvm::errs() << "Error writing metadata file " << jsonPath << ": "
                    << ec.message() << "\n";
     } else {
-      outFile << generateMetadataJSON(meta);
+      outFile << generateMetadataJSON(meta, originalOrder);
       outFile.close();
       llvm::outs() << "Wrote kernel metadata: " << jsonPath << "\n";
     }
@@ -958,6 +1025,28 @@ struct ConvertGPUToVortexPass
     MLIRContext *context = &getContext();
     ModuleOp module = getOperation();
 
+    // FIRST: Build argument order map from host wrapper functions BEFORE any changes
+    // This maps kernel base name -> list of (isPointer, type) for original args
+    llvm::StringMap<std::vector<bool>> originalArgIsPointer;
+
+    // Find host wrapper functions (func.func @__polygeist_launch_<name>)
+    for (auto funcOp : module.getOps<func::FuncOp>()) {
+      StringRef funcName = funcOp.getName();
+      if (!funcName.startswith("__polygeist_launch_"))
+        continue;
+
+      // Host wrapper args: user args... + blocks + threads (last 2 are launch params)
+      auto hostArgTypes = funcOp.getArgumentTypes();
+      unsigned numHostUserArgs = hostArgTypes.size() > 2 ? hostArgTypes.size() - 2 : 0;
+
+      std::vector<bool> isPointerVec;
+      for (unsigned i = 0; i < numHostUserArgs; ++i) {
+        isPointerVec.push_back(hostArgTypes[i].isa<MemRefType>() ||
+                               hostArgTypes[i].isa<LLVM::LLVMPointerType>());
+      }
+      originalArgIsPointer[funcName] = std::move(isPointerVec);
+    }
+
     // PREPROCESSING: Consolidate Polygeist auto-tuning artifacts
     // This must happen before any conversion patterns are applied
     consolidatePolygeistAlternatives(module);
@@ -967,10 +1056,11 @@ struct ConvertGPUToVortexPass
     // Files are written to current working directory:
     //   - <kernel_name>.meta.json (for runtime dynamic loading)
     //   - <kernel_name>_args.h (for compile-time type-safe usage)
+    // Pass pre-built argument order map for original argument positions
     module.walk([&](gpu::GPUModuleOp gpuModule) {
       for (auto gpuFunc : gpuModule.getOps<gpu::GPUFuncOp>()) {
         if (gpuFunc.isKernel()) {
-          emitKernelMetadata(gpuFunc, "" /* use current directory */);
+          emitKernelMetadata(gpuFunc, "" /* use current directory */, originalArgIsPointer);
         }
       }
     });