From 21893763c18d323af56d30db1bf878abbbdb9060 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Thu, 30 Jan 2025 15:05:16 -0800
Subject: Declare `VariablePointers` capability when needed. (#6231)

---
 source/slang/slang-emit-spirv.cpp  | 19 +++++++++++++++++++
 tests/spirv/variable-pointer.slang | 25 +++++++++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 tests/spirv/variable-pointer.slang

diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 2c36ae5f7..4a66cb4ef 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -5647,6 +5647,8 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
 
     SpvInst* emitPhi(SpvInstParent* parent, IRParam* inst)
     {
+        requireVariableBufferCapabilityIfNeeded(inst->getDataType());
+
         // An `IRParam` in an ordinary `IRBlock` represents a phi value.
         // We can translate them directly to SPIRV's `Phi` instruction.
         // In order to do that, we need to figure out the source values
@@ -5721,6 +5723,7 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
 
         // Does this function declare any requirements.
         handleRequiredCapabilities(funcValue);
+        requireVariableBufferCapabilityIfNeeded(inst->getDataType());
 
         // We want to detect any call to an intrinsic operation, and inline
         // the SPIRV snippet directly at the call site.
@@ -6113,6 +6116,8 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
 
     SpvInst* emitGetElement(SpvInstParent* parent, IRGetElement* inst)
     {
+        requireVariableBufferCapabilityIfNeeded(inst->getDataType());
+
         // Note: SPIRV only supports the case where `index` is constant.
         auto base = inst->getBase();
         const auto baseTy = base->getDataType();
@@ -6149,6 +6154,8 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
 
     SpvInst* emitLoad(SpvInstParent* parent, IRLoad* inst)
     {
+        requireVariableBufferCapabilityIfNeeded(inst->getDataType());
+
         auto ptrType = as<IRPtrTypeBase>(inst->getPtr()->getDataType());
         if (ptrType && addressSpaceToStorageClass(ptrType->getAddressSpace()) ==
                            SpvStorageClassPhysicalStorageBuffer)
@@ -8100,6 +8107,18 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
         }
     }
 
+    void requireVariableBufferCapabilityIfNeeded(IRInst* type)
+    {
+        if (auto ptrType = as<IRPtrTypeBase>(type))
+        {
+            if (ptrType->getAddressSpace() == AddressSpace::StorageBuffer)
+            {
+                ensureExtensionDeclaration(UnownedStringSlice("SPV_KHR_variable_pointers"));
+                requireSPIRVCapability(SpvCapabilityVariablePointers);
+            }
+        }
+    }
+
     // https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpExecutionMode
     Dictionary<SpvWord, OrderedHashSet<SpvExecutionMode>> m_executionModes;
     template<typename... Operands>
diff --git a/tests/spirv/variable-pointer.slang b/tests/spirv/variable-pointer.slang
new file mode 100644
index 000000000..47df42ce4
--- /dev/null
+++ b/tests/spirv/variable-pointer.slang
@@ -0,0 +1,25 @@
+//TEST:SIMPLE(filecheck=SPIRV): -target spirv
+
+//DISABLED_TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK): -vk -output-using-type
+
+// SPIRV: OpCapability VariablePointers
+
+//TEST_INPUT:set buffer1 = ubuffer(data=[1.0], stride=4)
+StructuredBuffer<float> buffer1;
+//TEST_INPUT:set buffer2 = ubuffer(data=[2.0], stride=4)
+StructuredBuffer<float> buffer2;
+//TEST_INPUT:set output = out ubuffer(data=[0.0], stride=4)
+RWStructuredBuffer<float> output;
+
+[numthreads(2,1,1)]
+void computeMain(int id : SV_DispatchThreadID)
+{
+    StructuredBuffer<float> buffer;
+    if (id == 0)
+        buffer = buffer1;
+    else
+        buffer = buffer2;
+    output[id] = buffer[0];
+    // CHECK: 1.0
+    // CHECK: 2.0
+}
\ No newline at end of file
-- 
cgit v1.2.3