[llvm] [mlir] [SPIR-V] Preserve module analysis for XeVM lowering (PR #188971)

[Aviral Garg via llvm-commits]

https://github.com/aviralgarg05 updated https://github.com/llvm/llvm-project/pull/188971

>From 09f6363bb7239747541e4a60d024392e7e65377e Mon Sep 17 00:00:00 2001
From: aviralgarg05 <gargaviral99 at gmail.com>
Date: Fri, 27 Mar 2026 16:30:48 +0530
Subject: [PATCH] [SPIR-V] Preserve module analysis for XeVM lowering

---
 llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp     |  9 ++++++++-
 llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp | 16 ++++++++++++++++
 llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h   |  4 ++++
 ...ower-to-xevm-declaration-only-kernels.mlir | 19 +++++++++++++++++++
 mlir/test/lit.cfg.py                          |  3 +++
 5 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/GPU/gpu-lower-to-xevm-declaration-only-kernels.mlir

diff --git a/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp b/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
index e1fceea086287..2225984bef3fc 100644
--- a/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
@@ -140,6 +140,7 @@ void SPIRVAsmPrinter::emitEndOfAsmFile(Module &M) {
 // Any cleanup actions with the Module after we don't care about its content
 // anymore.
 void SPIRVAsmPrinter::cleanUp(Module &M) {
+  clearCachedSPIRVModuleAnalysis(M);
   // Verifier disallows uses of intrinsic global variables.
   for (StringRef GVName :
        {"llvm.global_ctors", "llvm.global_dtors", "llvm.used"}) {
@@ -819,7 +820,12 @@ void SPIRVAsmPrinter::outputModuleSections() {
   // Get the global subtarget to output module-level info.
   ST = static_cast<const SPIRVTargetMachine &>(TM).getSubtargetImpl();
   TII = ST->getInstrInfo();
-  MAI = &getAnalysis<SPIRVModuleAnalysis>().MAI;
+  if (!MAI) {
+    if (auto *ModuleAnalysis = getAnalysisIfAvailable<SPIRVModuleAnalysis>())
+      MAI = &ModuleAnalysis->MAI;
+    else if (M)
+      MAI = getCachedSPIRVModuleAnalysis(*M);
+  }
   assert(ST && TII && MAI && M && "Module analysis is required");
   // Output instructions according to the Logical Layout of a Module:
   // 1,2. All OpCapability instructions, then optional OpExtension
@@ -861,6 +867,7 @@ void SPIRVAsmPrinter::outputModuleSections() {
 
 bool SPIRVAsmPrinter::doInitialization(Module &M) {
   ModuleSectionsEmitted = false;
+  MAI = nullptr;
   // We need to call the parent's one explicitly.
   return AsmPrinter::doInitialization(M);
 }
diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
index 6bd2eb552a936..19e4d0e9ce96c 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
@@ -28,11 +28,17 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
+#include <memory>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "spirv-module-analysis"
 
+namespace {
+DenseMap<const Module *, std::unique_ptr<SPIRV::ModuleAnalysisInfo>>
+    CachedModuleAnalyses;
+} // namespace
+
 static cl::opt<bool>
     SPVDumpDeps("spv-dump-deps",
                 cl::desc("Dump MIR with SPIR-V dependencies info"),
@@ -2910,6 +2916,16 @@ bool SPIRVModuleAnalysis::runOnModule(Module &M) {
 
   // Set maximum ID used.
   GR->setBound(MAI.MaxID);
+  CachedModuleAnalyses[&M] = std::make_unique<SPIRV::ModuleAnalysisInfo>(MAI);
 
   return false;
 }
+
+SPIRV::ModuleAnalysisInfo *llvm::getCachedSPIRVModuleAnalysis(const Module &M) {
+  auto It = CachedModuleAnalyses.find(&M);
+  return It == CachedModuleAnalyses.end() ? nullptr : It->second.get();
+}
+
+void llvm::clearCachedSPIRVModuleAnalysis(const Module &M) {
+  CachedModuleAnalyses.erase(&M);
+}
diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h
index b2efd148edf07..7c84161824e7e 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h
@@ -25,6 +25,7 @@ namespace llvm {
 class SPIRVSubtarget;
 class MachineFunction;
 class MachineModuleInfo;
+class Module;
 
 namespace SPIRV {
 // The enum contains logical module sections for the instruction collection.
@@ -258,5 +259,8 @@ struct SPIRVModuleAnalysis : public ModulePass {
   const SPIRVInstrInfo *TII;
   MachineModuleInfo *MMI;
 };
+
+SPIRV::ModuleAnalysisInfo *getCachedSPIRVModuleAnalysis(const Module &M);
+void clearCachedSPIRVModuleAnalysis(const Module &M);
 } // namespace llvm
 #endif // LLVM_LIB_TARGET_SPIRV_SPIRVMODULEANALYSIS_H
diff --git a/mlir/test/Dialect/GPU/gpu-lower-to-xevm-declaration-only-kernels.mlir b/mlir/test/Dialect/GPU/gpu-lower-to-xevm-declaration-only-kernels.mlir
new file mode 100644
index 0000000000000..7e33dd80024f5
--- /dev/null
+++ b/mlir/test/Dialect/GPU/gpu-lower-to-xevm-declaration-only-kernels.mlir
@@ -0,0 +1,19 @@
+// REQUIRES: xevm-conversions
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline | FileCheck %s
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: module attributes {gpu.container_module} {
+  // CHECK: gpu.binary @kernels
+  // CHECK-NOT: gpu.module @kernels
+  gpu.module @kernels {
+    func.func private @__hipblaslt_init_kernel(
+        %arg0: memref<?xf32, 1>, %arg1: index, %arg2: index, %arg3: index,
+        %arg4: index, %arg5: index) attributes {gpu.kernel}
+    func.func private @__hipblaslt_init_small_kernel(
+        %arg0: memref<?xf32, 1>, %arg1: index, %arg2: index, %arg3: index,
+        %arg4: index, %arg5: index) attributes {gpu.kernel}
+    func.func private @__hipblaslt_init_nan_tri_kernel(
+        %arg0: memref<?xf32, 1>, %arg1: index, %arg2: index, %arg3: index,
+        %arg4: index, %arg5: index, %arg6: i1) attributes {gpu.kernel}
+  }
+}
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index a716ba0adb480..6dae6def29435 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -60,6 +60,9 @@
 config.substitutions.append(("%host_cxx", config.host_cxx.strip()))
 config.substitutions.append(("%host_cc", config.host_cc.strip()))
 
+if config.run_xevm_tests:
+    config.available_features.add("xevm-conversions")
+
 
 # Searches for a runtime library with the given name and returns the found path.
 # Correctly handles the platforms shared library directory and naming conventions.