Drop ability to zip axii after construction

Author: robertmaynard

docs/benchmarks.md

@@ -259,9 +259,8 @@ using input_types = nvbench::type_list<char, int, unsigned int>;
 using output_types = nvbench::type_list<float, double>;
 NVBENCH_BENCH_TYPES(benchmark, NVBENCH_TYPE_AXES(input_types, output_types))
   .set_type_axes_names({"InputType", "OutputType"})
-  .add_int64_axis("NumInputs", {1000, 10000, 100000, 200000, 200000, 200000})
-  .add_float64_axis("Quality", {0.05, 0.1, 0.25, 0.5, 0.75, 1.})
-  .zip_axes({"NumInputs", "Quality"});
+  .add_zip_axes(nvbench::int64_axis{"NumInputs", {1000, 10000, 100000, 200000, 200000, 200000}},
+                nvbench::float64_axis{"Quality", {0.05, 0.1, 0.25, 0.5, 0.75, 1.}});
 ```
 
 Zipping these two axes reduces the total combinations from 216 to 36, reducing the

nvbench/axes_metadata.cuh

@@ -18,9 +18,9 @@
 
 #pragma once
 
-#include <nvbench/iteration_space_base.cuh>
 #include <nvbench/float64_axis.cuh>
 #include <nvbench/int64_axis.cuh>
+#include <nvbench/iteration_space_base.cuh>
 #include <nvbench/linear_axis_space.cuh>
 #include <nvbench/string_axis.cuh>
 #include <nvbench/type_axis.cuh>
@@ -65,30 +65,29 @@ struct axes_metadata
 
   void add_string_axis(std::string name, std::vector<std::string> data);
 
-  void add_axis(const axis_base& axis);
+  void add_axis(const axis_base &axis);
 
   template <typename... Args>
   void add_zip_axes(Args &&...args)
   {
-    (this->add_axis(std::forward<Args>(args)), ...);
-    this->zip_axes({args.get_name()...});
+    const std::size_t start = this->m_axes.size();
+    const std::size_t count = sizeof...(Args);
+    // (this->add_axis(std::forward<Args>(args)), ...);
+    (m_axes.push_back(args.clone()), ...);
+    this->add_zip_space(start, count);
   }
 
   template <typename... Args>
   void add_user_iteration_axes(
     std::function<nvbench::make_user_space_signature> make,
     Args &&...args)
   {
-    (this->add_axis(std::forward<Args>(args)), ...);
-    this->user_iteration_axes(std::move(make), {args.get_name()...});
+    const std::size_t start = this->m_axes.size();
+    const std::size_t count = sizeof...(Args);
+    (m_axes.push_back(args.clone()), ...);
+    this->add_user_iteration_space(std::move(make), start, count);
   }
 
-  void zip_axes(std::vector<std::string> names);
-
-  void
-  user_iteration_axes(std::function<nvbench::make_user_space_signature> make,
-                      std::vector<std::string> names);
-
   [[nodiscard]] const iteration_space_type &get_type_iteration_space() const
   {
     return m_type_space;
@@ -136,6 +135,12 @@ private:
   std::size_t m_type_axe_count = 0;
   iteration_space_type m_type_space;
   iteration_space_type m_value_space;
+
+  void add_zip_space(std::size_t first_index, std::size_t count);
+  void add_user_iteration_space(
+    std::function<nvbench::make_user_space_signature> make,
+    std::size_t first_index,
+    std::size_t count);
 };
 
 template <typename... TypeAxes>

nvbench/axes_metadata.cxx

@@ -24,6 +24,7 @@
 #include <fmt/ranges.h>
 
 #include <algorithm>
+#include <numeric>
 #include <stdexcept>
 #include <unordered_set>
 
@@ -117,128 +118,45 @@ catch (std::exception &e)
 void axes_metadata::add_float64_axis(std::string name,
                                      std::vector<nvbench::float64_t> data)
 {
-  this->add_axis(nvbench::float64_axis{name,data});
+  this->add_axis(nvbench::float64_axis{name, data});
 }
 
 void axes_metadata::add_int64_axis(std::string name,
                                    std::vector<nvbench::int64_t> data,
                                    nvbench::int64_axis_flags flags)
 {
-  this->add_axis(nvbench::int64_axis{name,data,flags});
+  this->add_axis(nvbench::int64_axis{name, data, flags});
 }
 
 void axes_metadata::add_string_axis(std::string name,
                                     std::vector<std::string> data)
 {
-  this->add_axis(nvbench::string_axis{name,data});
+  this->add_axis(nvbench::string_axis{name, data});
 }
 
-void axes_metadata::add_axis(const axis_base& axis)
+void axes_metadata::add_axis(const axis_base &axis)
 {
   m_value_space.push_back(
     std::make_unique<linear_axis_space>(m_axes.size(),
                                         m_axes.size() - m_type_axe_count));
   m_axes.push_back(axis.clone());
 }
 
-namespace
+void axes_metadata::add_zip_space(std::size_t first_index, std::size_t count)
 {
-std::tuple<std::vector<std::size_t>, std::vector<std::size_t>>
-get_axes_indices(std::size_t type_axe_count,
-                 const nvbench::axes_metadata::axes_type &axes,
-                 const std::vector<std::string> &names)
-{
-  std::vector<std::size_t> input_indices;
-  input_indices.reserve(names.size());
-  for (auto &n : names)
-  {
-    auto iter =
-      std::find_if(axes.cbegin(), axes.cend(), [&n](const auto &axis) {
-        return axis->get_name() == n;
-      });
-
-    // iter distance is input_indices
-    if (iter == axes.cend())
-    {
-      NVBENCH_THROW(std::runtime_error,
-                    "Unable to find the axes named ({}).",
-                    n);
-    }
-    auto index = std::distance(axes.cbegin(), iter);
-    input_indices.push_back(index);
-  }
-
-  std::vector<std::size_t> output_indices = input_indices;
-  for (auto &out : output_indices)
-  {
-    out -= type_axe_count;
-  }
-  return {std::move(input_indices), std::move(output_indices)};
-}
-
-void reset_iteration_space(
-  nvbench::axes_metadata::iteration_space_type &all_spaces,
-  const std::vector<std::size_t> &indices_to_remove)
-{
-  // 1. Find all spaces indices that
-  nvbench::axes_metadata::iteration_space_type reset_space;
-  nvbench::axes_metadata::iteration_space_type to_filter;
-  for (auto &space : all_spaces)
-  {
-    bool added = false;
-    for (auto &i : indices_to_remove)
-    {
-      if (space->contains(i))
-      {
-        // add each item back as linear_axis_space
-        auto as_linear = space->clone_as_linear();
-        to_filter.insert(to_filter.end(),
-                         std::make_move_iterator(as_linear.begin()),
-                         std::make_move_iterator(as_linear.end()));
-        added = true;
-        break;
-      }
-    }
-    if (!added)
-    {
-      // this space doesn't need to be removed
-      reset_space.push_back(std::move(space));
-    }
-  }
-
-  for (auto &iter : to_filter)
-  {
-    bool to_add = true;
-    for (auto &i : indices_to_remove)
-    {
-      if (iter->contains(i))
-      {
-        to_add = false;
-        break;
-      }
-    }
-    if (to_add)
-    {
-      reset_space.push_back(std::move(iter));
-      break;
-    }
-  }
-
-  all_spaces = std::move(reset_space);
-}
-} // namespace
-
-void axes_metadata::zip_axes(std::vector<std::string> names)
-{
-  NVBENCH_THROW_IF((names.size() < 2),
+  NVBENCH_THROW_IF((count < 2),
                    std::runtime_error,
-                   "At least two axi names ( {} provided ) need to be provided "
+                   "At least two axi ( {} provided ) need to be provided "
                    "when using zip_axes.",
-                   names.size());
+                   count);
 
   // compute the numeric indice for each name we have
-  auto [input_indices,
-        output_indices] = get_axes_indices(m_type_axe_count, m_axes, names);
+  std::vector<std::size_t> input_indices(count);
+  std::vector<std::size_t> output_indices(count);
+  std::iota(input_indices.begin(), input_indices.end(), first_index);
+  std::iota(input_indices.begin(),
+            input_indices.end(),
+            first_index - m_type_axe_count);
 
   const auto expected_size = m_axes[input_indices[0]]->get_size();
   for (auto i : input_indices)
@@ -255,22 +173,24 @@ void axes_metadata::zip_axes(std::vector<std::string> names)
                      expected_size);
   }
 
-  // remove any iteration spaces that have axes we need
-  reset_iteration_space(m_value_space, input_indices);
-
   // add the new tied iteration space
   auto tied = std::make_unique<zip_axis_space>(std::move(input_indices),
                                                std::move(output_indices));
   m_value_space.push_back(std::move(tied));
 }
 
-void axes_metadata::user_iteration_axes(
+void axes_metadata::add_user_iteration_space(
   std::function<nvbench::make_user_space_signature> make,
-  std::vector<std::string> names)
+  std::size_t first_index,
+  std::size_t count)
 {
   // compute the numeric indice for each name we have
-  auto [input_indices,
-        output_indices] = get_axes_indices(m_type_axe_count, m_axes, names);
+  std::vector<std::size_t> input_indices(count);
+  std::vector<std::size_t> output_indices(count);
+  std::iota(input_indices.begin(), input_indices.end(), first_index);
+  std::iota(input_indices.begin(),
+            input_indices.end(),
+            first_index - m_type_axe_count);
 
   for (auto i : input_indices)
   {
@@ -281,9 +201,6 @@ void axes_metadata::user_iteration_axes(
                      m_axes[i]->get_name());
   }
 
-  // remove any iteration spaces that have axes we need
-  reset_iteration_space(m_value_space, input_indices);
-
   auto user_func = make(std::move(input_indices), std::move(output_indices));
   m_value_space.push_back(std::move(user_func));
 }

nvbench/benchmark_base.cuh

@@ -280,38 +280,6 @@ struct benchmark_base
   /// @}
 
 protected:
-
-  /// Move existing Axis to being part of zip axis iteration space.
-  /// This will remove any existing iteration spaces that the named axis
-  /// are part of, while restoring all other axis in those spaces to
-  /// the default linear space
-  ///
-  /// This is meant to be used only by the option_parser
-  ///  @{
-  benchmark_base &zip_axes(std::vector<std::string> names)
-  {
-    m_axes.zip_axes(std::move(names));
-    return *this;
-  }
-  /// @}
-
-
-  /// Move existing Axis to being part of user axis iteration space.
-  /// This will remove any existing iteration spaces that the named axis
-  /// are part of, while restoring all other axis in those spaces to
-  /// the default linear space
-  ///
-  /// This is meant to be used only by the option_parser
-  ///  @{
-  benchmark_base &
-  user_iteration_axes(std::function<nvbench::make_user_space_signature> make,
-                      std::vector<std::string> names)
-  {
-    m_axes.user_iteration_axes(std::move(make), std::move(names));
-    return *this;
-  }
-  /// @}
-
   friend struct nvbench::runner_base;
 
   template <typename BenchmarkType>

nvbench/iteration_space_base.cuh

@@ -36,9 +36,17 @@ namespace nvbench
  *
  * * user_axis_space is equivalant to a transform iterator.
  *
- * We don't immediately construct the iterators as the active elements,
- * name, etc can be changed before execution. This class allows for
- * the deferred iterator creation while keeping the meta data insyc.
+ * The `nvbench::axes_metadata` stores all axes in a std::vector. To represent
+ * which axes each space is 'over' we store those indices. We don't store
+ * the pointers or names for the following reasons:
+ *
+ * * The names of an axis can change after being added. The `nvbench::axes_metadata`
+ * is not aware of the name change, and can't inform this class of it.
+ *
+ * * The `nvbench::axes_metadata` can be deep copied, which would invalidate
+ * any pointers held by this class. By holding onto the index we remove the need
+ * to do any form of fixup on deep copies of `nvbench::axes_metadata`.
+ *
  *
  */
 struct iteration_space_base
@@ -65,15 +73,6 @@ struct iteration_space_base
 
   [[nodiscard]] std::unique_ptr<iteration_space_base> clone() const;
 
-   /*!
-   * Returns a vector of linear spaces one for each axi held.
-   * This is required when a iteration_space is removed as we need
-   * to restore all the assoicated axes to default.
-   *
-   */
-  [[nodiscard]] std::vector<std::unique_ptr<iteration_space_base>>
-  clone_as_linear() const;
-
   /*!
    * Returns the iterator over the @a axis provided
    *
@@ -101,24 +100,6 @@ struct iteration_space_base
    */
   [[nodiscard]] std::size_t get_active_count(const axes_type &axes) const;
 
-  /*!
-   * Returns if this space was constructed with the input index specified
-   * by @a input_index.
-   *
-   * The `nvbench::axes_metadata` stores all axes in a std::vector. To represent
-   * which axes each space is 'over' we store those indices. We don't store
-   * the pointers or names for the following reasons:
-   *
-   * * The names of an axis can change after being added. The `nvbench::axes_metadata`
-   * is not aware of the name change, and can't inform this class of it.
-   *
-   * * The `nvbench::axes_metadata` can be deep copied, which would invalidate
-   * any pointers held by this class. By holding onto the index we remove the need
-   * to do any form of fixup on deep copies of `nvbench::axes_metadata`.
-   *
-   */
-  [[nodiscard]] bool contains(std::size_t input_index) const;
-
 protected:
   std::vector<std::size_t> m_input_indices;
   std::vector<std::size_t> m_output_indices;