linker: Allow modules to be partially linked

Fixes: https://github.com/KhronosGroup/SPIRV-Tools/issues/1144

include/spirv-tools/linker.hpp

@@ -26,7 +26,10 @@ namespace spvtools {
 
 class LinkerOptions {
  public:
-  LinkerOptions() : create_library_(false), verify_ids_(false) {}
+  LinkerOptions()
+      : create_library_(false),
+        verify_ids_(false),
+        allow_partial_linkage_(false) {}
 
   // Returns whether a library or an executable should be produced by the
   // linking phase.
@@ -50,9 +53,20 @@ class LinkerOptions {
   // context.
   void SetVerifyIds(bool verify_ids) { verify_ids_ = verify_ids; }
 
+  // Returns whether to allow for imported symbols to have no corresponding
+  // exported symbols
+  bool GetAllowPartialLinkage() const { return allow_partial_linkage_; }
+
+  // Sets whether to allow for imported symbols to have no corresponding
+  // exported symbols
+  void SetAllowPartialLinkage(bool allow_partial_linkage) {
+    allow_partial_linkage_ = allow_partial_linkage;
+  }
+
  private:
   bool create_library_;
   bool verify_ids_;
+  bool allow_partial_linkage_;
 };
 
 // Links one or more SPIR-V modules into a new SPIR-V module. That is, combine

source/link/linker.cpp

@@ -110,7 +110,8 @@ static spv_result_t MergeModules(const MessageConsumer& consumer,
 static spv_result_t GetImportExportPairs(
     const MessageConsumer& consumer, const ir::IRContext& linked_context,
     const DefUseManager& def_use_manager,
-    const DecorationManager& decoration_manager, LinkageTable* linkings_to_do);
+    const DecorationManager& decoration_manager, bool allow_partial_linkage,
+    LinkageTable* linkings_to_do);
 
 // Checks that for each pair of import and export, the import and export have
 // the same type as well as the same decorations.
@@ -134,7 +135,7 @@ static spv_result_t CheckImportExportCompatibility(
 // TODO(pierremoreau): Run a pass for removing dead instructions, for example
 //                     OpName for prototypes of imported funcions.
 static spv_result_t RemoveLinkageSpecificInstructions(
-    const MessageConsumer& consumer, bool create_executable,
+    const MessageConsumer& consumer, const LinkerOptions& options,
     const LinkageTable& linkings_to_do, DecorationManager* decoration_manager,
     ir::IRContext* linked_context);
 
@@ -222,9 +223,10 @@ spv_result_t Link(const Context& context, const uint32_t* const* binaries,
 
   // Phase 4: Find the import/export pairs
   LinkageTable linkings_to_do;
-  res = GetImportExportPairs(
-      consumer, linked_context, *linked_context.get_def_use_mgr(),
-      *linked_context.get_decoration_mgr(), &linkings_to_do);
+  res = GetImportExportPairs(consumer, linked_context,
+                             *linked_context.get_def_use_mgr(),
+                             *linked_context.get_decoration_mgr(),
+                             options.GetAllowPartialLinkage(), &linkings_to_do);
   if (res != SPV_SUCCESS) return res;
 
   // Phase 5: Ensure the import and export have the same types and decorations.
@@ -246,9 +248,9 @@ spv_result_t Link(const Context& context, const uint32_t* const* binaries,
 
   // Phase 8: Remove linkage specific instructions, such as import/export
   // attributes, linkage capability, etc. if applicable
-  res = RemoveLinkageSpecificInstructions(
-      consumer, !options.GetCreateLibrary(), linkings_to_do,
-      linked_context.get_decoration_mgr(), &linked_context);
+  res = RemoveLinkageSpecificInstructions(consumer, options, linkings_to_do,
+                                          linked_context.get_decoration_mgr(),
+                                          &linked_context);
   if (res != SPV_SUCCESS) return res;
 
   // Phase 9: Compact the IDs used in the module
@@ -481,7 +483,8 @@ static spv_result_t MergeModules(const MessageConsumer& consumer,
 static spv_result_t GetImportExportPairs(
     const MessageConsumer& consumer, const ir::IRContext& linked_context,
     const DefUseManager& def_use_manager,
-    const DecorationManager& decoration_manager, LinkageTable* linkings_to_do) {
+    const DecorationManager& decoration_manager, bool allow_partial_linkage,
+    LinkageTable* linkings_to_do) {
   spv_position_t position = {};
 
   if (linkings_to_do == nullptr)
@@ -561,7 +564,7 @@ static spv_result_t GetImportExportPairs(
     std::vector<LinkageSymbolInfo> possible_exports;
     const auto& exp = exports.find(import.name);
     if (exp != exports.end()) possible_exports = exp->second;
-    if (possible_exports.empty())
+    if (possible_exports.empty() && !allow_partial_linkage)
       return libspirv::DiagnosticStream(position, consumer,
                                         SPV_ERROR_INVALID_BINARY)
              << "Unresolved external reference to \"" << import.name << "\".";
@@ -571,6 +574,7 @@ static spv_result_t GetImportExportPairs(
              << "Too many external references, " << possible_exports.size()
              << ", were found for \"" << import.name << "\".";
 
+    if (!possible_exports.empty())
       linkings_to_do->emplace_back(import, possible_exports.front());
   }
 
@@ -623,7 +627,7 @@ static spv_result_t CheckImportExportCompatibility(
 }
 
 static spv_result_t RemoveLinkageSpecificInstructions(
-    const MessageConsumer& consumer, bool create_executable,
+    const MessageConsumer& consumer, const LinkerOptions& options,
     const LinkageTable& linkings_to_do, DecorationManager* decoration_manager,
     ir::IRContext* linked_context) {
   spv_position_t position = {};
@@ -689,21 +693,40 @@ static spv_result_t RemoveLinkageSpecificInstructions(
     }
   }
 
+  // If partial linkage is allowed, we need an efficient way to check whether
+  // an imported ID had a corresponding export symbol. As uses of the imported
+  // symbol have already been replaced by the exported symbol, use the exported
+  // symbol ID.
+  // TODO(pierremoreau): This will not work if the decoration is applied
+  //                     through a group, but the linker does not support that
+  //                     either.
+  std::unordered_set<SpvId> imports;
+  if (options.GetAllowPartialLinkage()) {
+    imports.reserve(linkings_to_do.size());
+    for (const auto& linking_entry : linkings_to_do)
+      imports.emplace(linking_entry.exported_symbol.id);
+  }
+
   // Remove import linkage attributes
   auto next = linked_context->annotation_begin();
   for (auto inst = next; inst != linked_context->annotation_end();
        inst = next) {
     ++next;
+    // If this is an import annotation:
+    // * if we do not allow partial linkage, remove all import annotations;
+    // * otherwise, remove the annotation only if there was a corresponding
+    //   export.
     if (inst->opcode() == SpvOpDecorate &&
         inst->GetSingleWordOperand(1u) == SpvDecorationLinkageAttributes &&
-        inst->GetSingleWordOperand(3u) == SpvLinkageTypeImport) {
+        inst->GetSingleWordOperand(3u) == SpvLinkageTypeImport &&
+        (!options.GetAllowPartialLinkage() ||
+         imports.find(inst->GetSingleWordOperand(0u)) != imports.end())) {
       linked_context->KillInst(&*inst);
     }
   }
 
-  // Remove export linkage attributes and Linkage capability if making an
-  // executable
-  if (create_executable) {
+  // Remove export linkage attributes if making an executable
+  if (!options.GetCreateLibrary()) {
     next = linked_context->annotation_begin();
     for (auto inst = next; inst != linked_context->annotation_end();
          inst = next) {
@@ -714,7 +737,11 @@ static spv_result_t RemoveLinkageSpecificInstructions(
         linked_context->KillInst(&*inst);
       }
     }
+  }
 
+  // Remove Linkage capability if making an executable and partial linkage is
+  // not allowed
+  if (!options.GetCreateLibrary() && !options.GetAllowPartialLinkage()) {
     for (auto& inst : linked_context->capabilities())
       if (inst.GetSingleWordInOperand(0u) == SpvCapabilityLinkage) {
         linked_context->KillInst(&inst);

test/link/CMakeLists.txt

@@ -46,3 +46,8 @@ add_spvtools_unittest(TARGET link_unique_ids
   SRCS unique_ids_test.cpp
   LIBS SPIRV-Tools-opt SPIRV-Tools-link
 )
+
+add_spvtools_unittest(TARGET link_partial_linkage
+  SRCS partial_linkage_test.cpp
+  LIBS SPIRV-Tools-opt SPIRV-Tools-link
+)

test/link/partial_linkage_test.cpp

@@ -0,0 +1,84 @@
+// Copyright (c) 2018 Pierre Moreau
+//
+// 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.
+
+#include "gmock/gmock.h"
+#include "linker_fixture.h"
+
+namespace {
+
+using ::testing::HasSubstr;
+using PartialLinkage = spvtest::LinkerTest;
+
+TEST_F(PartialLinkage, Allowed) {
+  const std::string body1 = R"(
+OpCapability Linkage
+OpDecorate %1 LinkageAttributes "foo" Import
+OpDecorate %2 LinkageAttributes "bar" Import
+%3 = OpTypeFloat 32
+%1 = OpVariable %3 Uniform
+%2 = OpVariable %3 Uniform
+)";
+  const std::string body2 = R"(
+OpCapability Linkage
+OpDecorate %1 LinkageAttributes "bar" Export
+%2 = OpTypeFloat 32
+%3 = OpConstant %2 3.14159
+%1 = OpVariable %2 Uniform %3
+)";
+
+  spvtest::Binary linked_binary;
+  spvtools::LinkerOptions linker_options;
+  linker_options.SetAllowPartialLinkage(true);
+  ASSERT_EQ(SPV_SUCCESS,
+            AssembleAndLink({body1, body2}, &linked_binary, linker_options));
+
+  const std::string expected_res = R"(OpCapability Linkage
+OpDecorate %1 LinkageAttributes "foo" Import
+%2 = OpTypeFloat 32
+%1 = OpVariable %2 Uniform
+%3 = OpConstant %2 3.14159
+%4 = OpVariable %2 Uniform %3
+)";
+  std::string res_body;
+  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER);
+  ASSERT_EQ(SPV_SUCCESS, Disassemble(linked_binary, &res_body))
+      << GetErrorMessage();
+  EXPECT_EQ(expected_res, res_body);
+}
+
+TEST_F(PartialLinkage, Disallowed) {
+  const std::string body1 = R"(
+OpCapability Linkage
+OpDecorate %1 LinkageAttributes "foo" Import
+OpDecorate %2 LinkageAttributes "bar" Import
+%3 = OpTypeFloat 32
+%1 = OpVariable %3 Uniform
+%2 = OpVariable %3 Uniform
+)";
+  const std::string body2 = R"(
+OpCapability Linkage
+OpDecorate %1 LinkageAttributes "bar" Export
+%2 = OpTypeFloat 32
+%3 = OpConstant %2 3.14159
+%1 = OpVariable %2 Uniform %3
+)";
+
+  spvtest::Binary linked_binary;
+  EXPECT_EQ(SPV_ERROR_INVALID_BINARY,
+            AssembleAndLink({body1, body2}, &linked_binary));
+  EXPECT_THAT(GetErrorMessage(),
+              HasSubstr("Unresolved external reference to \"foo\"."));
+}
+
+}  // anonymous namespace

tools/link/linker.cpp

@@ -36,6 +36,7 @@ Options:
   -h, --help              Print this help.
   -o                      Name of the resulting linked SPIR-V binary.
   --create-library        Link the binaries into a library, keeping all exported symbols.
+  --allow-partial-linkage Allow partial linkage by accepting imported symbols to be unresolved.
   --verify-ids            Verify that IDs in the resulting modules are truly unique.
   --version               Display linker version information
   --target-env            {vulkan1.0|spv1.0|spv1.1|spv1.2|opencl2.1|opencl2.2}
@@ -73,6 +74,8 @@ int main(int argc, char** argv) {
         options.SetCreateLibrary(true);
       } else if (0 == strcmp(cur_arg, "--verify-ids")) {
         options.SetVerifyIds(true);
+      } else if (0 == strcmp(cur_arg, "--allow-partial-linkage")) {
+        options.SetAllowPartialLinkage(true);
       } else if (0 == strcmp(cur_arg, "--version")) {
         printf("%s\n", spvSoftwareVersionDetailsString());
         // TODO(dneto): Add OpenCL 2.2 at least.