[llvm] [SPIR-V] Handle spirv.MemoryModel metadata (PR #186138)

[Arseniy Obolenskiy via llvm-commits]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/186138

>From dde8548bcb15f33c4a707a005a91a999114d0f44 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Thu, 12 Mar 2026 16:09:58 +0100
Subject: [PATCH] [SPIR-V] Handle Shader attr info from spirv.MemoryModel
 metadata

---
 llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp |  8 +++++-
 llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp      | 25 +++++++++++++++++++
 .../CodeGen/SPIRV/memory-model-md-glsl450.ll  | 14 +++++++++++
 .../CodeGen/SPIRV/memory-model-md-shader.ll   | 14 +++++++++++
 .../CodeGen/SPIRV/memory-model-md-vulkan.ll   | 14 +++++++++++
 5 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/SPIRV/memory-model-md-glsl450.ll
 create mode 100644 llvm/test/CodeGen/SPIRV/memory-model-md-shader.ll
 create mode 100644 llvm/test/CodeGen/SPIRV/memory-model-md-vulkan.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
index 38bc72ff154fa..091249c7f62d3 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
@@ -215,6 +215,10 @@ void SPIRVModuleAnalysis::setBaseInfo(const Module &M) {
   MAI.Reqs.getAndAddRequirements(SPIRV::OperandCategory::AddressingModelOperand,
                                  MAI.Addr, *ST);
 
+  if (MAI.Mem == SPIRV::MemoryModel::VulkanKHR)
+    MAI.Reqs.addExtension(
+        SPIRV::Extension::SPV_KHR_vulkan_memory_model);
+
   if (!ST->isShader()) {
     // TODO: check if it's required by default.
     MAI.ExtInstSetMap[static_cast<unsigned>(
@@ -940,7 +944,9 @@ void RequirementHandler::initAvailableCapabilitiesForVulkan(
                     Capability::StorageBufferArrayDynamicIndexing,
                     Capability::StorageImageArrayDynamicIndexing,
                     Capability::DerivativeControl, Capability::MinLod,
-                    Capability::ImageGatherExtended});
+                    Capability::ImageGatherExtended,
+                    Capability::Addresses,
+                    Capability::VulkanMemoryModelKHR});
 
   // Became core in Vulkan 1.2
   if (ST.isAtLeastSPIRVVer(VersionTuple(1, 5))) {
diff --git a/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp b/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp
index d65d8ec53c6d0..cb014002c1f23 100644
--- a/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "SPIRVSubtarget.h"
+#include "MCTargetDesc/SPIRVBaseInfo.h"
 #include "SPIRV.h"
 #include "SPIRVCommandLine.h"
 #include "SPIRVGlobalRegistry.h"
@@ -207,6 +208,30 @@ void SPIRVSubtarget::resolveEnvFromModule(const Module &M) {
     }
   }
 
+  if (!HasShaderAttr) {
+    if (auto *MemModel = M.getNamedMetadata("spirv.MemoryModel")) {
+      assert(MemModel->getNumOperands() > 0 && "Invalid spirv.MemoryModel");
+      auto *MemMD = MemModel->getOperand(0);
+      assert(MemMD->getNumOperands() > 1 &&
+             "Invalid spirv.MemoryModel operand");
+      unsigned MemModelVal =
+          mdconst::extract<ConstantInt>(MemMD->getOperand(1))->getZExtValue();
+      switch (MemModelVal) {
+      case SPIRV::MemoryModel::Simple:
+      case SPIRV::MemoryModel::GLSL450:
+        HasShaderAttr = true;
+        break;
+      case SPIRV::MemoryModel::VulkanKHR:
+        HasShaderAttr = true;
+        AvailableExtensions.insert(
+            SPIRV::Extension::SPV_KHR_vulkan_memory_model);
+        break;
+      case SPIRV::MemoryModel::OpenCL:
+        break;
+      }
+    }
+  }
+
   setEnv(HasShaderAttr ? Shader : Kernel);
 }
 
diff --git a/llvm/test/CodeGen/SPIRV/memory-model-md-glsl450.ll b/llvm/test/CodeGen/SPIRV/memory-model-md-glsl450.ll
new file mode 100644
index 0000000000000..9d2ce1128a43c
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/memory-model-md-glsl450.ll
@@ -0,0 +1,14 @@
+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpMemoryModel Logical GLSL450
+define void @main() {
+entry:
+  ret void
+}
+
+; AddressingModel=Logical (0), MemoryModel=GLSL450 (1)
+!spirv.MemoryModel = !{!0}
+!0 = !{i32 0, i32 1}
diff --git a/llvm/test/CodeGen/SPIRV/memory-model-md-shader.ll b/llvm/test/CodeGen/SPIRV/memory-model-md-shader.ll
new file mode 100644
index 0000000000000..d40fba3acd382
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/memory-model-md-shader.ll
@@ -0,0 +1,14 @@
+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpMemoryModel Physical32 Simple
+define void @main() {
+entry:
+  ret void
+}
+
+; AddressingModel=Physical32 (1), MemoryModel=Simple (0)
+!spirv.MemoryModel = !{!0}
+!0 = !{i32 1, i32 0}
diff --git a/llvm/test/CodeGen/SPIRV/memory-model-md-vulkan.ll b/llvm/test/CodeGen/SPIRV/memory-model-md-vulkan.ll
new file mode 100644
index 0000000000000..cbee31481ea42
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/memory-model-md-vulkan.ll
@@ -0,0 +1,14 @@
+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpMemoryModel Logical Vulkan
+define void @main() {
+entry:
+  ret void
+}
+
+; AddressingModel=Logical (0), MemoryModel=VulkanKHR (3)
+!spirv.MemoryModel = !{!0}
+!0 = !{i32 0, i32 3}