Allow specifying a dimension order in ynn_runtime::make_schedule.

This enables reordering loops before scheduling, so when the splits and number of workers are computed this can be taken into account.

PiperOrigin-RevId: 831653398

Author: vksnk

ynnpack/subgraph/dot.cc

@@ -13,6 +13,7 @@
 #include <cstdint>
 #include <cstring>
 #include <memory>
+#include <numeric>
 #include <optional>
 #include <tuple>
 #include <utility>
@@ -1036,23 +1037,23 @@ ynn_status ynn_define_dot(ynn_subgraph_t subgraph, size_t num_k_dims,
       split_n = {};
     }
 
+    std::vector<slinky::index_t> loop_order = {0, 1};
+    if (pack_b && !packed_b.is_static()) {
+      // Loop over n first so we don't redundantly compute the packing for each
+      // split of m.
+      std::swap(loop_order[0], loop_order[1]);
+    }
+
     slinky::expr splits[] = {split_n, split_m};
-    auto sched =
-        runtime.make_schedule(dims, output.buffer, node.outputs[0], splits);
+    auto sched = runtime.make_schedule(dims, output.buffer, node.outputs[0],
+                                       splits, 1, loop_order);
 
     // We want to use exactly these loop splits for two innermost dot loops.
     for (size_t i = 0; i < std::min<std::size_t>(2, sched->loop_splits.size());
          i++) {
       sched->loop_splits[i].step_is_required = true;
     }
 
-    if (pack_b && !packed_b.is_static() && sched->loop_splits.size() >= 2 &&
-        sched->loop_splits[1].axis == 1) {
-      // Loop over n first so we don't redundantly compute the packing for each
-      // split of m.
-      std::swap(sched->loop_splits[0], sched->loop_splits[1]);
-    }
-
     // Schedule the output buffer to be stored at the same level as it's
     // computed at.
     ynn::scheduled_buffer sched_output_buffer = {output.buffer, 0};

ynnpack/subgraph/runtime.cc

@@ -79,7 +79,8 @@ slinky::var ynn_runtime::make_global_variable(slinky::expr value,
 std::unique_ptr<ynn::scheduling_info> ynn_runtime::make_schedule(
     const std::vector<slinky::var>& dims, const slinky::buffer_expr_ptr output,
     uint32_t output_value, slinky::span<const slinky::expr> given_splits,
-    const slinky::expr& element_cost) {
+    const slinky::expr& element_cost,
+    const std::vector<slinky::index_t>& loop_order) {
   auto sched = std::make_unique<ynn::scheduling_info>();
 
   int max_threads = threadpool() ? threadpool()->thread_count() : 1;
@@ -104,7 +105,13 @@ std::unique_ptr<ynn::scheduling_info> ynn_runtime::make_schedule(
                                             output->elem_size() * element_cost);
   std::vector<slinky::expr> splits(rank);
   slinky::expr tile_area_so_far = 1;
-  for (int d = 0; d < rank; ++d) {
+
+  auto get_loop_dim = [&](int index_d) {
+    return index_d < loop_order.size() ? loop_order[index_d] : index_d;
+  };
+
+  for (int index_d = 0; index_d < rank; ++index_d) {
+    int d = get_loop_dim(index_d);
     if (!output_extents[d].defined()) continue;
     if (d < given_splits.size()) {
       splits[d] = given_splits[d];
@@ -124,7 +131,8 @@ std::unique_ptr<ynn::scheduling_info> ynn_runtime::make_schedule(
   std::vector<slinky::expr> workers(rank);
   slinky::expr threads_so_far = 1;
 
-  for (int d = rank - 1; d >= 0; --d) {
+  for (int index_d = rank - 1; index_d >= 0; --index_d) {
+    int d = get_loop_dim(index_d);
     if (max_threads == 1) {
       workers[d] = slinky::loop::serial;
     } else if (output_extents[d].defined() && splits[d].defined()) {
@@ -140,7 +148,8 @@ std::unique_ptr<ynn::scheduling_info> ynn_runtime::make_schedule(
     }
   }
 
-  for (int d = 0; d < rank; ++d) {
+  for (int index_d = 0; index_d < rank; ++index_d) {
+    int d = get_loop_dim(index_d);
     if (output_extents[d].defined() && splits[d].defined()) {
       sched->loop_splits.push_back({dims[d], splits[d], workers[d], d});
     }

ynnpack/subgraph/runtime.h

@@ -77,7 +77,8 @@ struct ynn_runtime {
   std::unique_ptr<ynn::scheduling_info> make_schedule(
       const std::vector<slinky::var>& dims, slinky::buffer_expr_ptr output,
       uint32_t output_value, slinky::span<const slinky::expr> given_splits = {},
-      const slinky::expr& element_cost = 1);
+      const slinky::expr& element_cost = 1,
+      const std::vector<slinky::index_t>& loop_order = {});
 
   slinky::buffer_expr_ptr null_buffer();