Add specific handling for inline spirv pointer types

tools/clang/lib/Headers/hlsl/vk/spirv.h

@@ -0,0 +1,34 @@
+// Copyright (c) 2024 Google LLC
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef _HLSL_VK_SPIRV_H_
+#define _HLSL_VK_SPIRV_H_
+
+namespace vk {
+
+enum StorageClass {
+  StorageClassWorkgroup = 4,
+};
+
+// An opaque type to represent a Spir-V pointer to the workgroup storage class.
+// clang-format off
+template <typename PointeeType>
+using WorkgroupSpirvPointer = const vk::SpirvOpaqueType<
+    /* OpTypePointer */ 32,
+    vk::Literal<vk::integral_constant<uint, StorageClassWorkgroup> >,
+    PointeeType>;
+// clang-format on
+
+// Returns an opaque Spir-V pointer to v. The memory object v's storage class
+// modifier must be groupshared. If the incorrect storage class is used, then
+// there will be a validation error, and it will not show the correct
+template <typename T>
+[[vk::ext_instruction(/* OpCopyObject */ 83)]] WorkgroupSpirvPointer<T>
+GetGroupSharedAddress([[vk::ext_reference]] T v);
+
+} // namespace vk
+
+#endif // _HLSL_VK_SPIRV_H_

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -732,6 +732,15 @@ const SpirvType *LowerTypeVisitor::lowerInlineSpirvType(
 
   auto args = specDecl->getTemplateArgs()[operandsIndex].getPackAsArray();
 
+  if (operandsIndex == 1 && args.size() == 2 &&
+      static_cast<spv::Op>(opcode) == spv::Op::OpTypePointer) {
+    const SpirvType *result =
+        getSpirvPointerFromInlineSpirvType(args, rule, isRowMajor, srcLoc);
+    if (result) {
+      return result;
+    }
+  }
+
   for (TemplateArgument arg : args) {
     switch (arg.getKind()) {
     case TemplateArgument::ArgKind::Type: {
@@ -1363,5 +1372,41 @@ LowerTypeVisitor::populateLayoutInformation(
   return result;
 }
 
+const SpirvType *LowerTypeVisitor::getSpirvPointerFromInlineSpirvType(
+    ArrayRef<TemplateArgument> args, SpirvLayoutRule rule,
+    Optional<bool> isRowMajor, SourceLocation location) {
+
+  assert(args.size() == 2 && "OpTypePointer requires exactly 2 arguments.");
+  QualType scLiteralType = args[0].getAsType();
+  SpirvConstant *constant = nullptr;
+  if (!getVkIntegralConstantValue(scLiteralType, constant, location) ||
+      !constant) {
+    return nullptr;
+  }
+  if (!constant->isLiteral())
+    return nullptr;
+
+  auto *intConstant = dyn_cast<SpirvConstantInteger>(constant);
+  if (!intConstant) {
+    return nullptr;
+  }
+
+  visitInstruction(constant);
+  spv::StorageClass storageClass =
+      static_cast<spv::StorageClass>(intConstant->getValue().getLimitedValue());
+
+  QualType pointeeType;
+  if (args[1].getKind() == TemplateArgument::ArgKind::Type) {
+    pointeeType = args[1].getAsType();
+  } else {
+    TemplateName templateName = args[1].getAsTemplate();
+    pointeeType = createASTTypeFromTemplateName(templateName);
+  }
+
+  const SpirvType *pointeeSpirvType =
+      lowerType(pointeeType, rule, isRowMajor, location);
+  return spvContext.getPointerType(pointeeSpirvType, storageClass);
+}
+
 } // namespace spirv
 } // namespace clang

tools/clang/lib/SPIRV/LowerTypeVisitor.h

@@ -124,6 +124,13 @@ class LowerTypeVisitor : public Visitor {
                                    SpirvLayoutRule rule,
                                    const uint32_t fieldIndex);
 
+  /// Get a lowered SpirvPointer from the args to a SpirvOpaqueType.
+  /// The pointer will use the given layout rule. `isRowMajor` is used to
+  /// lower the pointee type.
+  const SpirvType *getSpirvPointerFromInlineSpirvType(
+      ArrayRef<TemplateArgument> args, SpirvLayoutRule rule,
+      Optional<bool> isRowMajor, SourceLocation location);
+
 private:
   ASTContext &astContext;                /// AST context
   SpirvContext &spvContext;              /// SPIR-V context

tools/clang/test/CMakeLists.txt

@@ -9,6 +9,8 @@ endif ()
 
 string(REPLACE ${CMAKE_CFG_INTDIR} ${LLVM_BUILD_MODE} CLANG_TOOLS_DIR ${LLVM_RUNTIME_OUTPUT_INTDIR})
 
+set(HLSL_HEADERS_DIR ${LLVM_SOURCE_DIR}/tools/clang/lib/Headers/hlsl) # HLSL Change
+
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.in
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg

tools/clang/test/CodeGenSPIRV/workgroupspirvpointer.hlsl

@@ -0,0 +1,26 @@
+// RUN: dxc -fspv-target-env=vulkan1.3 -T cs_6_0 -E main -spirv -HV 2021 -I %hlsl_headers %s 2>&1 | FileCheck %s
+
+#include "vk/spirv.h"
+
+// CHECK-NOT: OpCapability VariablePointers
+
+RWStructuredBuffer<int> data;
+
+groupshared int shared_data[64];
+
+[[vk::ext_instruction(/* OpLoad */ 61)]] int
+Load(vk::WorkgroupSpirvPointer<int> p);
+
+int foo(vk::WorkgroupSpirvPointer<int> param) {
+  return Load(param);
+}
+
+[numthreads(64, 1, 1)] void main() {
+// CHECK: [[ac:%[0-9]+]] = OpAccessChain %_ptr_Workgroup_int %shared_data %int_0
+// CHECK: [[ld:%[0-9]+]] = OpLoad %int [[ac]]
+// CHECK: [[ac:%[0-9]+]] = OpAccessChain %_ptr_StorageBuffer_int %data %int_0 %uint_0
+// CHECK: OpStore [[ac]] [[ld]]
+
+  vk::WorkgroupSpirvPointer<int> p = vk::GetGroupSharedAddress(shared_data[0]);
+  data[0] = foo(p);
+}

tools/clang/test/CodeGenSPIRV/workgroupspirvpointer.varpointer.hlsl

@@ -0,0 +1,23 @@
+// RUN: dxc -fspv-target-env=vulkan1.3 -T cs_6_0 -E main -spirv -HV 2021 -I %hlsl_headers %s 2>&1 | FileCheck %s
+
+#include "vk/spirv.h"
+
+// CHECK: OpCapability VariablePointers
+
+RWStructuredBuffer<int> data;
+
+groupshared int shared_data[64];
+
+[[vk::ext_instruction(/* OpLoad */ 61)]] int
+Load(vk::WorkgroupSpirvPointer<int> p);
+
+[[noinline]]
+int foo(vk::WorkgroupSpirvPointer<int> param) {
+  return Load(param);
+}
+
+[[vk::ext_capability(/* VariablePointersCapability */ 4442)]]
+[numthreads(64, 1, 1)] void main() {
+  vk::WorkgroupSpirvPointer<int> p = vk::GetGroupSharedAddress(shared_data[0]);
+  data[0] = foo(p);
+}

tools/clang/test/lit.cfg

@@ -210,6 +210,7 @@ if has_plugins and config.llvm_plugin_ext:
 
 config.substitutions.append( ('%llvmshlibdir', config.llvm_shlib_dir) )
 config.substitutions.append( ('%pluginext', config.llvm_plugin_ext) )
+config.substitutions.append( ('%hlsl_headers', config.hlsl_headers_dir) ) #HLSL change
 
 if config.clang_examples:
     config.available_features.add('examples')

tools/clang/test/lit.site.cfg.in

@@ -21,6 +21,7 @@ config.enable_shared = @ENABLE_SHARED@
 config.enable_backtrace = "@ENABLE_BACKTRACES@"
 config.host_arch = "@HOST_ARCH@"
 config.spirv = "@ENABLE_SPIRV_CODEGEN@" =="ON"
+config.hlsl_headers_dir = "@HLSL_HEADERS_DIR@" # HLSL change
 
 # Support substitution of the tools and libs dirs with user parameters. This is
 # used when we can't determine the tool dir at configuration time.