Add specific handling for inline spirv pointer types (#6873)

Add specific handling for inline spirv pointer types
    
We currently blindly create a new type for all vk::SpirvType* types.
This can cause problems when the type is suppose to match another type.
In this case, we want a spirv pointer type to match the pointer type of
the
variable. The OpTypePointer for the variable is implicitly created when
declaring the variable, but, if we want to explicitly declare the
pointer as a SpirvType, we get two different OpTypePointer instructions
in the module. So technically the types do not match.
    
To fix this, I add special handling in the SPIR-V backend to be able to
merge the implicit pointer type created and the SpirvType when they
match.

This implements the SPIR-V Pointers in [HLSL spec proposal
0021](https://github.com/microsoft/hlsl-specs/blob/main/proposals/0021-vk-coop-matrix.md#spir-v-pointers)

---------

Co-authored-by: Nathan Gauër <github@keenuts.net>

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.