From 4fc055e24899a3a05d2212143fabb441b7e2cdf5 Mon Sep 17 00:00:00 2001
From: Ben Barsdell <bbarsdell@nvidia.com>
Date: Tue, 30 May 2023 11:00:16 +1000
Subject: [PATCH 1/2] Overhaul options handling

- Refactors all options handling to use a single parser implementation
  and new Options and OptionsVec classes.
- This makes the code significantly cleaner and more robust.
- Maintains backwards compatibility with string vectors at API
  boundaries via implicit conversions.
---
 jitify2.hpp     | 905 +++++++++++++++++++++++++++---------------------
 jitify2_test.cu |  81 ++++-
 2 files changed, 594 insertions(+), 392 deletions(-)

diff --git a/jitify2.hpp b/jitify2.hpp
index 485fee6..04908da 100644
--- a/jitify2.hpp
+++ b/jitify2.hpp
@@ -45,7 +45,10 @@
 // unnecessary dependencies.
 #ifdef JITIFY_SERIALIZATION_ONLY
 
+#include <algorithm>
+#include <cassert>
 #include <climits>
+#include <initializer_list>
 #include <iostream>
 #include <sstream>
 #include <streambuf>
@@ -100,11 +103,13 @@
 #endif
 
 #include <algorithm>
+#include <cassert>
 #include <cctype>
 #include <climits>
 #include <cstring>
 #include <fstream>
 #include <functional>
+#include <initializer_list>
 #include <iomanip>
 #include <iostream>
 #include <list>
@@ -220,6 +225,283 @@ using StringRef = const std::string&;
 using StringSlice = std::string;
 #endif
 
