implement breadth-first sorting of compute graphs

Author: theoheimel

madspace/include/madspace/madcode/function.h

@@ -53,6 +53,8 @@ class Function {
     ValueVec _locals;
     std::unordered_map<std::string, Value> _globals;
     std::vector<InstructionCall> _instructions;
+
+    friend Function sort_breadth_first(const Function& function);
 };
 
 std::ostream& operator<<(std::ostream& out, const Value& value);

madspace/include/madspace/madcode/optimizer.h

@@ -10,24 +10,27 @@ class InstructionDependencies {
 public:
     InstructionDependencies(const Function& function);
     bool depends(std::size_t test_index, std::size_t dependency_index) {
-        return matrix[test_index * size + dependency_index];
+        return _matrix[test_index * _size + dependency_index];
     }
+    const std::vector<int>& ranks() const { return _ranks; }
 
 private:
-    std::size_t size;
-    std::vector<bool> matrix;
-    std::vector<int> ranks;
+    std::size_t _size;
+    std::vector<bool> _matrix;
+    std::vector<int> _ranks;
 };
 
 class LastUseOfLocals {
 public:
     LastUseOfLocals(const Function& function);
     std::vector<std::size_t>& local_indices(std::size_t index) {
-        return last_used.at(index);
+        return _last_used.at(index);
     }
 
 private:
-    std::vector<std::vector<std::size_t>> last_used;
+    std::vector<std::vector<std::size_t>> _last_used;
 };
 
+Function sort_breadth_first(const Function& function);
+
 } // namespace madspace

madspace/src/madcode/optimizer.cpp

@@ -10,7 +10,7 @@
 using namespace madspace;
 
 InstructionDependencies::InstructionDependencies(const Function& function) :
-    size(function.instructions().size()), matrix(size * size) {
+    _size(function.instructions().size()), _matrix(_size * _size) {
     std::vector<int> local_source(function.locals().size(), -1);
     int index = 0;
     for (auto& instr : function.instructions()) {
@@ -20,26 +20,26 @@ InstructionDependencies::InstructionDependencies(const Function& function) :
             if (source_index == -1) {
                 continue;
             }
-            matrix.at(index * size + source_index) = true;
-            for (int i = 0; i < size; ++i) {
-                matrix.at(index * size + i) =
-                    matrix.at(index * size + i) | matrix.at(source_index * size + i);
+            _matrix.at(index * _size + source_index) = true;
+            for (int i = 0; i < _size; ++i) {
+                _matrix.at(index * _size + i) =
+                    _matrix.at(index * _size + i) | _matrix.at(source_index * _size + i);
             }
-            int source_rank = ranks.at(source_index);
+            int source_rank = _ranks.at(source_index);
             if (rank < source_rank) {
                 rank = source_rank;
             }
         }
         for (auto& output : instr.outputs) {
             local_source.at(output.local_index) = index;
         }
-        ranks.push_back(rank + 1);
+        _ranks.push_back(rank + 1);
         ++index;
     }
 }
 
 LastUseOfLocals::LastUseOfLocals(const Function& function) :
-    last_used(function.instructions().size()) {
+    _last_used(function.instructions().size()) {
     std::vector<bool> seen_locals;
     for (auto& local : function.locals()) {
         seen_locals.push_back(
@@ -50,7 +50,7 @@ LastUseOfLocals::LastUseOfLocals(const Function& function) :
         seen_locals.at(output.local_index) = true;
     }
     auto instr = function.instructions().rbegin();
-    auto indices = last_used.begin();
+    auto indices = _last_used.begin();
     for (; instr != function.instructions().rend(); ++instr, ++indices) {
         for (auto& input : instr->inputs) {
             auto index = input.local_index;
@@ -60,5 +60,22 @@ LastUseOfLocals::LastUseOfLocals(const Function& function) :
             }
         }
     }
-    std::reverse(last_used.begin(), last_used.end());
+    std::reverse(_last_used.begin(), _last_used.end());
+}
+
+Function madspace::sort_breadth_first(const Function& function) {
+    Function func_out = function;
+    InstructionDependencies dependencies(function);
+    auto order = dependencies.ranks();
+    std::vector<std::size_t> instruction_perm(function.instructions().size());
+    std::iota(instruction_perm.begin(), instruction_perm.end(), 0);
+    std::stable_sort(
+        instruction_perm.begin(), instruction_perm.end(),
+        [&](std::size_t i, std::size_t j) { return order.at(i) < order.at(j); }
+    );
+    func_out._instructions.clear();
+    for (std::size_t index : instruction_perm) {
+        func_out._instructions.push_back(function._instructions.at(index));
+    }
+    return func_out;
 }

madspace/src/phasespace/integrand.cpp

@@ -651,13 +651,16 @@ ValueVec MultiChannelIntegrand::build_function_impl(
         common_args.max_weight = args.at(1);
     }
     std::vector<Integrand::ChannelResult> results;
-    for (auto [integrand, chan_size] : zip(_integrands, all_batch_sizes)) {
+    for (std::size_t index = 0; auto [integrand, chan_size] : zip(_integrands, all_batch_sizes)) {
+        fb.set_current_stream(index + 1);
         auto channel_args = common_args;
         channel_args.r = fb.random(chan_size, integrand->_random_dim);
         channel_args.batch_size = chan_size;
         channel_args.has_permutations = integrand->_mapping.channel_count() > 1;
         results.push_back(integrand->build_channel_part(fb, channel_args));
+        ++index;
     }
+    fb.set_current_stream(0);
 
     Integrand::ChannelResult common_results;
     for (std::size_t i = 0; i < common_results.values.size(); ++i) {