Added option --preserve-numeric-ids to tools/spirv-as

Added options to tools::Assemble

include/spirv-tools/libspirv.h

@@ -244,6 +244,15 @@ typedef enum spv_number_kind_t {
   SPV_NUMBER_FLOATING,
 } spv_number_kind_t;
 
+typedef enum spv_text_to_binary_options_t {
+  SPV_TEXT_TO_BINARY_OPTION_NONE = SPV_BIT(0),
+  // Numeric IDs in the binary will have the same values as in the source.
+  // Non-numeric IDs are allocated by filling in the gaps, starting with 1
+  // and going up.
+  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS = SPV_BIT(1),
+  SPV_FORCE_32_BIT_ENUM(spv_text_to_binary_options_t)
+} spv_text_to_binary_options_t;
+
 typedef enum spv_binary_to_text_options_t {
   SPV_BINARY_TO_TEXT_OPTION_NONE = SPV_BIT(0),
   SPV_BINARY_TO_TEXT_OPTION_PRINT = SPV_BIT(1),
@@ -416,6 +425,13 @@ spv_result_t spvTextToBinary(const spv_const_context context, const char* text,
                              const size_t length, spv_binary* binary,
                              spv_diagnostic* diagnostic);
 
+// Encodes the given SPIR-V assembly text to its binary representation. Same as
+// spvTextToBinary but with options. The options parameter is a bit field of
+// spv_text_to_binary_options_t.
+spv_result_t spvTextToBinaryWithOptions(
+    const spv_const_context context, const char* text, const size_t length,
+    const uint32_t options, spv_binary* binary, spv_diagnostic* diagnostic);
+
 // Frees an allocated text stream. This is a no-op if the text parameter
 // is a null pointer.
 void spvTextDestroy(spv_text text);

include/spirv-tools/libspirv.hpp

@@ -56,6 +56,9 @@ class ValidatorOptions {
 class SpirvTools {
  public:
   enum {
+    // Default assembling option used by assemble():
+    kDefaultAssembleOption = SPV_TEXT_TO_BINARY_OPTION_NONE,
+
     // Default disassembling option used by Disassemble():
     // * Avoid prefix comments from decoding the SPIR-V module header, and
     // * Use friendly names for variables.
@@ -86,11 +89,13 @@ class SpirvTools {
   // Assembles the given assembly |text| and writes the result to |binary|.
   // Returns true on successful assembling. |binary| will be kept untouched if
   // assembling is unsuccessful.
-  bool Assemble(const std::string& text, std::vector<uint32_t>* binary) const;
+  bool Assemble(const std::string& text, std::vector<uint32_t>* binary,
+                uint32_t options = kDefaultAssembleOption) const;
   // |text_size| specifies the number of bytes in |text|. A terminating null
   // character is not required to present in |text| as long as |text| is valid.
   bool Assemble(const char* text, size_t text_size,
-                std::vector<uint32_t>* binary) const;
+                std::vector<uint32_t>* binary,
+                uint32_t options = kDefaultAssembleOption) const;
 
   // Disassembles the given SPIR-V |binary| with the given |options| and writes
   // the assembly to |text|. Returns ture on successful disassembling. |text|
@@ -109,7 +114,8 @@ class SpirvTools {
   // |binary_size| specifies the number of words in |binary|.
   bool Validate(const uint32_t* binary, size_t binary_size) const;
   // Like the previous overload, but takes an options object.
-  bool Validate(const uint32_t* binary, size_t binary_size, const ValidatorOptions& options) const;
+  bool Validate(const uint32_t* binary, size_t binary_size,
+                const ValidatorOptions& options) const;
 
  private:
   struct Impl;  // Opaque struct for holding the data fields used by this class.

source/libspirv.cpp

@@ -39,15 +39,17 @@ void SpirvTools::SetMessageConsumer(MessageConsumer consumer) {
 }
 
 bool SpirvTools::Assemble(const std::string& text,
-                          std::vector<uint32_t>* binary) const {
+                          std::vector<uint32_t>* binary,
-  return Assemble(text.data(), text.size(), binary);
+                          uint32_t options) const {
+  return Assemble(text.data(), text.size(), binary, options);
 }
 
 bool SpirvTools::Assemble(const char* text, const size_t text_size,
-                          std::vector<uint32_t>* binary) const {
+                          std::vector<uint32_t>* binary,
+                          uint32_t options) const {
   spv_binary spvbinary = nullptr;
-  spv_result_t status =
+  spv_result_t status = spvTextToBinaryWithOptions(
-      spvTextToBinary(impl_->context, text, text_size, &spvbinary, nullptr);
+      impl_->context, text, text_size, options, &spvbinary, nullptr);
   if (status == SPV_SUCCESS) {
     binary->assign(spvbinary->code, spvbinary->code + spvbinary->wordCount);
   }

source/text.cpp

@@ -659,13 +659,57 @@ spv_result_t SetHeader(spv_target_env env, const uint32_t bound,
   return SPV_SUCCESS;
 }
 
+// Collects all numeric ids in the module source into |numeric_ids|.
+// This function is essentially a dry-run of spvTextToBinary.
+spv_result_t GetNumericIds(const libspirv::AssemblyGrammar& grammar,
+                           const spvtools::MessageConsumer& consumer,
+                           const spv_text text,
+                           std::set<uint32_t>* numeric_ids) {
+  libspirv::AssemblyContext context(text, consumer);
+
+  if (!text->str) return context.diagnostic() << "Missing assembly text.";
+
+  if (!grammar.isValid()) {
+    return SPV_ERROR_INVALID_TABLE;
+  }
+
+  // Skip past whitespace and comments.
+  context.advance();
+
+  while (context.hasText()) {
+    spv_instruction_t inst;
+
+    if (spvTextEncodeOpcode(grammar, &context, &inst)) {
+      return SPV_ERROR_INVALID_TEXT;
+    }
+
+    if (context.advance()) break;
+  }
+
+  *numeric_ids = context.GetNumericIds();
+  return SPV_SUCCESS;
+}
+
 // Translates a given assembly language module into binary form.
 // If a diagnostic is generated, it is not yet marked as being
 // for a text-based input.
-spv_result_t spvTextToBinaryInternal(const libspirv::AssemblyGrammar& grammar,
+spv_result_t spvTextToBinaryInternal(
-                                     const spvtools::MessageConsumer& consumer,
+    const libspirv::AssemblyGrammar& grammar,
-                                     const spv_text text, spv_binary* pBinary) {
+    const spvtools::MessageConsumer& consumer, const spv_text text,
-  libspirv::AssemblyContext context(text, consumer);
+    const uint32_t options, spv_binary* pBinary) {
+  // The ids in this set will have the same values both in source and binary.
+  // All other ids will be generated by filling in the gaps.
+  std::set<uint32_t> ids_to_preserve;
+
+  if (options & SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS) {
+    // Collect all numeric ids from the source into ids_to_preserve.
+    const spv_result_t result =
+        GetNumericIds(grammar, consumer, text, &ids_to_preserve);
+    if (result != SPV_SUCCESS) return result;
+  }
+
+  libspirv::AssemblyContext context(text, consumer, std::move(ids_to_preserve));
+
   if (!text->str) return context.diagnostic() << "Missing assembly text.";
 
   if (!grammar.isValid()) {
@@ -725,6 +769,15 @@ spv_result_t spvTextToBinary(const spv_const_context context,
                              const char* input_text,
                              const size_t input_text_size, spv_binary* pBinary,
                              spv_diagnostic* pDiagnostic) {
+  return spvTextToBinaryWithOptions(
+      context, input_text, input_text_size, SPV_BINARY_TO_TEXT_OPTION_NONE,
+      pBinary, pDiagnostic);
+}
+
+spv_result_t spvTextToBinaryWithOptions(
+    const spv_const_context context, const char* input_text,
+    const size_t input_text_size, const uint32_t options, spv_binary* pBinary,
+    spv_diagnostic* pDiagnostic) {
   spv_context_t hijack_context = *context;
   if (pDiagnostic) {
     *pDiagnostic = nullptr;
@@ -734,8 +787,8 @@ spv_result_t spvTextToBinary(const spv_const_context context,
   spv_text_t text = {input_text, input_text_size};
   libspirv::AssemblyGrammar grammar(&hijack_context);
 
-  spv_result_t result =
+  spv_result_t result = spvTextToBinaryInternal(
-      spvTextToBinaryInternal(grammar, hijack_context.consumer, &text, pBinary);
+      grammar, hijack_context.consumer, &text, options, pBinary);
   if (pDiagnostic && *pDiagnostic) (*pDiagnostic)->isTextSource = true;
 
   return result;

source/text_handler.cpp

@@ -14,6 +14,7 @@
 
 #include "text_handler.h"
 
+#include <algorithm>
 #include <cassert>
 #include <cstdlib>
 #include <cstring>
@@ -154,11 +155,33 @@ const IdType kUnknownType = {0, false, IdTypeClass::kBottom};
 // This represents all of the data that is only valid for the duration of
 // a single compilation.
 uint32_t AssemblyContext::spvNamedIdAssignOrGet(const char* textValue) {
-  if (named_ids_.end() == named_ids_.find(textValue)) {
+  if (!ids_to_preserve_.empty()) {
-    named_ids_[std::string(textValue)] = bound_++;
+    uint32_t id = 0;
+    if (spvutils::ParseNumber(textValue, &id)) {
+      if (ids_to_preserve_.find(id) != ids_to_preserve_.end()) {
+        bound_ = std::max(bound_, id + 1);
+        return id;
+      }
+    }
   }
-  return named_ids_[textValue];
+
+  const auto it = named_ids_.find(textValue);
+  if (it == named_ids_.end()) {
+    uint32_t id = next_id_++;
+    if (!ids_to_preserve_.empty()) {
+      while (ids_to_preserve_.find(id) != ids_to_preserve_.end()) {
+        id = next_id_++;
+      }
+    }
+
+    named_ids_.emplace(textValue, id);
+    bound_ = std::max(bound_, id + 1);
+    return id;
+  }
+
+  return it->second;
 }
+
 uint32_t AssemblyContext::getBound() const { return bound_; }
 
 spv_result_t AssemblyContext::advance() {
@@ -362,4 +385,14 @@ spv_ext_inst_type_t AssemblyContext::getExtInstTypeForId(uint32_t id) const {
   return std::get<1>(*type);
 }
 
+std::set<uint32_t> AssemblyContext::GetNumericIds() const {
+  std::set<uint32_t> ids;
+  for (const auto& kv : named_ids_) {
+    uint32_t id;
+    if (spvutils::ParseNumber(kv.first.c_str(), &id))
+      ids.insert(id);
+  }
+  return ids;
+}
+
 }  // namespace libspirv

source/text_handler.h

@@ -117,8 +117,10 @@ class ClampToZeroIfUnsignedType<
 // Encapsulates the data used during the assembly of a SPIR-V module.
 class AssemblyContext {
  public:
-  AssemblyContext(spv_text text, const spvtools::MessageConsumer& consumer)
+  AssemblyContext(spv_text text, const spvtools::MessageConsumer& consumer,
-      : current_position_({}), consumer_(consumer), text_(text), bound_(1) {}
+                  std::set<uint32_t>&& ids_to_preserve = std::set<uint32_t>())
+      : current_position_({}), consumer_(consumer), text_(text), bound_(1),
+        next_id_(1), ids_to_preserve_(std::move(ids_to_preserve))  {}
 
   // Assigns a new integer value to the given text ID, or returns the previously
   // assigned integer value if the ID has been seen before.
@@ -224,6 +226,11 @@ class AssemblyContext {
   // id is not the id for an extended instruction type.
   spv_ext_inst_type_t getExtInstTypeForId(uint32_t id) const;
 
+  // Returns a set consisting of each ID generated by spvNamedIdAssignOrGet from
+  // a numeric ID text representation. For example, generated from "%12" but not
+  // from "%foo".
+  std::set<uint32_t> GetNumericIds() const;
+
  private:
   // Maps ID names to their corresponding numerical ids.
   using spv_named_id_table = std::unordered_map<std::string, uint32_t>;
@@ -241,6 +248,8 @@ class AssemblyContext {
   spvtools::MessageConsumer consumer_;
   spv_text text_;
   uint32_t bound_;
+  uint32_t next_id_;
+  std::set<uint32_t> ids_to_preserve_;
 };
 }
 #endif  // _LIBSPIRV_TEXT_HANDLER_H_

source/util/string_utils.cpp

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <algorithm>
 #include <cstdint>
 #include <type_traits>
 
@@ -36,4 +37,3 @@ std::string CardinalToOrdinal(size_t cardinal) {
 }
 
 }  // namespace spvutils
-

test/CMakeLists.txt

@@ -159,6 +159,11 @@ add_spvtools_unittest(
   SRCS log_test.cpp
   LIBS ${SPIRV_TOOLS})
 
+add_spvtools_unittest(
+  TARGET preserve_numeric_ids
+  SRCS preserve_numeric_ids_test.cpp
+  LIBS ${SPIRV_TOOLS})
+
 add_subdirectory(opt)
 add_subdirectory(val)
 add_subdirectory(stats)

test/preserve_numeric_ids_test.cpp

@@ -0,0 +1,158 @@
+// Copyright (c) 2017 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Tests for unique type declaration rules validator.
+
+#include <string>
+
+#include "source/text.h"
+#include "source/text_handler.h"
+#include "test_fixture.h"
+
+namespace {
+
+using spvtest::ScopedContext;
+
+// Converts code to binary and then back to text.
+spv_result_t ToBinaryAndBack(
+    const std::string& before, std::string* after,
+    uint32_t text_to_binary_options = SPV_TEXT_TO_BINARY_OPTION_NONE,
+    uint32_t binary_to_text_options = SPV_BINARY_TO_TEXT_OPTION_NONE,
+    spv_target_env env = SPV_ENV_UNIVERSAL_1_0) {
+  ScopedContext ctx(env);
+  spv_binary binary;
+  spv_text text;
+
+  spv_result_t result = spvTextToBinaryWithOptions(
+      ctx.context, before.c_str(), before.size(), text_to_binary_options,
+      &binary, nullptr);
+  if (result != SPV_SUCCESS) {
+    return result;
+  }
+
+  result = spvBinaryToText(
+      ctx.context, binary->code, binary->wordCount, binary_to_text_options,
+      &text, nullptr);
+  if (result != SPV_SUCCESS) {
+    return result;
+  }
+
+  *after = std::string(text->str, text->length);
+
+  spvBinaryDestroy(binary);
+  spvTextDestroy(text);
+
+  return SPV_SUCCESS;
+}
+
+TEST(ToBinaryAndBack, DontPreserveNumericIds) {
+  const std::string before =
+R"(OpCapability Addresses
+OpCapability Kernel
+OpCapability GenericPointer
+OpCapability Linkage
+OpMemoryModel Physical32 OpenCL
+%i32 = OpTypeInt 32 1
+%u32 = OpTypeInt 32 0
+%f32 = OpTypeFloat 32
+%200 = OpTypeVoid
+%300 = OpTypeFunction %200
+%main = OpFunction %200 None %300
+%entry = OpLabel
+%100 = OpConstant %u32 100
+%1 = OpConstant %u32 200
+%2 = OpConstant %u32 300
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string expected =
+R"(OpCapability Addresses
+OpCapability Kernel
+OpCapability GenericPointer
+OpCapability Linkage
+OpMemoryModel Physical32 OpenCL
+%1 = OpTypeInt 32 1
+%2 = OpTypeInt 32 0
+%3 = OpTypeFloat 32
+%4 = OpTypeVoid
+%5 = OpTypeFunction %4
+%6 = OpFunction %4 None %5
+%7 = OpLabel
+%8 = OpConstant %2 100
+%9 = OpConstant %2 200
+%10 = OpConstant %2 300
+OpReturn
+OpFunctionEnd
+)";
+
+  std::string after;
+  EXPECT_EQ(SPV_SUCCESS, ToBinaryAndBack(before, &after,
+                                         SPV_TEXT_TO_BINARY_OPTION_NONE,
+                                         SPV_BINARY_TO_TEXT_OPTION_NO_HEADER));
+
+  EXPECT_EQ(expected, after);
+}
+
+TEST(TextHandler, PreserveNumericIds) {
+  const std::string before =
+R"(OpCapability Addresses
+OpCapability Kernel
+OpCapability GenericPointer
+OpCapability Linkage
+OpMemoryModel Physical32 OpenCL
+%i32 = OpTypeInt 32 1
+%u32 = OpTypeInt 32 0
+%f32 = OpTypeFloat 32
+%200 = OpTypeVoid
+%300 = OpTypeFunction %200
+%main = OpFunction %200 None %300
+%entry = OpLabel
+%100 = OpConstant %u32 100
+%1 = OpConstant %u32 200
+%2 = OpConstant %u32 300
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string expected =
+R"(OpCapability Addresses
+OpCapability Kernel
+OpCapability GenericPointer
+OpCapability Linkage
+OpMemoryModel Physical32 OpenCL
+%3 = OpTypeInt 32 1
+%4 = OpTypeInt 32 0
+%5 = OpTypeFloat 32
+%200 = OpTypeVoid
+%300 = OpTypeFunction %200
+%6 = OpFunction %200 None %300
+%7 = OpLabel
+%100 = OpConstant %4 100
+%1 = OpConstant %4 200
+%2 = OpConstant %4 300
+OpReturn
+OpFunctionEnd
+)";
+
+  std::string after;
+  EXPECT_EQ(SPV_SUCCESS,
+            ToBinaryAndBack(before, &after,
+                            SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS,
+                            SPV_BINARY_TO_TEXT_OPTION_NO_HEADER));
+
+  EXPECT_EQ(expected, after);
+}
+
+}  // namespace

tools/as/as.cpp

@@ -39,6 +39,10 @@ Options:
   --version       Display assembler version information.
   --target-env {vulkan1.0|spv1.0|spv1.1}
                   Use Vulkan1.0/SPIR-V1.0/SPIR-V1.1 validation rules.
+  --preserve-numeric-ids
+                  Numeric IDs in the binary will have the same values as in the
+                  source. Non-numeric IDs are allocated by filling in the gaps,
+                  starting with 1 and going up.
 )",
       argv0, argv0);
 }
@@ -46,6 +50,7 @@ Options:
 int main(int argc, char** argv) {
   const char* inFile = nullptr;
   const char* outFile = nullptr;
+  uint32_t options = 0;
   spv_target_env target_env = SPV_ENV_UNIVERSAL_1_1;
   for (int argi = 1; argi < argc; ++argi) {
     if ('-' == argv[argi][0]) {
@@ -83,6 +88,9 @@ int main(int argc, char** argv) {
             print_usage(argv[0]);
             return 0;
           }
+          if (0 == strcmp(argv[argi], "--preserve-numeric-ids")) {
+            options |= SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS;
+          }
           if (0 == strcmp(argv[argi], "--target-env")) {
             if (argi + 1 < argc) {
               const auto env_str = argv[++argi];
@@ -121,8 +129,8 @@ int main(int argc, char** argv) {
   spv_binary binary;
   spv_diagnostic diagnostic = nullptr;
   spv_context context = spvContextCreate(target_env);
-  spv_result_t error = spvTextToBinary(context, contents.data(),
+  spv_result_t error = spvTextToBinaryWithOptions(
-                                       contents.size(), &binary, &diagnostic);
+      context, contents.data(), contents.size(), options, &binary, &diagnostic);
   spvContextDestroy(context);
   if (error) {
     spvDiagnosticPrint(diagnostic);