+class Option {
+ public:
+  Option() = default;
+  explicit Option(std::string _key, std::string _value = {},
+                  StringVec _repr = {})
+      : key_(std::move(_key)),
+        value_(std::move(_value)),
+        repr_(std::move(_repr)) {
+    if (repr_.empty()) {
+      repr_ = {key_};
+      if (!value_.empty()) {
+        repr_.front() += "=" + value_;
+      }
+    }
+    // TODO: Consider changing key and value to be views into key_and_value to
+    // avoid double-storage.
+    if (value_.empty()) {
+      key_and_value_ = key_;
+    } else {
+      key_and_value_.reserve(key_.size() + 1 + value_.size());
+      key_and_value_.append(key_);
+      key_and_value_.append("=");
+      key_and_value_.append(value_);
+    }
+  }
+
+  const std::string& key() const { return key_; }
+  const std::string& value() const { return value_; }
+  const std::string& key_and_value() const { return key_and_value_; }
+  const StringVec& original_representation() const { return repr_; }
+
+  friend const std::string& to_string(const Option& option) {
+    return option.key_and_value();
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const Option& option) {
+    return os << option.key_and_value();
+  }
+
+  friend bool operator==(const Option& lhs, const Option& rhs) {
+    return lhs.key_ == rhs.key_ && lhs.value_ == rhs.value_;
+  }
+  friend bool operator!=(const Option& lhs, const Option& rhs) {
+    return !(lhs == rhs);
+  }
+
+ private:
+  std::string key_;
+  std::string value_;
+  std::string key_and_value_;
+  StringVec repr_;
+};
+
+namespace detail {
+
+// Strip whitespace from string in-place.
+inline void ltrim(std::string* s) {
+  s->erase(s->begin(), std::find_if(s->begin(), s->end(), [](unsigned char c) {
+             return !std::isspace(c);
+           }));
+}
+inline void rtrim(std::string* s) {
+  s->erase(std::find_if(s->rbegin(), s->rend(),
+                        [](unsigned char c) { return !std::isspace(c); })
+               .base(),
+           s->end());
+}
+inline void trim(std::string* s) {
+  ltrim(s);
+  rtrim(s);
+}
+
+// Strip whitespace from a string view.
+inline StringSlice ltrim(StringRef s) {
+  size_t beg = std::find_if(s.begin(), s.end(),
+                            [](unsigned char c) { return !std::isspace(c); }) -
+               s.begin();
+  return s.substr(beg);
+}
+inline StringSlice rtrim(StringRef s) {
+  size_t end = std::find_if(s.rbegin(), s.rend(),
+                            [](unsigned char c) { return !std::isspace(c); })
+                   .base() -
+               s.begin();
+  return s.substr(0, end);
+}
+inline StringSlice trim(StringRef s) { return rtrim(ltrim(s)); }
+
+}  // namespace detail
+
+class OptionsVec {
+  using vec_type = std::vector<Option>;
+
+ public:
+  using iterator = vec_type::iterator;
+  using const_iterator = vec_type::const_iterator;
+
+  OptionsVec() = default;
+  // Allow implicit conversion.
+  OptionsVec(std::vector<Option> _options) : options_(std::move(_options)) {}
+  // Allow implicit conversion (to avoid breaking the old options API).
+  OptionsVec(const StringVec& raw_options) : ok_(parse(raw_options)) {}
+  // Allow implicit conversion (to avoid breaking the old options API).
+  OptionsVec(std::initializer_list<std::string> raw_options)
+      : OptionsVec(StringVec(raw_options)) {}
+
+  explicit operator bool() const { return ok_; }
+
+  bool ok() const { return ok_; }
+
+  size_t size() const {
+    assert(ok_);
+    return options_.size();
+  }
+
+  StringVec serialize() const {
+    assert(ok_);
+    return serialize_impl(true);
+  }
+
+  StringVec serialize_canonical() const {
+    assert(ok_);
+    return serialize_impl(false);
+  }
+
+  // Allow implicit conversion (to avoid breaking the old options API).
+  operator StringVec() const { return serialize(); }
+
+  // Removes all options with any of the specified keys, and returns whether any
+  // were removed.
+  bool pop(std::initializer_list<std::string> keys) {
+    assert(ok_);
+    auto iter = std::remove_if(
+        options_.begin(), options_.end(), [&](const Option& option) {
+          return std::find(keys.begin(), keys.end(), option.key()) !=
+                 keys.end();
+        });
+    if (iter == options_.end()) return false;
+    options_.resize(iter - options_.begin());
+    return true;
+  }
+
+  void pop_back() {
+    assert(ok_);
+    options_.pop_back();
+  }
+
+  // Returns the indexes of all options that match any of the specified keys.
+  std::vector<int> find(std::initializer_list<std::string> keys,
+                        size_t max_results = (size_t)-1) const {
+    assert(ok_);
+    std::vector<int> results;
+    for (int i = 0; i < (int)options_.size(); ++i) {
+      if (results.size() == max_results) break;
+      // Note: Using std::find instead of a hashmap because keys will typically
+      // be only 2-3 elements long.
+      if (std::find(keys.begin(), keys.end(), options_[i].key()) !=
+          keys.end()) {
+        results.push_back(i);
+      }
+    }
+    return results;
+  }
+
+  iterator erase(size_t idx) {
+    assert(ok_);
+    return options_.erase(options_.begin() + idx);
+  }
+
+  template <class InputIt>
+  iterator insert(const_iterator pos, InputIt first, InputIt last) {
+    assert(ok_);
+    return options_.insert(pos, first, last);
+  }
+
+  void push_back(Option option) {
+    assert(ok_);
+    options_.push_back(std::move(option));
+  }
+
+  template <typename... Args>
+  void emplace_back(Args&&... args) {
+    assert(ok_);
+    options_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  iterator begin() {
+    assert(ok_);
+    return options_.begin();
+  }
+  iterator end() {
+    assert(ok_);
+    return options_.end();
+  }
+  const_iterator begin() const {
+    assert(ok_);
+    return options_.begin();
+  }
+  const_iterator end() const {
+    assert(ok_);
+    return options_.end();
+  }
+
+  const Option& operator[](size_t idx) const {
+    assert(ok_);
+    return options_[idx];
+  }
+  Option& operator[](size_t idx) {
+    assert(ok_);
+    return options_[idx];
+  }
+
+  friend bool operator==(const OptionsVec& lhs, const OptionsVec& rhs) {
+    return lhs.ok_ == rhs.ok_ && lhs.options_ == rhs.options_;
+  }
+  friend bool operator!=(const OptionsVec& lhs, const OptionsVec& rhs) {
+    return !(lhs == rhs);
+  }
+
+ private:
+  // Parses a vector of option strings into a vector of Option objects. Also
+  // strips whitespace surrounding keys and values. Returns false on failure.
+  bool parse(const StringVec& options) {
+    for (size_t i = 0; i < options.size(); ++i) {
+      std::string option = options[i];
+      detail::trim(&option);  // Strip whitespace
+      if (option[0] != '-') {
+        return false;  // "Expected an option, got " + option
+      }
+      std::string key, val;
+      StringVec repr = {option};
+      const size_t eql = option.find('=');
+      if (i + 1 < options.size() && options[i + 1][0] != '-') {
+        // Parse "-key" "val".
+        key = option;
+        val = options[++i];
+        repr = {key, val};
+      } else if (eql != std::string::npos) {
+        // Parse "-key=val".
+        key = option.substr(0, eql);
+        val = option.substr(eql + 1);
+      } else if (option.size() > 2 &&
+                 // HACK: Special case for '-l<lib>' linker flag.
+                 (std::isupper(static_cast<unsigned char>(option[1])) ||
+                  (option[1] == 'l' && option.substr(0, 9) != "-lineinfo"))) {
+        // Parse "-Kval".
+        key = option.substr(0, 2);
+        val = option.substr(2);
+      } else {
+        // Parse "-key" (no value).
+        key = option;
+      }
+      detail::trim(&val);  // Strip whitespace
+      options_.emplace_back(std::move(key), std::move(val), std::move(repr));
+    }
+    return true;
+  }
+
+  // This is the inverse of parse().
+  StringVec serialize_impl(bool use_original_representation) const {
+    StringVec results;
+    results.reserve(options_.size());
+    for (const Option& option : options_) {
+      if (use_original_representation) {
+        const StringVec& repr = option.original_representation();
+        results.insert(results.end(), repr.begin(), repr.end());
+      } else {
+        results.push_back(option.key_and_value());
+      }
+    }
+    return results;
+  }
+
+  vec_type options_;
+  bool ok_ = true;
+};
+
 namespace serialization {
 
 // Stream buffer that can be initialized with data without copying.
@@ -313,6 +595,17 @@ inline bool deserialize(std::istream& stream, StringVec* v) {
   return true;
 }
 
+inline void serialize(std::ostream& stream, const OptionsVec& ov) {
+  serialize(stream, ov.serialize());
+}
+
+inline bool deserialize(std::istream& stream, OptionsVec* ov) {
+  StringVec v;
+  if (!deserialize(stream, &v)) return false;
+  *ov = OptionsVec(v);
+  return static_cast<bool>(*ov);
+}
+
 inline void serialize(std::ostream& stream, const StringMap& m) {
   serialize(stream, m.size());
   for (const auto& kv : m) {
@@ -503,39 +796,6 @@ inline std::string string_join(const StringVec& args, StringRef sep = ",",
   return result;
 }
 
-// Strip whitespace from string in-place.
-inline void ltrim(std::string* s) {
-  s->erase(s->begin(), std::find_if(s->begin(), s->end(), [](unsigned char c) {
-             return !std::isspace(c);
-           }));
-}
-inline void rtrim(std::string* s) {
-  s->erase(std::find_if(s->rbegin(), s->rend(),
-                        [](unsigned char c) { return !std::isspace(c); })
-               .base(),
-           s->end());
-}
-inline void trim(std::string* s) {
-  ltrim(s);
-  rtrim(s);
-}
-
-// Strip whitespace from a string view.
-inline StringSlice ltrim(StringRef s) {
-  size_t beg = std::find_if(s.begin(), s.end(),
-                            [](unsigned char c) { return !std::isspace(c); }) -
-               s.begin();
-  return s.substr(beg);
-}
-inline StringSlice rtrim(StringRef s) {
-  size_t end = std::find_if(s.rbegin(), s.rend(),
-                            [](unsigned char c) { return !std::isspace(c); })
-                   .base() -
-               s.begin();
-  return s.substr(0, end);
-}
-inline StringSlice trim(StringRef s) { return rtrim(ltrim(s)); }
-
 }  // namespace detail
 
 /*! Reflection utilities namespace. */
@@ -979,6 +1239,10 @@ inline HashType hash_value(const std::unordered_map<Key, Val>& m,
   }
   return result;
 }
+template <typename HashType>
+inline HashType hash_value(const OptionsVec& options, HashType seed = {}) {
+  return hash_value<HashType>(options.serialize_canonical(), seed);
+}
 
 // Based on fast-hash. See https://github.com/ztanml/fast-hash
 inline uint64_t fasthash64(uint64_t h) {
@@ -1944,8 +2208,8 @@ class LinkedProgramData
     : public serialization::Serializable<LinkedProgramData> {
   std::string cubin_;
   StringMap lowered_name_map_;
-  std::string log_;           // Linker log
-  StringVec linker_options_;  // Linker options that were used
+  std::string log_;            // Linker log
+  OptionsVec linker_options_;  // Linker options that were used
 
   JITIFY_DEFINE_SERIALIZABLE_MEMBERS(LinkedProgramData, cubin_,
                                      lowered_name_map_)
@@ -1953,7 +2217,7 @@ class LinkedProgramData
  public:
   LinkedProgramData() = default;
   LinkedProgramData(std::string cubin, StringMap lowered_name_map = {},
-                    std::string log = {}, StringVec linker_options = {})
+                    std::string log = {}, OptionsVec linker_options = {})
       : cubin_(std::move(cubin)),
         lowered_name_map_(std::move(lowered_name_map)),
         log_(std::move(log)),
@@ -1966,7 +2230,7 @@ class LinkedProgramData
   /*! Get the log returned from the linker. */
   const std::string& log() const { return log_; }
   /*! Get the options that were passed to the linker. */
-  const StringVec& linker_options() const { return linker_options_; }
+  const OptionsVec& linker_options() const { return linker_options_; }
 
   /*! Load the program as a module into the current CUDA context.
    *  \return A LoadedProgram object that contains either a valid
@@ -1992,66 +2256,28 @@ class LinkedProgram
   static LinkedProgram link(const std::string& program,
                             CUjitInputType program_type,
                             StringMap lowered_name_map = {},
-                            StringVec options = {});
+                            OptionsVec options = {});
   /*! Link multiple programs.
    * \note Remaining linker options in each program must match.
    * \see CompiledProgramData::link */
   static LinkedProgram link(size_t num_programs,
                             const CompiledProgramData* compiled_programs[],
-                            StringVec options = {});
+                            OptionsVec options = {});
 
   static LinkedProgram link(
       const std::vector<const CompiledProgram*>& compiled_programs,
-      StringVec options = {});
+      OptionsVec options = {});
 
  private:
   static LinkedProgram link_impl(size_t num_programs,
                                  const std::string* programs[],
                                  const CUjitInputType program_types[],
-                                 StringMap lowered_name_map, StringVec options);
+                                 StringMap lowered_name_map,
+                                 OptionsVec options);
 };
 
 namespace detail {
 
-using OptionsMap = std::unordered_map<std::string, StringVec>;
-
-// Parses a vector of option strings into a map of key -> values (one value for
-// each time the key is repeated in the input). Also strips whitespace
-// surrounding keys and values. Returns an error message on failure, otherwise
-// an empty string.
-// std::string
-inline bool parse_options(const StringVec& options, OptionsMap* options_map) {
-  for (size_t i = 0; i < options.size(); ++i) {
-    std::string option = options[i];
-    trim(&option);                       // Strip whitespace
-    if (option[0] != '-') return false;  //"Expected an option, got " + option;
-    std::string key, val;
-    size_t eql = option.find('=');
-    if (i + 1 < options.size() && options[i + 1][0] != '-') {
-      // Parse "-key" "val".
-      key = option;
-      val = options[++i];
-    } else if (eql != std::string::npos) {
-      // Parse "-key=val".
-      key = option.substr(0, eql);
-      val = option.substr(eql + 1);
-    } else if (option.size() > 2 &&
-               // HACK: Special case for '-l<lib>' linker flag.
-               (std::isupper(option[1]) || option[1] == 'l')) {
-      // Parse "-Kval".
-      key = option.substr(0, 2);
-      val = option.substr(2);
-    } else {
-      // Parse "-key" (no value).
-      key = option;
-    }
-    trim(&val);  // Strip whitespace
-    (*options_map)[key].push_back(std::move(val));
-  }
-  // return {};
-  return true;
-}
-
 inline std::string path_base(const std::string& p) {
   // "/usr/local/myfile.dat" -> "/usr/local"
   // "foo/bar"  -> "foo"
@@ -2163,7 +2389,7 @@ inline CUjitInputType get_cuda_jit_input_type(std::string* filename) {
 inline bool link_programs_culink(size_t num_programs,
                                  const std::string* programs[],
                                  const CUjitInputType program_types[],
-                                 const OptionsMap& options_map,
+                                 const OptionsVec& options_vec,
                                  std::string* error, std::string* log,
                                  std::string* linked_cubin) {
   std::vector<CUjit_option> option_keys;
@@ -2178,17 +2404,15 @@ inline bool link_programs_culink(size_t num_programs,
   }
 #endif
   StringVec link_files, link_paths;
-  for (const auto& key_val : options_map) {
-    const std::string& key = key_val.first;
-    const StringVec& vals = key_val.second;
-    std::string val = !vals.empty() ? vals.back() : "";
+  for (const Option& option : options_vec) {
+    const std::string& key = option.key();
+    const std::string& val = option.value();
     // Note: ptxas actually uses "-g" (lowercase), but we use "-G" to be
     // consistent with NVRTC and NVCC.
     if (key == "-G" || key == "--device-debug") {
       option_keys.push_back(CU_JIT_GENERATE_DEBUG_INFO);
       option_vals.push_back((void*)(intptr_t)1);
-      // HACK: Can't allow -lineinfo due to ambiguity with "-l<lib>".
-    } else if (/*key == "-lineinfo" ||*/ key == "--generate-line-info") {
+    } else if (key == "-lineinfo" || key == "--generate-line-info") {
       option_keys.push_back(CU_JIT_GENERATE_LINE_INFO);
       option_vals.push_back((void*)(intptr_t)1);
     } else if (key == "-arch" || key == "--gpu-name" ||
@@ -2216,9 +2440,9 @@ inline bool link_programs_culink(size_t num_programs,
       option_keys.push_back(CU_JIT_LOG_VERBOSE);
       option_vals.push_back((void*)(intptr_t)1);
     } else if (key == "-l") {
-      link_files = vals;
+      link_files.push_back(val);
     } else if (key == "-L") {
-      link_paths = vals;
+      link_paths.push_back(val);
 #if CUDA_VERSION >= 11040
       // LTO optimization options.
     } else if (key == "-ftz" || key == "--ftz") {
@@ -2357,7 +2581,7 @@ inline bool link_programs_culink(size_t num_programs,
 inline bool link_programs_nvjitlink(size_t num_programs,
                                     const std::string* programs[],
                                     const CUjitInputType program_types[],
-                                    const OptionsMap& options_map,
+                                    const OptionsVec& options_vec,
                                     std::string* error, std::string* log,
                                     std::string* linked_cubin) {
   StringVec options;
@@ -2370,16 +2594,14 @@ inline bool link_programs_nvjitlink(size_t num_programs,
   }
 #endif
   StringVec link_files, link_paths;
-  for (const auto& key_val : options_map) {
-    const std::string& key = key_val.first;
-    const StringVec& vals = key_val.second;
-    std::string val = !vals.empty() ? vals.back() : "";
+  for (const Option& option : options_vec) {
+    const std::string& key = option.key();
+    const std::string& val = option.value();
     // Note: ptxas actually uses "-g" (lowercase), but we use "-G" to be
     // consistent with NVRTC and NVCC.
     if (key == "-G" || key == "--device-debug") {
       options.push_back("-g");
-      // HACK: Can't allow -lineinfo due to ambiguity with "-l<lib>".
-    } else if (/*key == "-lineinfo" ||*/ key == "--generate-line-info") {
+    } else if (key == "-lineinfo" || key == "--generate-line-info") {
       options.push_back("-lineinfo");
     } else if (key == "-arch" || key == "--gpu-name" ||
                key == "--gpu-architecture") {
@@ -2399,9 +2621,9 @@ inline bool link_programs_nvjitlink(size_t num_programs,
     } else if (key == "-v" || key == "--verbose") {
       options.push_back("-verbose");
     } else if (key == "-l") {
-      link_files = vals;
+      link_files.push_back(val);
     } else if (key == "-L") {
-      link_paths = vals;
+      link_paths.push_back(val);
     } else if (key == "-ftz" || key == "--ftz") {
       options.push_back(is_true_value(val) ? "-ftz=1" : "-ftz=0");
     } else if (key == "-prec-div" || key == "--prec-div") {
@@ -2416,13 +2638,9 @@ inline bool link_programs_nvjitlink(size_t num_programs,
       options.push_back("-prec-sqrt=0");
       options.push_back("-fma=1");
     } else if (key == "-Xptxas" || key == "--ptxas-options") {
-      for (const std::string& _val : vals) {
-        options.push_back("-Xptxas=" + _val);
-      }
+      options.push_back("-Xptxas=" + val);
     } else if (key == "-Xnvvm" || key == "--nvvm-options") {
-      for (const std::string& _val : vals) {
-        options.push_back("-Xnvvm=" + _val);
-      }
+      options.push_back("-Xnvvm=" + val);
       // TODO: Work out how we should handle "-optimize-unused-variables", which
       // sounds like the same thing as our "--remove-unused-globals" support.
     } else {
@@ -2538,20 +2756,19 @@ inline bool link_programs_nvjitlink(size_t num_programs,
 // Note that this appends to *log if it is provided.
 inline bool link_programs(size_t num_programs, const std::string* programs[],
                           const CUjitInputType program_types[],
-                          const StringVec& options, std::string* error,
+                          OptionsVec options, std::string* error,
                           std::string* log, std::string* linked_cubin) {
   if (num_programs == 0) {
     if (error) *error = "Require at least one program to link";
     return false;
   }
 
-  OptionsMap options_map;
-  if (!parse_options(options, &options_map)) {
+  if (!options) {
     if (error) *error = "Syntax error in linker options";
     return false;
   }
   if (num_programs == 1 && program_types[0] == CU_JIT_INPUT_CUBIN &&
-      !options_map.count("-l")) {
+      options.find({"-l"}).empty()) {
     // No linking required, just return the given CUBIN.
     if (linked_cubin) *linked_cubin = *programs[0];
     return true;
@@ -2560,25 +2777,21 @@ inline bool link_programs(size_t num_programs, const std::string* programs[],
 #if CUDA_VERSION >= 12000
   bool use_culink = false;
 #endif
-  for (const std::string flag : {"-use-culink", "--use-culink"}) {
-    auto iter = options_map.find(flag);
-    if (iter != options_map.end()) {
+  if (options.pop({"-use-culink", "--use-culink"})) {
 #if CUDA_VERSION >= 12000
-      use_culink = true;
+    use_culink = true;
 #endif
-      options_map.erase(iter);
-    }
   }
 
   const bool result =
 #if CUDA_VERSION >= 12000
       !use_culink
           ? link_programs_nvjitlink(num_programs, programs, program_types,
-                                    options_map, error, log, linked_cubin)
+                                    options, error, log, linked_cubin)
           :
 #endif
-          link_programs_culink(num_programs, programs, program_types,
-                               options_map, error, log, linked_cubin);
+          link_programs_culink(num_programs, programs, program_types, options,
+                               error, log, linked_cubin);
   return result;
 }
 
@@ -2594,9 +2807,9 @@ class CompiledProgramData
   std::string nvvm_;   // Only available with NVRTC version >= 11.4
   // Maps name expressions to lowered symbol names (aka. unmangled to mangled).
   StringMap lowered_name_map_;
-  StringVec remaining_linker_options_;  // Passed on to LinkedProgram::link.
-  std::string log_;                     // Compilation log
-  StringVec compiler_options_;          // Compiler options that were used.
+  OptionsVec remaining_linker_options_;  // Passed on to LinkedProgram::link.
+  std::string log_;                      // Compilation log
+  OptionsVec compiler_options_;          // Compiler options that were used.
 
   JITIFY_DEFINE_SERIALIZABLE_MEMBERS(CompiledProgramData, ptx_, cubin_, nvvm_,
                                      lowered_name_map_,
@@ -2606,8 +2819,8 @@ class CompiledProgramData
   CompiledProgramData() = default;
   CompiledProgramData(std::string ptx, std::string cubin = {},
                       std::string nvvm = {}, StringMap lowered_name_map = {},
-                      StringVec linker_options = {}, std::string log = {},
-                      StringVec compiler_options = {})
+                      OptionsVec linker_options = {}, std::string log = {},
+                      OptionsVec compiler_options = {})
       : ptx_(std::move(ptx)),
         cubin_(std::move(cubin)),
         nvvm_(std::move(nvvm)),
@@ -2639,20 +2852,20 @@ class CompiledProgramData
   /*! Get the map of name expressions to lowered (mangled) symbol names. */
   const StringMap& lowered_name_map() const { return lowered_name_map_; }
   /*! Get the remaining options that will be passed on to the compiler. */
-  const StringVec& remaining_linker_options() const {
+  const OptionsVec& remaining_linker_options() const {
     return remaining_linker_options_;
   }
   /*! Get the log returned from the compiler. */
   const std::string& log() const { return log_; }
   /*! Get the options that were passed to the compiler. */
-  const StringVec& compiler_options() const { return compiler_options_; }
+  const OptionsVec& compiler_options() const { return compiler_options_; }
 
   /*! Link the program into a binary CUBIN object.
    *  \param extra_linker_options List of additional linker options.
    *  \return A LinkedProgram object that contains either a valid
    *    LinkedProgramData object or an error state.
    */
-  LinkedProgram link(StringVec extra_linker_options = {}) const {
+  LinkedProgram link(OptionsVec extra_linker_options = {}) const {
     const CompiledProgramData* compiled_programs[] = {this};
     return LinkedProgram::link(1, compiled_programs,
                                std::move(extra_linker_options));
@@ -2673,16 +2886,16 @@ class CompiledProgram
                                  const std::string& source,
                                  const StringMap& header_sources = {},
                                  const StringVec& name_expressions = {},
-                                 StringVec compiler_options = {},
-                                 StringVec linker_options = {});
+                                 OptionsVec compiler_options = {},
+                                 OptionsVec linker_options = {});
 
   /*! \see PreprocessedProgramData::compile */
   static CompiledProgram compile(const std::string& name,
                                  const std::string& source,
                                  const StringMap& header_sources = {},
                                  const std::string& name_expression = {},
-                                 StringVec compiler_options = {},
-                                 StringVec linker_options = {}) {
+                                 OptionsVec compiler_options = {},
+                                 OptionsVec linker_options = {}) {
     return compile(name, source, header_sources, StringVec({name_expression}),
                    std::move(compiler_options), std::move(linker_options));
   }
@@ -2690,9 +2903,9 @@ class CompiledProgram
 
 inline LinkedProgram LinkedProgram::link(
     size_t num_programs, const CompiledProgramData* compiled_programs[],
-    StringVec options) {
+    OptionsVec options) {
   if (num_programs == 0) return Error("Must have at least one program to link");
-  const StringVec& prog_linker_options =
+  const OptionsVec& prog_linker_options =
       compiled_programs[0]->remaining_linker_options();
   StringMap lowered_name_map = compiled_programs[0]->lowered_name_map();
   size_t total_lowered_names = lowered_name_map.size();
@@ -2744,7 +2957,7 @@ inline LinkedProgram LinkedProgram::link(
 
 inline LinkedProgram LinkedProgram::link(
     const std::vector<const CompiledProgram*>& compiled_programs,
-    StringVec options) {
+    OptionsVec options) {
   std::vector<const CompiledProgramData*> prog_ptrs;
   prog_ptrs.reserve(compiled_programs.size());
   for (const CompiledProgram* compiled_program_ptr : compiled_programs) {
@@ -2758,7 +2971,7 @@ inline LinkedProgram LinkedProgram::link(
 inline LinkedProgram LinkedProgram::link(const std::string& program,
                                          CUjitInputType program_type,
                                          StringMap lowered_name_map,
-                                         StringVec options) {
+                                         OptionsVec options) {
   const std::string* programs[] = {&program};
   return link_impl(1, programs, &program_type, std::move(lowered_name_map),
                    std::move(options));
@@ -2767,7 +2980,7 @@ inline LinkedProgram LinkedProgram::link(const std::string& program,
 inline LinkedProgram LinkedProgram::link_impl(
     size_t num_programs, const std::string* programs[],
     const CUjitInputType program_types[], StringMap lowered_name_map,
-    StringVec options) {
+    OptionsVec options) {
   std::string error, log, linked_cubin;
   log = detail::string_join(options, " ", "Linker options: \"", "\"\n");
   if (!detail::link_programs(num_programs, programs, program_types, options,
@@ -2915,73 +3128,44 @@ inline LibNvrtc& nvrtc() {
 
 namespace detail {
 
-// Parses and removes the (first) architecture flag from the given vector of
-// options.
+// Parses the (first) architecture flag from the given vector of options.
 // Returns 0 on failure or if no architecture option is found.
 // Sets *error on failure (if provided).
 // Returns -1 if the arch value is the special string "compute_." or "sm_.".
 // Otherwise returns the integer arch value.
 // On success, sets *is_virtual to true if a "compute_" value was found, or
-// false for an "arch_" value, and *beg_idx and *end_idx are set to the range
-// of the arch entries in the options vector (e.g., so that they can be erased
-// by the caller).
-inline int parse_arch_flag(const StringVec& options, bool* is_virtual,
+// false for an "arch_" value, and *idx is set to the index of the option within
+// the options vector (e.g., so that it can be erased by the caller).
+inline int parse_arch_flag(const OptionsVec& options, bool* is_virtual,
                            std::string* error = nullptr,
-                           size_t* beg_idx = nullptr,
-                           size_t* end_idx = nullptr) {
-  for (int i = 0; i < (int)options.size(); ++i) {
-    StringRef option = ltrim(options[i]);
-    size_t key_end = 0;
-    if (startswith(option, "-arch")) {
-      key_end = std::strlen("-arch");
-    } else if (startswith(option, "--gpu-architecture")) {
-      key_end = std::strlen("--gpu-architecture");
-    } else if (startswith(option, "--gpu-name")) {  // ptxas flag name
-      key_end = std::strlen("--gpu-name");
-    } else {
-      continue;
-    }
-    size_t eql_beg = option.find("=", key_end);
-    StringSlice value;
-    bool found_eql = eql_beg != std::string::npos;
-    if (found_eql) {
-      value = trim(option.substr(eql_beg + 1));
-    } else {
-      if (i + 1 == (int)options.size()) {
-        if (error) *error = "Expected value after option.";
-        return 0;
-      }
-      value = trim(options[i + 1]);
-    }
-    if (startswith(value, "compute_")) {
-      *is_virtual = true;
-      value = value.substr(std::strlen("compute_"));
-    } else if (startswith(value, "sm_")) {
-      *is_virtual = false;
-      value = value.substr(std::strlen("sm_"));
-    } else {
-      if (error) *error = "Expected value to begin with 'compute_' or 'sm_'.";
+                           size_t* idx = nullptr) {
+  const std::vector<int> idxs =
+      options.find({"-arch", "--gpu-architecture", "--gpu-name"}, 1);
+  if (idxs.empty()) return 0;  // // No architecture flag found
+  std::string value = options[idxs[0]].value();
+  if (startswith(value, "compute_")) {
+    *is_virtual = true;
+    value = value.substr(std::strlen("compute_"));
+  } else if (startswith(value, "sm_")) {
+    *is_virtual = false;
+    value = value.substr(std::strlen("sm_"));
+  } else {
+    if (error) *error = "Expected value to begin with 'compute_' or 'sm_'.";
+    return 0;
+  }
+  int result;
+  if (value == ".") {
+    result = -1;
+  } else {
+    int cc = std::atoi(std::string(value).c_str());
+    if (cc == 0) {
+      if (error) *error = "Failed to parse a valid architecture number.";
       return 0;
     }
-    int result;
-    if (value == ".") {
-      result = -1;
-    } else {
-      int cc = std::atoi(std::string(value).c_str());
-      if (cc == 0) {
-        if (error) *error = "Failed to parse a valid architecture number.";
-        return 0;
-      }
-      result = cc;
-    }
-    // Store iterators to the arch entries in options.
-    if (beg_idx) *beg_idx = i;
-    if (end_idx) *end_idx = i + (1 + !found_eql);
-    return result;
+    result = cc;
   }
-  if (beg_idx) *beg_idx = 0;
-  if (end_idx) *end_idx = 0;
-  return 0;  // No architecture option found
+  if (idx) *idx = idxs[0];
+  return result;
 }
 
 // Returns 0 on failure and sets *error if provided. Otherwise returns a compute
@@ -3063,8 +3247,8 @@ inline int limit_to_supported_compute_capability(int cc,
 // necessary, filling in the architecture flag in both compiler_options and
 // linker_options.
 // Returns false on failure and sets *error if provided.
-inline bool process_architecture_flags(StringVec* compiler_options,
-                                       StringVec* linker_options,
+inline bool process_architecture_flags(OptionsVec* compiler_options,
+                                       OptionsVec* linker_options,
                                        std::string* error_ptr = nullptr) {
   std::string error;
   auto check_error = [&]() {
@@ -3093,13 +3277,13 @@ inline bool process_architecture_flags(StringVec* compiler_options,
     return false;
   }
   // Now parse compiler options.
-  size_t beg_idx, end_idx;
-  int given_cc = parse_arch_flag(*compiler_options, &is_virtual, &error,
-                                 &beg_idx, &end_idx);
+  size_t idx;
+  int given_cc = parse_arch_flag(*compiler_options, &is_virtual, &error, &idx);
   if (!check_error()) return false;
   // Remove the parsed arch flag entries; they are replaced below.
-  compiler_options->erase(compiler_options->begin() + beg_idx,
-                          compiler_options->begin() + end_idx);
+  if (given_cc != 0) {
+    compiler_options->erase(idx);
+  }
   int real_cc;
   if (linker_cc != 0) {
     real_cc = linker_cc;
@@ -3146,52 +3330,40 @@ inline bool process_architecture_flags(StringVec* compiler_options,
   // Add the computed arch flag back to the compiler options and to the linker
   // options.
   if (virt_cc) {
-    compiler_options->push_back("-arch=compute_" + std::to_string(virt_cc));
+    compiler_options->push_back(
+        Option("-arch", "compute_" + std::to_string(virt_cc)));
   } else {
-    compiler_options->push_back("-arch=sm_" + std::to_string(real_cc));
+    compiler_options->push_back(
+        Option("-arch", "sm_" + std::to_string(real_cc)));
   }
   if (linker_cc == 0) {
-    linker_options->push_back("-arch=sm_" + std::to_string(real_cc));
+    linker_options->push_back(Option("-arch", "sm_" + std::to_string(real_cc)));
   }
   return true;
 }
 
-inline void add_std_flag_if_not_specified(StringVec* options,
-                                          std::string value = "c++11") {
-  for (const std::string& option : *options) {
-    if (option.find("--std") != std::string::npos ||
-        option.find("-std") != std::string::npos) {
-      // A standard was explicitly specified, don't change anything.
-      return;
-    }
+// Returns the standard year (e.g., 11 for c++11).
+inline int add_std_flag_if_not_specified(OptionsVec* options,
+                                         int default_standard_year = 11) {
+  auto option_inds = options->find({"--std", "-std"});
+  if (!option_inds.empty()) {
+    const char* value = (*options)[option_inds.back()].value().c_str();
+    if (*value++ != 'c' || *value++ != '+' || *value++ != '+') return 0;
+    int standard_year = std::atoi(value);
+    return standard_year;
   }
   // Jitify must be compiled with C++11 support, so we default to enabling it
   // for the JIT-compiled code too.
-  options->push_back("-std=" + value);
+  std::string value = "c++" + std::to_string(default_standard_year);
+  options->emplace_back("-std", value);
+  return default_standard_year;
 }
 
-inline void add_default_device_flag_if_not_specified(StringVec* options) {
-  for (const std::string& option : *options) {
-    if (option.find("--device-as-default-execution-space") !=
-            std::string::npos ||
-        option.find("-default-device") != std::string::npos) {
-      return;  // Already specified, do nothing.
-    }
+inline void add_default_device_flag_if_not_specified(OptionsVec* options) {
+  if (options->find({"--device-as-default-execution-space", "-default-device"})
+          .empty()) {
+    options->emplace_back("-default-device");
   }
-  options->push_back("-default-device");
-}
-
-inline bool pop_flag(StringVec* options, const std::string& short_flag,
-                     const std::string& long_flag) {
-  auto it = std::remove_if(options->begin(), options->end(),
-                           [&](const std::string& opt) {
-                             return opt == short_flag || opt == long_flag;
-                           });
-  if (it != options->end()) {
-    options->resize(it - options->begin());
-    return true;
-  }
-  return false;
 }
 
 // Demangles nested variable names using the PTX name mangling scheme
@@ -3299,7 +3471,7 @@ inline bool ptx_remove_unused_globals(std::string* ptx);  // Defined below
 //   well as any global definitions found in the generated PTX.
 inline bool compile_program(
     const std::string& name, const std::string& source,
-    const StringMap& header_sources, const StringVec& options,
+    const StringMap& header_sources, const OptionsVec& options,
     std::string* error = nullptr, std::string* log = nullptr,
     std::string* ptx = nullptr, std::string* cubin = nullptr,
     std::string* nvvm = nullptr, const StringVec& name_expressions = {},
@@ -3321,13 +3493,15 @@ inline bool compile_program(
 
   std::vector<const char*> options_c;
   options_c.reserve(options.size());
-  for (const std::string& option : options) {
+  for (const Option& option : options) {
     if (nvrtc().get_version() < 11010) {
       // This NVRTC doesn't support specifying c++03 explicitly, so remove it.
-      // TODO: Should also support "-std" and "c++03" as separate entries.
-      if (option == "-std=c++03" || option == "--std=c++03") continue;
+      if ((option.key() == "-std" || option.key() == "--std") &&
+          option.value() == "c++03") {
+        continue;
+      }
     }
-    options_c.push_back(option.c_str());
+    options_c.push_back(option.key_and_value().c_str());
   }
 
 #define JITIFY_CHECK_NVRTC(call)                                      \
@@ -3533,49 +3707,20 @@ inline bool ptx_remove_unused_globals(std::string* ptx) {
   return true;
 }
 
-// Returns false if there is a syntax error in the options.
-inline bool copy_compiler_option_for_linker_ptxas(
-    const StringVec& compiler_options, StringVec* linker_options,
-    bool has_value, StringRef short_key, StringRef long_key,
-    StringRef output_key = {}) {
+inline void copy_compiler_flag_for_linker_ptxas(
+    const OptionsVec& compiler_options, OptionsVec* linker_options,
+    std::initializer_list<std::string> keys, StringRef output_key = {}) {
   // First check if the option is already specified in linker_options.
-  for (const std::string& raw_option : *linker_options) {
-    StringRef option = ltrim(raw_option);
-    if (startswith(option, short_key) || startswith(option, long_key)) {
-      return true;  // Already specified, don't do anything
-    }
-  }
-  // Now find the option in compiler_options.
-  StringSlice key, val;
-  for (size_t i = 0; i < compiler_options.size(); ++i) {
-    StringRef option = ltrim(compiler_options[i]);
-    if (startswith(option, short_key) || startswith(option, long_key)) {
-      if (!has_value) {
-        key = option;
-        break;
-      }
-      size_t eql_pos = option.find('=', short_key.size());
-      if (eql_pos != std::string::npos) {
-        key = option.substr(0, eql_pos);
-        val = option.substr(eql_pos + 1);
-      } else {
-        key = option;
-        if (i + 1 == compiler_options.size()) return false;  // Syntax error
-        val = compiler_options[i + 1];
-      }
-      break;
-    }
+  if (!linker_options->find(keys, 1).empty()) return;
+  // Copy any matching compiler option to the linker options.
+  const std::vector<int> idxs = compiler_options.find(keys, 1);
+  for (int idx : idxs) {
+    const Option& compiler_option = compiler_options[idx];
+    const Option linker_option =
+        output_key.empty() ? compiler_option
+                           : Option(output_key, compiler_option.value());
+    linker_options->push_back(linker_option);
   }
-  if (key.empty()) return true;  // Nothing to copy
-  if (!output_key.empty()) {
-    key = output_key;
-  }
-  if (has_value) {
-    linker_options->push_back(string_concat(key, "=", val));
-  } else {
-    linker_options->push_back(static_cast<std::string>(key));
-  }
-  return true;
 }
 
 }  // namespace detail
@@ -3583,16 +3728,18 @@ inline bool copy_compiler_option_for_linker_ptxas(
 inline CompiledProgram CompiledProgram::compile(
     const std::string& name, const std::string& source,
     const StringMap& header_sources, const StringVec& name_expressions,
-    StringVec compiler_options, StringVec linker_options) {
+    OptionsVec compiler_options, OptionsVec linker_options) {
+  if (!compiler_options) return Error("Failed to parse compiler options");
+  if (!linker_options) return Error("Failed to parse linker options");
   std::string error;
   if (!detail::process_architecture_flags(&compiler_options, &linker_options,
                                           &error)) {
     return Error("Failed to process architecture flags: " + error);
   }
-  detail::add_std_flag_if_not_specified(&compiler_options, "c++11");
+  detail::add_std_flag_if_not_specified(&compiler_options, 11);
   detail::add_default_device_flag_if_not_specified(&compiler_options);
-  bool should_remove_unused_globals = detail::pop_flag(
-      &compiler_options, "-remove-unused-globals", "--remove-unused-globals");
+  bool should_remove_unused_globals = compiler_options.pop(
+      {"-remove-unused-globals", "--remove-unused-globals"});
   std::string log, ptx, cubin, nvvm;
   StringMap lowered_name_map;
   if (!detail::compile_program(name, source, header_sources, compiler_options,
@@ -3617,35 +3764,27 @@ inline CompiledProgram CompiledProgram::compile(
   // the linker does ptx->cubin compilation prior to linking. This allows users
   // to specify these options in compiler_options without having to worry about
   // whether they also need to be passed in linker_options.
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/false, "-G",
-      "--device-debug");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/false, "-lineinfo",
-      "--generate-line-info",
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-G", "--device-debug"});
+  detail::copy_compiler_flag_for_linker_ptxas(
+      compiler_options, &linker_options, {"-lineinfo", "--generate-line-info"},
       "--generate-line-info");  // Note that linker doesn't support "-lineinfo"
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-maxrregcount",
-      "--maxrregcount");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-ftz", "--ftz");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-prec-div",
-      "--prec-div");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-prec-sqrt",
-      "--prec-sqrt");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-fmad", "--fmad");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-use_fast_math",
-      "--use_fast_math");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-Xptxas",
-      "--ptxas-options");
-  detail::copy_compiler_option_for_linker_ptxas(
-      compiler_options, &linker_options, /*has_value=*/true, "-Xnvvm",
-      "--nvvm-options");
+  detail::copy_compiler_flag_for_linker_ptxas(
+      compiler_options, &linker_options, {"-maxrregcount", "--maxrregcount"});
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-ftz", "--ftz"});
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-prec-div", "--prec-div"});
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-prec-sqrt", "--prec-sqrt"});
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-fmad", "--fmad"});
+  detail::copy_compiler_flag_for_linker_ptxas(
+      compiler_options, &linker_options, {"-use_fast_math", "--use_fast_math"});
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-Xptxas", "--ptxas-options"});
+  detail::copy_compiler_flag_for_linker_ptxas(compiler_options, &linker_options,
+                                              {"-Xnvvm", "--nvvm-options"});
 
   return CompiledProgram(std::move(ptx), std::move(cubin), std::move(nvvm),
                          std::move(lowered_name_map), std::move(linker_options),
@@ -3682,8 +3821,8 @@ class PreprocessedProgramData
   std::string source_;
   StringMap header_sources_;
   // Note that these are the options to be passed on to Compiled/LinkedProgram.
-  StringVec remaining_compiler_options_;
-  StringVec remaining_linker_options_;
+  OptionsVec remaining_compiler_options_;
+  OptionsVec remaining_linker_options_;
   std::string header_log_;
   std::string compile_log_;
 
@@ -3696,8 +3835,8 @@ class PreprocessedProgramData
   PreprocessedProgramData() = default;
   PreprocessedProgramData(std::string name, std::string source,
                           StringMap header_sources = {},
-                          StringVec remaining_compiler_options = {},
-                          StringVec remaining_linker_options = {},
+                          OptionsVec remaining_compiler_options = {},
+                          OptionsVec remaining_linker_options = {},
                           std::string header_log = {},
                           std::string compile_log = {})
       : name_(std::move(name)),
@@ -3738,11 +3877,11 @@ class PreprocessedProgramData
   /*! Get the header sources map. */
   const StringMap& header_sources() const { return header_sources_; }
   /*! Get the remaining options that will be passed on to the compiler. */
-  const StringVec& remaining_compiler_options() const {
+  const OptionsVec& remaining_compiler_options() const {
     return remaining_compiler_options_;
   }
   /*! Get the remaining options that will be passed on to the linker. */
-  const StringVec& remaining_linker_options() const {
+  const OptionsVec& remaining_linker_options() const {
     return remaining_linker_options_;
   }
   /*! Get the log of header lookups made during preprocessing. */
@@ -3764,8 +3903,8 @@ class PreprocessedProgramData
    */
   CompiledProgram compile(const StringVec& name_expressions = {},
                           const StringMap& extra_header_sources = {},
-                          StringVec extra_compiler_options = {},
-                          StringVec extra_linker_options = {}) const {
+                          OptionsVec extra_compiler_options = {},
+                          OptionsVec extra_linker_options = {}) const {
     StringMap combined_header_sources;
     const StringMap& combined_header_sources_ref = detail::merge(
         header_sources_, extra_header_sources, &combined_header_sources);
@@ -3793,8 +3932,8 @@ class PreprocessedProgramData
    */
   CompiledProgram compile(const std::string& name_expression,
                           const StringMap& extra_header_sources = {},
-                          StringVec extra_compiler_options = {},
-                          StringVec extra_linker_options = {}) const {
+                          OptionsVec extra_compiler_options = {},
+                          OptionsVec extra_linker_options = {}) const {
     // Allow name_expression="" to be passed instead of name_expression={}
     // (which is ambiguous with the overload above that takes a StringVec).
     StringVec name_expressions =
@@ -3811,8 +3950,8 @@ class PreprocessedProgramData
    */
   LoadedProgram load(const StringVec& name_expressions = {},
                      const StringMap& extra_header_sources = {},
-                     StringVec extra_compiler_options = {},
-                     StringVec extra_linker_options = {}) const {
+                     OptionsVec extra_compiler_options = {},
+                     OptionsVec extra_linker_options = {}) const {
     CompiledProgram compiled = compile(name_expressions, extra_header_sources,
                                        std::move(extra_compiler_options),
                                        std::move(extra_linker_options));
@@ -3829,8 +3968,8 @@ class PreprocessedProgramData
    */
   Kernel get_kernel(std::string name, StringVec other_name_expressions = {},
                     const StringMap& extra_header_sources = {},
-                    StringVec extra_compiler_options = {},
-                    StringVec extra_linker_options = {}) const {
+                    OptionsVec extra_compiler_options = {},
+                    OptionsVec extra_linker_options = {}) const {
     other_name_expressions.push_back(name);
     CompiledProgram compiled = compile(
         other_name_expressions, extra_header_sources,
@@ -3859,8 +3998,8 @@ class PreprocessedProgram
   /*! \see ProgramData::preprocess */
   static PreprocessedProgram preprocess(std::string name, std::string source,
                                         StringMap header_sources = {},
-                                        StringVec compiler_options = {},
-                                        StringVec linker_options = {},
+                                        OptionsVec compiler_options = {},
+                                        OptionsVec linker_options = {},
                                         FileCallback header_callback = nullptr);
 };
 
@@ -5960,17 +6099,13 @@ inline std::string minify_cpp_source(const std::string& source) {
       remove_cpp_comments_and_line_continuations(source));
 }
 
-inline void extract_include_paths(StringVec* options,
+inline void extract_include_paths(OptionsVec* options,
                                   StringVec* include_paths) {
-  auto options_iter = options->begin();
-  while (options_iter != options->end()) {
-    const std::string& option = *options_iter;
-    if (option.substr(0, 2) == "-I") {
-      include_paths->push_back(option.substr(2));
-      options_iter = options->erase(options_iter);
-    } else {
-      ++options_iter;
-    }
+  const std::vector<int> idxs = options->find({"-I"});
+  for (int i = (int)idxs.size() - 1; i >= 0; --i) {
+    const int idx = idxs[i];
+    include_paths->push_back((*options)[idx].value());
+    options->erase(idx);
   }
 }
 
@@ -5978,12 +6113,11 @@ inline void extract_include_paths(StringVec* options,
 
 inline PreprocessedProgram PreprocessedProgram::preprocess(
     std::string name, std::string source, StringMap header_sources,
-    StringVec compiler_options, StringVec linker_options,
+    OptionsVec compiler_options, OptionsVec linker_options,
     FileCallback header_callback) {
   // Add pre-include built-in JIT-safe headers.
-  bool use_system_headers_war =
-      !detail::pop_flag(&compiler_options, "-no-system-headers-workaround",
-                        "--no-system-headers-workaround");
+  bool use_system_headers_war = !compiler_options.pop(
+      {"-no-system-headers-workaround", "--no-system-headers-workaround"});
 #if CUDA_VERSION >= 11000
   // This issue with /usr/include always being searched is fixed in this NVRTC.
   use_system_headers_war = false;
@@ -5997,32 +6131,31 @@ inline PreprocessedProgram PreprocessedProgram::preprocess(
       header_sources.emplace(header_name, header_source);
     }
   }
-  if (!detail::pop_flag(&compiler_options, "-no-preinclude-workarounds",
-                        "--no-preinclude-workarounds")) {
+  if (!compiler_options.pop(
+          {"-no-preinclude-workarounds", "--no-preinclude-workarounds"})) {
     header_sources.emplace(
         "jitify_preinclude.h",
         detail::get_jitsafe_headers_map().at("jitify_preinclude.h"));
-    compiler_options.push_back("-include=jitify_preinclude.h");
+    compiler_options.push_back(Option("-include", "jitify_preinclude.h"));
   }
-  detail::add_std_flag_if_not_specified(&compiler_options, "c++11");
+  detail::add_std_flag_if_not_specified(&compiler_options, 11);
   detail::add_default_device_flag_if_not_specified(&compiler_options);
-  bool minify = detail::pop_flag(&compiler_options, "-m", "--minify");
+  bool minify = compiler_options.pop({"-m", "--minify"});
   // TODO: This flag is experimental, because the implementation does not
   // support transformations of "namespace std {" (as used for specializations).
-  bool use_cuda_std =
-      detail::pop_flag(&compiler_options, "-cuda-std", "--cuda-std");
-  bool replace_pragma_once = !detail::pop_flag(
-      &compiler_options, "-no-replace-pragma-once", "--no-replace-pragma-once");
-  bool use_builtin_headers = !detail::pop_flag(
-      &compiler_options, "-no-builtin-headers", "--no-builtin-headers");
+  bool use_cuda_std = compiler_options.pop({"-cuda-std", "--cuda-std"});
+  bool replace_pragma_once = !compiler_options.pop(
+      {"-no-replace-pragma-once", "--no-replace-pragma-once"});
+  bool use_builtin_headers = !compiler_options.pop(
+      {"-no-builtin-headers", "--no-builtin-headers", "-nostdinc"});
   // Allow -nostdinc as an alias for -no-builtin-headers.
-  if (detail::pop_flag(&compiler_options, "-nostdinc", "--nostdinc")) {
+  if (compiler_options.pop({"-nostdinc", "--nostdinc"})) {
     use_builtin_headers = false;
   }
 
   // This is re-added to the remaining options below.
-  bool should_remove_unused_globals = detail::pop_flag(
-      &compiler_options, "-remove-unused-globals", "--remove-unused-globals");
+  bool should_remove_unused_globals = compiler_options.pop(
+      {"-remove-unused-globals", "--remove-unused-globals"});
 
   // Patch all given sources.
   source = detail::patch_cuda_source(source, use_cuda_std, replace_pragma_once);
@@ -6065,9 +6198,9 @@ inline PreprocessedProgram PreprocessedProgram::preprocess(
   struct ArchFlag {
     int cc;
     bool is_virtual;
-    explicit operator std::string() const {
-      return std::string("-arch=") + (is_virtual ? "compute_" : "sm_") +
-             std::to_string(cc);
+    explicit operator Option() const {
+      return Option("-arch",
+                    (is_virtual ? "compute_" : "sm_") + std::to_string(cc));
     }
     bool operator==(const ArchFlag& other) const {
       return cc == other.cc && is_virtual == other.is_virtual;
@@ -6081,10 +6214,10 @@ inline PreprocessedProgram PreprocessedProgram::preprocess(
   std::unordered_set<ArchFlag, ArchFlag::Hash> arch_flags;
   while (true) {
     std::string error;
-    size_t beg_idx, end_idx;
+    size_t idx;
     bool is_virtual = false;
-    int given_cc = detail::parse_arch_flag(compiler_options, &is_virtual,
-                                           &error, &beg_idx, &end_idx);
+    int given_cc =
+        detail::parse_arch_flag(compiler_options, &is_virtual, &error, &idx);
     if (!error.empty()) {
       return Error("Failed to parse architecture flag: " + error);
     }
@@ -6105,8 +6238,7 @@ inline PreprocessedProgram PreprocessedProgram::preprocess(
     }
     arch_flags.insert({given_cc, is_virtual});
     // Remove the parsed arch flag entries; they are replaced below.
-    compiler_options.erase(compiler_options.begin() + beg_idx,
-                           compiler_options.begin() + end_idx);
+    compiler_options.erase(idx);
   }
   if (arch_flags.empty()) {
     // Push a placeholder entry so that preprocessing still runs (with the
@@ -6120,7 +6252,7 @@ inline PreprocessedProgram PreprocessedProgram::preprocess(
   for (const ArchFlag& arch_flag : arch_flags) {
     if (arch_flag.cc) {
       // Temporarily add this arch flag.
-      compiler_options.push_back(static_cast<std::string>(arch_flag));
+      compiler_options.push_back(static_cast<Option>(arch_flag));
     }
 
     std::string compiler_options_msg = detail::string_join(
@@ -6215,7 +6347,7 @@ inline PreprocessedProgram PreprocessedProgram::preprocess(
 
   // Re-add the -remove-unused-globals flag if it was provided.
   if (should_remove_unused_globals) {
-    compiler_options.push_back("-remove-unused-globals");
+    compiler_options.push_back(Option("-remove-unused-globals"));
   }
 
   return PreprocessedProgram(
@@ -6268,8 +6400,8 @@ class ProgramData : public serialization::Serializable<ProgramData> {
    *  \return A PreprocessedProgram object that contains either a valid
    *    PreprocessedProgramData object or an error state.
    */
-  PreprocessedProgram preprocess(StringVec compiler_options = {},
-                                 StringVec linker_options = {},
+  PreprocessedProgram preprocess(OptionsVec compiler_options = {},
+                                 OptionsVec linker_options = {},
                                  FileCallback header_callback = nullptr) const {
     return PreprocessedProgram::preprocess(name_, source_, header_sources_,
                                            compiler_options, linker_options,
@@ -6886,12 +7018,12 @@ class KeyWithUInt64 {
 class AutoKey {
   StringVec name_expressions_;
   StringMap extra_header_sources_;
-  StringVec extra_compiler_options_;
-  StringVec extra_linker_options_;
+  OptionsVec extra_compiler_options_;
+  OptionsVec extra_linker_options_;
 
  public:
   AutoKey(StringVec name_expressions, StringMap extra_header_sources,
-          StringVec extra_compiler_options, StringVec extra_linker_options)
+          OptionsVec extra_compiler_options, OptionsVec extra_linker_options)
       : name_expressions_(std::move(name_expressions)),
         extra_header_sources_(std::move(extra_header_sources)),
         extra_compiler_options_(std::move(extra_compiler_options)),
@@ -6946,8 +7078,9 @@ class AutoKey {
     }
     key_str += '\0';
     for (const StringVec& vec :
-         {key.name_expressions_, key.extra_compiler_options_,
-          key.extra_linker_options_}) {
+         {key.name_expressions_,
+          key.extra_compiler_options_.serialize_canonical(),
+          key.extra_linker_options_.serialize_canonical()}) {
       key_str += std::to_string(vec.size());
       key_str += '\0';
       for (const std::string& str : vec) {
@@ -7001,14 +7134,14 @@ class ProgramCache {
   size_t num_hits_ = 0;
   size_t num_misses_ = 0;
 
-  StringVec merge_compiler_options(StringVec extra_compiler_options) const {
+  OptionsVec merge_compiler_options(OptionsVec extra_compiler_options) const {
     extra_compiler_options.insert(extra_compiler_options.begin(),
                                   preprog_.remaining_compiler_options().begin(),
                                   preprog_.remaining_compiler_options().end());
     return extra_compiler_options;
   }
 
-  StringVec merge_linker_options(StringVec extra_linker_options) const {
+  OptionsVec merge_linker_options(OptionsVec extra_linker_options) const {
     extra_linker_options.insert(extra_linker_options.begin(),
                                 preprog_.remaining_linker_options().begin(),
                                 preprog_.remaining_linker_options().end());
@@ -7029,14 +7162,14 @@ class ProgramCache {
 
   LinkedProgram build_linked_program(const StringVec& name_expressions,
                                      const StringMap& extra_header_sources,
-                                     StringVec extra_compiler_options,
-                                     StringVec extra_linker_options) const {
+                                     OptionsVec extra_compiler_options,
+                                     OptionsVec extra_linker_options) const {
     StringMap tmp_all_header_sources;
     const StringMap& all_header_sources =
         merge_header_sources(extra_header_sources, &tmp_all_header_sources);
-    StringVec all_compiler_options =
+    OptionsVec all_compiler_options =
         merge_compiler_options(extra_compiler_options);
-    StringVec all_linker_options = merge_linker_options(extra_linker_options);
+    OptionsVec all_linker_options = merge_linker_options(extra_linker_options);
 
     auto compiled = CompiledProgram::compile(
         preprog_.name(), preprog_.source(), all_header_sources,
@@ -7108,8 +7241,8 @@ class ProgramCache {
   LoadedProgram get_program(const key_type& key,
                             const StringVec& name_expressions,
                             const StringMap& extra_header_sources = {},
-                            StringVec extra_compiler_options = {},
-                            StringVec extra_linker_options = {}) {
+                            OptionsVec extra_compiler_options = {},
+                            OptionsVec extra_linker_options = {}) {
     // Add the current CUDA context to the key, as modules are context-specific.
     CUcontext context;
     if (!cuda()) return LoadedProgram::Error(cuda().error());
@@ -7129,7 +7262,7 @@ class ProgramCache {
     } else {
       ++num_misses_;
       // Add the SM architecture to the key, as cubins are arch-specific.
-      StringVec all_compiler_options =
+      OptionsVec all_compiler_options =
           merge_compiler_options(extra_compiler_options);
       std::string error;
       bool is_virtual;
@@ -7186,8 +7319,8 @@ class ProgramCache {
                                     int>::type = 0>
   LoadedProgram get_program(const StringVec& name_expressions,
                             const StringMap& extra_header_sources = {},
-                            StringVec extra_compiler_options = {},
-                            StringVec extra_linker_options = {}) {
+                            OptionsVec extra_compiler_options = {},
+                            OptionsVec extra_linker_options = {}) {
     return get_program(
         detail::AutoKey(name_expressions, extra_header_sources,
                         extra_compiler_options, extra_linker_options),
@@ -7220,8 +7353,8 @@ class ProgramCache {
   Kernel get_kernel(const key_type& key, std::string name,
                     StringVec other_name_expressions = {},
                     const StringMap& extra_header_sources = {},
-                    StringVec extra_compiler_options = {},
-                    StringVec extra_linker_options = {}) {
+                    OptionsVec extra_compiler_options = {},
+                    OptionsVec extra_linker_options = {}) {
     other_name_expressions.push_back(name);
     LoadedProgram program = get_program(
         key, other_name_expressions, extra_header_sources,
@@ -7240,8 +7373,8 @@ class ProgramCache {
                                     int>::type = 0>
   Kernel get_kernel(std::string name, StringVec other_name_expressions = {},
                     const StringMap& extra_header_sources = {},
-                    StringVec extra_compiler_options = {},
-                    StringVec extra_linker_options = {}) {
+                    OptionsVec extra_compiler_options = {},
+                    OptionsVec extra_linker_options = {}) {
     other_name_expressions.push_back(name);
     LoadedProgram program = get_program(
         detail::AutoKey(other_name_expressions, extra_header_sources,
diff --git a/jitify2_test.cu b/jitify2_test.cu
index 64a43d0..999a49a 100644
--- a/jitify2_test.cu
+++ b/jitify2_test.cu
@@ -736,9 +736,11 @@ __global__ void my_kernel() {}
   EXPECT_EQ(compiled->lowered_name_map().size(), size_t(1));
   ASSERT_EQ(compiled->lowered_name_map().count(instantiation), size_t(1));
   EXPECT_EQ(compiled->lowered_name_map().at(instantiation), lowered_name);
+  std::vector<std::string> serialized_linker_options =
+      compiled->remaining_linker_options().serialize();
   std::unordered_multiset<std::string> linker_options;
-  linker_options.insert(compiled->remaining_linker_options().begin(),
-                        compiled->remaining_linker_options().end());
+  linker_options.insert(serialized_linker_options.begin(),
+                        serialized_linker_options.end());
   EXPECT_EQ(linker_options.count("-lfoo"), 1);
   EXPECT_EQ(compiled->log(), "");
 }
@@ -1255,6 +1257,70 @@ __global__ void foo_kernel(int* data) {
   CHECK_CUDART(cudaFree(d_data));
 }
 
+TEST(Jitify2Test, OptionsVec) {
+  OptionsVec options0;
+  EXPECT_TRUE(options0.ok());
+  OptionsVec options1({Option("-arch", "sm_50"), Option("-G")});
+  EXPECT_TRUE(options1.ok());
+  StringVec options_sv({"-arch", "sm_50", "-G"});
+  OptionsVec options2(options_sv);
+  EXPECT_TRUE(options2.ok());
+  OptionsVec options3({"-arch", "sm_50", "-G"});
+  EXPECT_TRUE(options3.ok());
+
+  OptionsVec options({"--gpu-architecture", "compute_50", "-arch", "sm_50",
+                      "-maxrregcount=100", "-Ifoo", "-I=foo2", "--device-debug",
+                      "-G", "--restrict", "-restrict", "-lbar", "-l=bar2",
+                      "-lineinfo"});
+  EXPECT_TRUE(options.ok());
+
+  EXPECT_EQ(options.size(), 12);
+  EXPECT_EQ(options.serialize(),
+            StringVec({"--gpu-architecture", "compute_50", "-arch", "sm_50",
+                       "-maxrregcount=100", "-Ifoo", "-I=foo2",
+                       "--device-debug", "-G", "--restrict", "-restrict",
+                       "-lbar", "-l=bar2", "-lineinfo"}));
+  EXPECT_EQ(options.serialize_canonical(),
+            StringVec({"--gpu-architecture=compute_50", "-arch=sm_50",
+                       "-maxrregcount=100", "-I=foo", "-I=foo2",
+                       "--device-debug", "-G", "--restrict", "-restrict",
+                       "-l=bar", "-l=bar2", "-lineinfo"}));
+  StringVec implicit_conversion = options;
+  EXPECT_EQ(implicit_conversion, options.serialize());
+  EXPECT_TRUE(options.pop({"--gpu-architecture", "-arch"}));
+  EXPECT_EQ(options.serialize(),
+            StringVec({"-maxrregcount=100", "-Ifoo", "-I=foo2",
+                       "--device-debug", "-G", "--restrict", "-restrict",
+                       "-lbar", "-l=bar2", "-lineinfo"}));
+  options.push_back(Option("-foo=bar"));
+  EXPECT_EQ(options.serialize(),
+            StringVec({"-maxrregcount=100", "-Ifoo", "-I=foo2",
+                       "--device-debug", "-G", "--restrict", "-restrict",
+                       "-lbar", "-l=bar2", "-lineinfo", "-foo=bar"}));
+  options.pop_back();
+  EXPECT_EQ(options.serialize(),
+            StringVec({"-maxrregcount=100", "-Ifoo", "-I=foo2",
+                       "--device-debug", "-G", "--restrict", "-restrict",
+                       "-lbar", "-l=bar2", "-lineinfo"}));
+  options.erase(1);
+  options.erase(1);
+  EXPECT_EQ(
+      options.serialize(),
+      StringVec({"-maxrregcount=100", "--device-debug", "-G", "--restrict",
+                 "-restrict", "-lbar", "-l=bar2", "-lineinfo"}));
+  OptionsVec extra({"-Ifoo", "-I=foo2"});
+  options.insert(options.begin() + 1, extra.begin(), extra.end());
+  EXPECT_EQ(options.serialize(),
+            StringVec({"-maxrregcount=100", "-Ifoo", "-I=foo2",
+                       "--device-debug", "-G", "--restrict", "-restrict",
+                       "-lbar", "-l=bar2", "-lineinfo"}));
+  EXPECT_EQ(options[2], Option("-I", "foo2"));
+  EXPECT_EQ(options, options);
+  EXPECT_NE(options, extra);
+  std::vector<int> inds = options.find({"--restrict", "-l"});
+  EXPECT_EQ(inds, std::vector<int>({5, 7, 8}));
+}
+
 TEST(Jitify2Test, ArchFlags) {
   static const char* const source = R"(
 const int arch = __CUDA_ARCH__ / 10;
@@ -1334,9 +1400,11 @@ const int arch = __CUDA_ARCH__ / 10;
           ->compile()
           ->link();
   ASSERT_EQ(get_error(linked), "");
+  std::vector<std::string> serialized_linker_options =
+      linked->linker_options().serialize();
   std::unordered_multiset<std::string> linker_options(
-      linked->linker_options().begin(), linked->linker_options().end());
-  EXPECT_EQ(linker_options.count("-maxrregcount=100"), 1);
+      serialized_linker_options.begin(), serialized_linker_options.end());
+  ASSERT_EQ(linker_options.count("-maxrregcount=100"), 1);
   EXPECT_EQ(linker_options.count("--generate-line-info"), 1);
   EXPECT_EQ(linker_options.count("-G"), 1);
 
@@ -1347,9 +1415,10 @@ const int arch = __CUDA_ARCH__ / 10;
                ->compile()
                ->link();
   ASSERT_EQ(get_error(linked), "");
+  serialized_linker_options = linked->linker_options().serialize_canonical();
   linker_options.clear();
-  linker_options.insert(linked->linker_options().begin(),
-                        linked->linker_options().end());
+  linker_options.insert(serialized_linker_options.begin(),
+                        serialized_linker_options.end());
   EXPECT_EQ(linker_options.count("--maxrregcount=100"), 1);
   EXPECT_EQ(linker_options.count("--generate-line-info"), 1);
   EXPECT_EQ(linker_options.count("--device-debug"), 1);

From 7bc924ed3bcf5dcf2e76ba1ac56180ba448551b2 Mon Sep 17 00:00:00 2001
From: Ben Barsdell <bbarsdell@nvidia.com>
Date: Wed, 22 May 2024 19:15:17 +1000
Subject: [PATCH 2/2] Factor out common headers to avoid repetition

---
 jitify2.hpp | 18 ++++--------------
 1 file changed, 4 insertions(+), 14 deletions(-)

diff --git a/jitify2.hpp b/jitify2.hpp
index 04908da..c73cc33 100644
--- a/jitify2.hpp
+++ b/jitify2.hpp
@@ -41,10 +41,6 @@
 #ifndef JITIFY2_HPP_INCLUDE_GUARD
 #define JITIFY2_HPP_INCLUDE_GUARD
 
-// This macro is used by source files generated by jitify_preprocess to avoid
-// unnecessary dependencies.
-#ifdef JITIFY_SERIALIZATION_ONLY
-
 #include <algorithm>
 #include <cassert>
 #include <climits>
@@ -56,6 +52,10 @@
 #include <unordered_map>
 #include <vector>
 
+// This macro is used by source files generated by jitify_preprocess to avoid
+// unnecessary dependencies.
+#ifdef JITIFY_SERIALIZATION_ONLY
+
 #if __cplusplus >= 201703L
 #include <string_view>
 #endif
@@ -102,28 +102,18 @@
 #include <nv_decode.h>  // For __cu_demangle (requires linking with libcufilt.a)
 #endif
 
-#include <algorithm>
-#include <cassert>
 #include <cctype>
-#include <climits>
 #include <cstring>
 #include <fstream>
 #include <functional>
-#include <initializer_list>
 #include <iomanip>
-#include <iostream>
 #include <list>
 #include <map>
 #include <memory>
 #include <regex>
-#include <sstream>
-#include <streambuf>
-#include <string>
 #include <thread>
 #include <type_traits>
-#include <unordered_map>
 #include <unordered_set>
-#include <vector>
 
 #if JITIFY_THREAD_SAFE
 #include <mutex>