Switch amber to parse clspv's embedded reflection (#899)

* No longer uses the descriptor map, instead parses the embedded
  reflection instructions directly
* strips reflection instructions after parsing to prevent needing
  VK_KHR_non_semantic_info
* update deps

6 files changed, 255 insertions, 97 deletions

diff --git a/DEPS b/DEPS
index c0afb40..a73655e 100644
--- a/DEPS
+++ b/DEPS
@@ -9,18 +9,18 @@ vars = {
   'nlohmann_git': 'https://github.com/nlohmann',
   'swiftshader_git': 'https://swiftshader.googlesource.com',
 
-  'clspv_llvm_revision': '23d1800433d93b94a0c098e97f6f26889b22cb6a',
-  'clspv_revision': '55d518644f61470c421df4229e89abb7a0d5e974',
+  'clspv_llvm_revision': '394db2259575ef3cac8d3d37836b11eb2373c435',
+  'clspv_revision': '86ce19c0130bd13a70862a50a9aa9676eba6548c',
   'cppdap_revision': '1fd23dda91e01550be1a421de307e6fedb2035a9',
   'cpplint_revision': '26470f9ccb354ff2f6d098f831271a1833701b28',
-  'dxc_revision': '174fb21fa97f678e1b76a0dfcf110e59e8de15c3',
+  'dxc_revision': '711e5e2969b3809d7f2254ca17a67d862b359097',
   'glslang_revision': '3ee5f2f1d3316e228916788b300d786bb574d337',
   'googletest_revision': 'a781fe29bcf73003559a3583167fe3d647518464',
   'json_revision': 'fbec662afab55019654e471b65a846a47a696722',
   'lodepng_revision': '34628e89e80cd007179b25b0b2695e6af0f57fac',
   'shaderc_revision': 'ba92b11e1fcaf4c38a64f84d643d6429175bf650',
-  'spirv_headers_revision': '979924c8bc839e4cb1b69d03d48398551f369ce7',
-  'spirv_tools_revision': '9dc1bfa31314731399287690573e2803bec20660',
+  'spirv_headers_revision': '3fdabd0da2932c276b25b9b4a988ba134eba1aa6',
+  'spirv_tools_revision': 'fd3cabd8b5fc43ce83884ac06486c283b9902b4f',
   'swiftshader_revision': '6a8a74986c357b0c6fa0dfd2b4b9230af8d39d1a',
   'vulkan_headers_revision': '83825d55c7d522931124696ecb07ed48f2693e5c',
   'vulkan_loader_revision': 'bfe4f378aee6a2825b8112429cd46529d936babf',
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 7dc0e89..b8dce3d 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -65,7 +65,9 @@ endif()
 
 add_library(libamber ${AMBER_SOURCES})
 amber_default_compile_options(libamber)
-target_include_directories(libamber PRIVATE "${CMAKE_BINARY_DIR}")
+target_include_directories(libamber PRIVATE
+  "${CMAKE_BINARY_DIR}"
+  ${SPIRV-Headers_SOURCE_DIR}/include)
 set_target_properties(libamber PROPERTIES OUTPUT_NAME "amber")
 
 if (${AMBER_ENABLE_DXC})
diff --git a/src/clspv_helper.cc b/src/clspv_helper.cc
index 1db7785..a6c4d50 100644
--- a/src/clspv_helper.cc
+++ b/src/clspv_helper.cc
@@ -14,117 +14,267 @@
 
 #include "src/clspv_helper.h"
 
+#include <unordered_map>
 #include <utility>
 
 #include "clspv/ArgKind.h"
 #include "clspv/Compiler.h"
+#include "clspv/Sampler.h"
+#include "spirv-tools/libspirv.hpp"
+#include "spirv-tools/optimizer.hpp"
+#include "spirv/unified1/NonSemanticClspvReflection.h"
+#include "spirv/unified1/spirv.hpp"
+
+using amber::Pipeline;
+
+namespace {
+
+struct ReflectionHelper {
+  Pipeline::ShaderInfo* shader_info = nullptr;
+  Pipeline* pipeline = nullptr;
+  uint32_t uint_id = 0;
+  std::unordered_map<uint32_t, std::string> strings;
+  std::unordered_map<uint32_t, uint32_t> constants;
+  std::string error_message;
+};
+
+Pipeline::ShaderInfo::DescriptorMapEntry::Kind GetArgKindFromExtInst(
+    uint32_t value) {
+  switch (static_cast<NonSemanticClspvReflectionInstructions>(value)) {
+    case NonSemanticClspvReflectionArgumentStorageBuffer:
+    case NonSemanticClspvReflectionConstantDataStorageBuffer:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SSBO;
+    case NonSemanticClspvReflectionArgumentUniform:
+    case NonSemanticClspvReflectionConstantDataUniform:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::UBO;
+    case NonSemanticClspvReflectionArgumentPodStorageBuffer:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD;
+    case NonSemanticClspvReflectionArgumentPodUniform:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD_UBO;
+    case NonSemanticClspvReflectionArgumentPodPushConstant:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD_PUSHCONSTANT;
+    case NonSemanticClspvReflectionArgumentSampledImage:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::RO_IMAGE;
+    case NonSemanticClspvReflectionArgumentStorageImage:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::WO_IMAGE;
+    case NonSemanticClspvReflectionArgumentSampler:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SAMPLER;
+    case NonSemanticClspvReflectionArgumentWorkgroup:
+    default:
+      return Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SSBO;
+  }
+}
+
+spv_result_t ParseExtendedInst(ReflectionHelper* helper,
+                               const spv_parsed_instruction_t* inst) {
+  auto ext_inst = inst->words[inst->operands[3].offset];
+  switch (ext_inst) {
+    case NonSemanticClspvReflectionKernel: {
+      // Remap the name string to the declaration's result id.
+      const auto& name = helper->strings[inst->words[inst->operands[5].offset]];
+      helper->strings[inst->result_id] = name;
+      break;
+    }
+    case NonSemanticClspvReflectionArgumentInfo: {
+      // Remap the name string to the info's result id.
+      const auto& name = helper->strings[inst->words[inst->operands[4].offset]];
+      helper->strings[inst->result_id] = name;
+      break;
+    }
+    case NonSemanticClspvReflectionArgumentStorageBuffer:
+    case NonSemanticClspvReflectionArgumentUniform:
+    case NonSemanticClspvReflectionArgumentSampledImage:
+    case NonSemanticClspvReflectionArgumentStorageImage:
+    case NonSemanticClspvReflectionArgumentSampler: {
+      // These arguments have descriptor set and binding.
+      auto kernel_id = inst->words[inst->operands[4].offset];
+      auto ordinal_id = inst->words[inst->operands[5].offset];
+      auto ds_id = inst->words[inst->operands[6].offset];
+      auto binding_id = inst->words[inst->operands[7].offset];
+      std::string arg_name;
+      if (inst->num_operands == 9) {
+        arg_name = helper->strings[inst->words[inst->operands[8].offset]];
+      }
+      auto kind = GetArgKindFromExtInst(ext_inst);
+      Pipeline::ShaderInfo::DescriptorMapEntry entry{
+          arg_name,
+          kind,
+          helper->constants[ds_id],
+          helper->constants[binding_id],
+          helper->constants[ordinal_id],
+          /* offset */ 0,
+          /* size */ 0};
+      helper->shader_info->AddDescriptorEntry(helper->strings[kernel_id],
+                                              std::move(entry));
+      break;
+    }
+    case NonSemanticClspvReflectionArgumentPodStorageBuffer:
+    case NonSemanticClspvReflectionArgumentPodUniform: {
+      // These arguments have descriptor set, binding, offset and size.
+      auto kernel_id = inst->words[inst->operands[4].offset];
+      auto ordinal_id = inst->words[inst->operands[5].offset];
+      auto ds_id = inst->words[inst->operands[6].offset];
+      auto binding_id = inst->words[inst->operands[7].offset];
+      auto offset_id = inst->words[inst->operands[8].offset];
+      auto size_id = inst->words[inst->operands[9].offset];
+      std::string arg_name;
+      if (inst->num_operands == 11) {
+        arg_name = helper->strings[inst->words[inst->operands[10].offset]];
+      }
+      auto kind = GetArgKindFromExtInst(ext_inst);
+      Pipeline::ShaderInfo::DescriptorMapEntry entry{
+          arg_name,
+          kind,
+          helper->constants[ds_id],
+          helper->constants[binding_id],
+          helper->constants[ordinal_id],
+          helper->constants[offset_id],
+          helper->constants[size_id]};
+      helper->shader_info->AddDescriptorEntry(helper->strings[kernel_id],
+                                              std::move(entry));
+      break;
+    }
+    case NonSemanticClspvReflectionArgumentPodPushConstant: {
+      // These arguments have offset and size.
+      auto kernel_id = inst->words[inst->operands[4].offset];
+      auto ordinal_id = inst->words[inst->operands[5].offset];
+      auto offset_id = inst->words[inst->operands[6].offset];
+      auto size_id = inst->words[inst->operands[7].offset];
+      std::string arg_name;
+      if (inst->num_operands == 9) {
+        arg_name = helper->strings[inst->words[inst->operands[8].offset]];
+      }
+      auto kind = GetArgKindFromExtInst(ext_inst);
+      Pipeline::ShaderInfo::DescriptorMapEntry entry{
+          arg_name,
+          kind,
+          /* descriptor set */ 0,
+          /* binding */ 0,
+          helper->constants[ordinal_id],
+          helper->constants[offset_id],
+          helper->constants[size_id]};
+      helper->shader_info->AddDescriptorEntry(helper->strings[kernel_id],
+                                              std::move(entry));
+      break;
+    }
+    case NonSemanticClspvReflectionArgumentWorkgroup:
+      helper->error_message = "Workgroup arguments are not currently supported";
+      return SPV_ERROR_INVALID_DATA;
+    case NonSemanticClspvReflectionConstantDataStorageBuffer:
+    case NonSemanticClspvReflectionConstantDataUniform:
+      helper->error_message =
+          "Constant descriptor entries are not currently supported";
+      return SPV_ERROR_INVALID_DATA;
+    case NonSemanticClspvReflectionSpecConstantWorkgroupSize:
+    case NonSemanticClspvReflectionSpecConstantGlobalOffset:
+    case NonSemanticClspvReflectionSpecConstantWorkDim:
+      // Nothing to do. Amber currently requires script authors to know the
+      // spec ids and use them directly.
+      break;
+    case NonSemanticClspvReflectionPushConstantGlobalOffset: {
+      auto offset_id = inst->words[inst->operands[4].offset];
+      auto size_id = inst->words[inst->operands[5].offset];
+      Pipeline::ShaderInfo::PushConstant push_constant{
+          Pipeline::ShaderInfo::PushConstant::PushConstantType::kGlobalOffset,
+          helper->constants[offset_id], helper->constants[size_id]};
+      helper->shader_info->AddPushConstant(std::move(push_constant));
+      break;
+    }
+    case NonSemanticClspvReflectionPushConstantEnqueuedLocalSize:
+    case NonSemanticClspvReflectionPushConstantGlobalSize:
+    case NonSemanticClspvReflectionPushConstantRegionOffset:
+    case NonSemanticClspvReflectionPushConstantNumWorkgroups:
+    case NonSemanticClspvReflectionPushConstantRegionGroupOffset:
+      helper->error_message = "Unsupported push constant";
+      return SPV_ERROR_INVALID_DATA;
+    case NonSemanticClspvReflectionLiteralSampler: {
+      auto ds_id = inst->words[inst->operands[4].offset];
+      auto binding_id = inst->words[inst->operands[5].offset];
+      auto mask_id = inst->words[inst->operands[6].offset];
+      helper->pipeline->AddSampler(helper->constants[mask_id],
+                                   helper->constants[ds_id],
+                                   helper->constants[binding_id]);
+      break;
+    } break;
+  }
+
+  return SPV_SUCCESS;
+}
+
+spv_result_t ParseReflection(void* user_data,
+                             const spv_parsed_instruction_t* inst) {
+  auto* helper = reinterpret_cast<ReflectionHelper*>(user_data);
+  switch (inst->opcode) {
+    case spv::OpTypeInt:
+      if (inst->words[inst->operands[1].offset] == 32 &&
+          inst->words[inst->operands[2].offset] == 0) {
+        // Track the result id of OpTypeInt 32 0.
+        helper->uint_id = inst->result_id;
+      }
+      break;
+    case spv::OpConstant:
+      if (inst->words[inst->operands[0].offset] == helper->uint_id) {
+        // Record the values for all uint32_t constants.
+        uint32_t value = inst->words[inst->operands[2].offset];
+        helper->constants[inst->result_id] = value;
+      }
+      break;
+    case spv::OpString: {
+      // Track all strings.
+      std::string value =
+          reinterpret_cast<const char*>(inst->words + inst->operands[1].offset);
+      helper->strings[inst->result_id] = value;
+      break;
+    }
+    case spv::OpExtInst:
+      if (inst->ext_inst_type ==
+          SPV_EXT_INST_TYPE_NONSEMANTIC_CLSPVREFLECTION) {
+        return ParseExtendedInst(helper, inst);
+      }
+      break;
+  }
+  return SPV_SUCCESS;
+}
+
+}  // anonymous namespace
 
 namespace amber {
 namespace clspvhelper {
 
 Result Compile(Pipeline::ShaderInfo* shader_info,
                Pipeline* pipeline,
+               spv_target_env env,
                std::vector<uint32_t>* generated_binary) {
-  std::vector<clspv::version0::DescriptorMapEntry> entries;
   const auto& src_str = shader_info->GetShader()->GetData();
   std::string options;
   for (const auto& option : shader_info->GetCompileOptions()) {
     options += option + " ";
   }
   if (clspv::CompileFromSourceString(src_str, /* sampler map */ "", options,
-                                     generated_binary, &entries)) {
+                                     generated_binary)) {
     return Result("Clspv compile failed");
   }
 
-  for (auto& entry : entries) {
-    if (entry.kind == clspv::version0::DescriptorMapEntry::Constant) {
-      return Result("Constant descriptor entries are not currently supported");
-    }
-
-    if (entry.kind == clspv::version0::DescriptorMapEntry::KernelArg) {
-      Pipeline::ShaderInfo::DescriptorMapEntry descriptor_entry;
-      descriptor_entry.descriptor_set = entry.descriptor_set;
-      descriptor_entry.binding = entry.binding;
-      descriptor_entry.pod_offset = 0;
-      descriptor_entry.pod_arg_size = 0;
-      switch (entry.kernel_arg_data.arg_kind) {
-        case clspv::ArgKind::Buffer:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SSBO;
-          break;
-        case clspv::ArgKind::BufferUBO:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::UBO;
-          break;
-        case clspv::ArgKind::Pod:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD;
-          break;
-        case clspv::ArgKind::PodUBO:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD_UBO;
-          break;
-        case clspv::ArgKind::PodPushConstant:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::POD_PUSHCONSTANT;
-          break;
-        case clspv::ArgKind::ReadOnlyImage:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::RO_IMAGE;
-          break;
-        case clspv::ArgKind::WriteOnlyImage:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::WO_IMAGE;
-          break;
-        case clspv::ArgKind::Sampler:
-          descriptor_entry.kind =
-              Pipeline::ShaderInfo::DescriptorMapEntry::Kind::SAMPLER;
-          break;
-        case clspv::ArgKind::Local:
-          // Local arguments are handled via specialization constants.
-          break;
-        default:
-          return Result("Unsupported kernel argument descriptor entry");
-      }
-
-      if (entry.kernel_arg_data.arg_kind == clspv::ArgKind::Pod ||
-          entry.kernel_arg_data.arg_kind == clspv::ArgKind::PodUBO ||
-          entry.kernel_arg_data.arg_kind == clspv::ArgKind::PodPushConstant) {
-        descriptor_entry.pod_offset = entry.kernel_arg_data.pod_offset;
-        descriptor_entry.pod_arg_size = entry.kernel_arg_data.pod_arg_size;
-      }
+  // Parse the reflection instructions.
+  ReflectionHelper helper;
+  helper.shader_info = shader_info;
+  helper.pipeline = pipeline;
+  spv_context context(spvContextCreate(env));
+  if (spvBinaryParse(context, &helper, generated_binary->data(),
+                     generated_binary->size(), nullptr, ParseReflection,
+                     nullptr)) {
+    return Result(helper.error_message);
+  }
 
-      descriptor_entry.arg_name = entry.kernel_arg_data.arg_name;
-      descriptor_entry.arg_ordinal = entry.kernel_arg_data.arg_ordinal;
-
-      shader_info->AddDescriptorEntry(entry.kernel_arg_data.kernel_name,
-                                      std::move(descriptor_entry));
-    } else if (entry.kind ==
-               clspv::version0::DescriptorMapEntry::PushConstant) {
-      Pipeline::ShaderInfo::PushConstant push_constant;
-      switch (entry.push_constant_data.pc) {
-        case clspv::PushConstant::Dimensions:
-          push_constant.type =
-              Pipeline::ShaderInfo::PushConstant::PushConstantType::kDimensions;
-          break;
-        case clspv::PushConstant::GlobalOffset:
-          push_constant.type = Pipeline::ShaderInfo::PushConstant::
-              PushConstantType::kGlobalOffset;
-          break;
-        default:
-          return Result("unhandled push constant generated by Clspv");
-      }
-      push_constant.offset = entry.push_constant_data.offset;
-      push_constant.size = entry.push_constant_data.size;
-      shader_info->AddPushConstant(std::move(push_constant));
-    } else if (entry.kind == clspv::version0::DescriptorMapEntry::Sampler) {
-      // Create a new sampler info.
-      pipeline->AddSampler(entry.sampler_data.mask, entry.descriptor_set,
-                           entry.binding);
-    } else if (entry.kind == clspv::version0::DescriptorMapEntry::KernelDecl) {
-      // Nothing to do for declarations.
-    }
+  // Strip the reflection instructions to avoid requiring the implementation to
+  // support VK_KHR_shader_non_semantic_info.
+  spvtools::Optimizer opt(env);
+  opt.RegisterPass(spvtools::CreateStripReflectInfoPass());
+  std::vector<uint32_t> stripped;
+  if (!opt.Run(generated_binary->data(), generated_binary->size(), &stripped)) {
+    return Result("failed to strip reflection instructions");
   }
+  generated_binary->swap(stripped);
 
   return Result();
 }
diff --git a/src/clspv_helper.h b/src/clspv_helper.h
index 82d8db3..3a06c74 100644
--- a/src/clspv_helper.h
+++ b/src/clspv_helper.h
@@ -19,6 +19,7 @@
 #include <vector>
 
 #include "amber/result.h"
+#include "spirv-tools/libspirv.h"
 #include "src/pipeline.h"
 
 namespace amber {
@@ -28,6 +29,7 @@ namespace clspvhelper {
 // Returns the generated SPIR-V binary via |generated_binary| argument.
 Result Compile(Pipeline::ShaderInfo* shader_info,
                Pipeline* pipeline,
+               spv_target_env env,
                std::vector<uint32_t>* generated_binary);
 
 }  // namespace clspvhelper
diff --git a/src/shader_compiler.cc b/src/shader_compiler.cc
index 4eb328b..b0bb44a 100644
--- a/src/shader_compiler.cc
+++ b/src/shader_compiler.cc
@@ -148,7 +148,7 @@ std::pair<Result, std::vector<uint32_t>> ShaderCompiler::Compile(
 
 #if AMBER_ENABLE_CLSPV
   } else if (shader->GetFormat() == kShaderFormatOpenCLC) {
-    Result r = CompileOpenCLC(shader_info, pipeline, &results);
+    Result r = CompileOpenCLC(shader_info, pipeline, target_env, &results);
     if (!r.IsSuccess())
       return {r, {}};
 #endif  // AMBER_ENABLE_CLSPV
@@ -288,12 +288,14 @@ Result ShaderCompiler::CompileHlsl(const Shader*,
 #if AMBER_ENABLE_CLSPV
 Result ShaderCompiler::CompileOpenCLC(Pipeline::ShaderInfo* shader_info,
                                       Pipeline* pipeline,
+                                      spv_target_env env,
                                       std::vector<uint32_t>* result) const {
-  return clspvhelper::Compile(shader_info, pipeline, result);
+  return clspvhelper::Compile(shader_info, pipeline, env, result);
 }
 #else
 Result ShaderCompiler::CompileOpenCLC(Pipeline::ShaderInfo*,
                                       Pipeline*,
+                                      spv_target_env,
                                       std::vector<uint32_t>*) const {
   return {};
 }
diff --git a/src/shader_compiler.h b/src/shader_compiler.h
index e4caa62..e5c6ee2 100644
--- a/src/shader_compiler.h
+++ b/src/shader_compiler.h
@@ -21,6 +21,7 @@
 
 #include "amber/amber.h"
 #include "amber/result.h"
+#include "spirv-tools/libspirv.h"
 #include "src/pipeline.h"
 #include "src/shader.h"
 #include "src/virtual_file_store.h"
@@ -60,6 +61,7 @@ class ShaderCompiler {
   Result CompileHlsl(const Shader* shader, std::vector<uint32_t>* result) const;
   Result CompileOpenCLC(Pipeline::ShaderInfo* shader,
                         Pipeline* pipeline,
+                        spv_target_env env,
                         std::vector<uint32_t>* result) const;
 
   std::string spv_env_;