flexflow · lockshaw · May 15, 2026 · May 14, 2026
diff --git a/bin/export-model-arch/src/export-model-arch/main.cc b/bin/export-model-arch/src/export-model-arch/main.cc
@@ -14,7 +14,7 @@
 #include "utils/cli/cli_parse.h"
 #include "utils/cli/cli_parse_result.h"
 #include "utils/cli/cli_spec.h"
-#include "utils/graph/open_dataflow_graph/algorithms/as_dot.h"
+#include "utils/graph/open_dataflow_graph/algorithms/open_dataflow_graph_as_dot.h"
 #include "utils/graph/series_parallel/binary_sp_decomposition_tree/right_associative_binary_sp_tree_from_nary.h"
 #include "utils/graph/series_parallel/get_series_parallel_decomposition.h"
 

diff --git a/bin/substitution-to-dot/src/substitution-to-dot/main.cc b/bin/substitution-to-dot/src/substitution-to-dot/main.cc
@@ -1,5 +1,5 @@
 #include "substitution-generator/legacy_rules.h"
-#include "utils/dot_file.h"
+#include "utils/dot/dot_file.h"
 #include <cassert>
 #include <iostream>
 

diff --git a/flake.nix b/flake.nix
@@ -5,10 +5,10 @@
     bash-prompt-prefix = "(ff) ";
     extra-substituters = [
       "https://ff.cachix.org"
-      "https://cuda-maintainers.cachix.org/"
+      #"https://cuda-maintainers.cachix.org/"
     ];
     extra-trusted-public-keys = [
-      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
+      #"cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
       "ff.cachix.org-1:IRdsNEnht4YKGUasP6SX5DfpaOTBckhpJDEODz7wMFM="
     ];
   };
@@ -119,6 +119,7 @@
               gbenchmark
               libtorch-bin
               graphviz # for documentation
+              texliveBasic # for documentation
             ])
             (with proj-repo.packages.${system}; [
               proj
@@ -162,6 +163,8 @@
               jq
               gh
               expect
+              universal-ctags
+              ninja
             ])
             (with pkgs.python3Packages; [
               gitpython

diff --git a/lib/compiler/include/compiler/unity_algorithm/graph_optimize_state.h b/lib/compiler/include/compiler/unity_algorithm/graph_optimize_state.h
@@ -26,15 +26,6 @@ struct GraphOptimizeState {
 std::string format_as(GraphOptimizeState const &);
 std::ostream &operator<<(std::ostream &, GraphOptimizeState const &);
 
-// TODO(@lockshaw)(#pr): Delete this if still unused
-// std::optional<GraphOptimizeState>
-//   graph_optimize_state_from_machine_mapping_result(ParallelComputationGraph
-//   const &,
-//                                                    PCGBinarySPDecomposition
-//                                                    const &,
-//                                                    MachineMappingResult const
-//                                                    &);
-
 } // namespace FlexFlow
 
 namespace std {

diff --git a/...mapping/abstracted_tensor_set_movement/get_abstracted_tensor_set_movement_across_split.cc b/...mapping/abstracted_tensor_set_movement/get_abstracted_tensor_set_movement_across_split.cc
@@ -101,13 +101,13 @@ AbstractedTensorSetMovement get_abstracted_tensor_set_movement_across_split(
   };
 
   return AbstractedTensorSetMovement{
-      transform(
-          edges_by_tensor.right_groups(),
-          [&](std::unordered_set<ParallelComputationGraphEdge> const &edges) {
-            return merge_abstracted_single_tensor_movements(
-                transform(unordered_multiset_of(edges),
-                          to_abstracted_single_tensor_movement));
-          }),
+      transform(edges_by_tensor.right_groups(),
+                [&](nonempty_unordered_set<ParallelComputationGraphEdge> const
+                        &edges) {
+                  return merge_abstracted_single_tensor_movements(transform(
+                      unordered_multiset_of(edges.unwrap_as_unordered_set()),
+                      to_abstracted_single_tensor_movement));
+                }),
   };
 }
 

diff --git a/lib/compiler/src/compiler/machine_mapping/machine_mapping.cc b/lib/compiler/src/compiler/machine_mapping/machine_mapping.cc
@@ -2,6 +2,9 @@
 #include "compiler/machine_mapping/machine_view.h"
 #include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.h"
 #include "op-attrs/computation_graph_op_attrs.h"
+#include "op-attrs/pcg_operator_attrs.h"
+#include "pcg/machine_compute_resource_slice.h"
+#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.h"
 #include "utils/bidict/algorithms/bidict_from_map.h"
 #include "utils/containers/are_disjoint.h"
 #include "utils/containers/binary_merge_disjoint_maps.h"
@@ -15,30 +18,32 @@ MappedParallelComputationGraph
 
   std::unordered_set<parallel_layer_guid_t> pcg_layers =
       get_parallel_layers(pcg);
+
   std::unordered_set<parallel_layer_guid_t> mapped_layers =
       keys(mapping.machine_views);
-  ASSERT(pcg_layers == mapped_layers);
-
-  return MappedParallelComputationGraph{
-      /*pcg=*/pcg,
-      /*mapped_tasks=*/
-      generate_map(
-          get_parallel_layers(pcg),
-          [&](parallel_layer_guid_t l) -> MappedOperatorTaskGroup {
-            ComputationGraphOpAttrs op_attrs =
-                compgraph_op_attrs_from_pcg_op_attrs(pcg_get_op_attrs(pcg, l))
-                    .value();
-
-            std::unordered_map<TensorSlotName, ParallelTensorDimDegrees>
-                inputs_dim_degrees = get_incoming_input_degrees(pcg, l);
-
-            ASSERT(contains_key(mapping.machine_views, l));
-            MachineView machine_view = mapping.machine_views.at(l);
-
-            return mapped_operator_task_group_from_machine_view(
-                op_attrs, inputs_dim_degrees, machine_view);
-          }),
+
+  ASSERT(mapped_layers == pcg_layers);
+
+  auto mapping_for_layer =
+      [&](parallel_layer_guid_t l) -> MappedOperatorTaskGroup {
+    ComputationGraphOpAttrs op_attrs = assert_unwrap(
+        compgraph_op_attrs_from_pcg_op_attrs(pcg_get_op_attrs(pcg, l)));
+
+    std::unordered_map<TensorSlotName, ParallelTensorDimDegrees>
+        inputs_dim_degrees = get_incoming_input_degrees(pcg, l);
+
+    ASSERT(contains_key(mapping.machine_views, l));
+    MachineView machine_view = mapping.machine_views.at(l);
+
+    return mapped_operator_task_group_from_machine_view(
+        op_attrs, inputs_dim_degrees, machine_view);
   };
+
+  std::unordered_map<parallel_layer_guid_t, MappedOperatorTaskGroup>
+      mapped_op_task_groups = generate_map(mapped_layers, mapping_for_layer);
+
+  return mapped_pcg_from_pcg_and_mapped_op_task_groups(pcg,
+                                                       mapped_op_task_groups);
 }
 
 MachineMapping combine_disjoint_mappings(MachineMapping const &m1,

diff --git a/lib/compiler/src/compiler/search_result.cc b/lib/compiler/src/compiler/search_result.cc
@@ -10,7 +10,7 @@ MappedParallelComputationGraph
 
 std::string format_as(SearchResult const &r) {
   return fmt::format("<SearchResult\npcg={}\nmachine_mapping={}>",
-                     as_dot(r.pcg),
+                     pcg_as_dot(r.pcg),
                      r.machine_mapping);
 }
 

diff --git a/.../series_parallel/computation_graph/get_computation_graph_series_parallel_decomposition.cc b/.../series_parallel/computation_graph/get_computation_graph_series_parallel_decomposition.cc
@@ -35,23 +35,15 @@ std::string render_preprocessed_computation_graph_for_sp_decomposition(
     preprocessed_digraph.add_edge(DirectedEdge{fake_node, dst.raw_node});
   }
 
-  std::function<std::string(Node const &)> get_node_label =
-      [&](Node const &n) -> std::string {
+  std::function<nlohmann::json(Node const &)> get_node_label =
+      [&](Node const &n) -> nlohmann::json {
     if (n == fake_node) {
       return "FAKE";
     }
-    LayerAttrs a = cg.raw_graph.at(n);
-    RecordFormatter r = as_dot(a.op_attrs);
 
-    if (a.name.has_value()) {
-      RecordFormatter rr;
-      rr << "Name" << a.name.value();
-      r << rr;
-    }
+    nlohmann::json result = cg.raw_graph.at(n);
 
-    std::ostringstream oss;
-    oss << r;
-    return oss.str();
+    return result;
   };
   std::string preprocessed_dot = digraph_as_dot(
       transitive_reduction(preprocessed_digraph), get_node_label);

diff --git a/lib/compiler/src/compiler/unity_algorithm/graph_optimize_state.cc b/lib/compiler/src/compiler/unity_algorithm/graph_optimize_state.cc
@@ -84,71 +84,13 @@ bool GraphOptimizeState::operator<(GraphOptimizeState const &other) const {
 
 std::string format_as(GraphOptimizeState const &s) {
   return fmt::format(
-      "<SearchResult\nruntime={}\npcg={}>", s.runtime, as_dot(s.pcg));
+      "<SearchResult\nruntime={}\npcg={}>", s.runtime, pcg_as_dot(s.pcg));
 }
 
 std::ostream &operator<<(std::ostream &s, GraphOptimizeState const &x) {
   return (s << fmt::to_string(x));
 }
 
-// TODO(@lockshaw)(#pr): Delete this if still unused
-// std::optional<GraphOptimizeState>
-//   graph_optimize_state_from_machine_mapping_result(ParallelComputationGraph
-//   const &pcg,
-//                                                    PCGBinarySPDecomposition
-//                                                    const
-//                                                    &binary_sp_decomposition,
-//                                                    MachineMappingResult const
-//                                                    &machine_mapping_result) {
-//
-//   FeasibleMachineMappingResult feasible_mapping = ({
-//     if (is_infeasible(machine_mapping_result)) {
-//       return std::nullopt;
-//     }
-//
-//     require_feasible(machine_mapping_result);
-//   });
-//
-//   bidict<BinaryTreePath, parallel_layer_guid_t> path_to_leaf_map =
-//     bidict_from_map(pcg_sp_tree_get_path_to_leaf_map(binary_sp_decomposition));
-//
-//   std::unordered_map<BinaryTreePath, MappedOperatorTaskGroup>
-//     mapped_tasks_by_path = zip_values_strict_with(
-//         path_to_leaf_map.as_unordered_map(),
-//         feasible_mapping.machine_mapping.raw_mapping,
-//         [&](parallel_layer_guid_t const &layer_guid, MachineView const &mv)
-//           -> MappedOperatorTaskGroup
-//         {
-//           ComputationGraphOpAttrs comp_graph_op_attrs =
-//             assert_unwrap(compgraph_op_attrs_from_pcg_op_attrs(pcg_get_op_attrs(pcg,
-//             layer_guid)));
-//
-//           return mapped_operator_task_group_from_machine_view(
-//             comp_graph_op_attrs,
-//             get_incoming_input_degrees(pcg, layer_guid),
-//             mv);
-//         });
-//
-//   std::unordered_map<parallel_layer_guid_t, MappedOperatorTaskGroup>
-//     mapped_tasks = map_keys(mapped_tasks_by_path,
-//                             [&](BinaryTreePath const &path) ->
-//                             parallel_layer_guid_t {
-//                               return path_to_leaf_map.at_l(path);
-//                             });
-//
-//   GraphOptimizeResult result = GraphOptimizeResult{
-//     MappedParallelComputationGraph{
-//       pcg,
-//       mapped_tasks,
-//     },
-//   };
-//
-//   return GraphOptimizeState{
-//     result,
-//     feasible_mapping.runtime,
-//   };
-// }
-
 } // namespace FlexFlow
 
 namespace std {

diff --git a/lib/op-attrs/include/op-attrs/computation_graph_op_attrs.h b/lib/op-attrs/include/op-attrs/computation_graph_op_attrs.h
@@ -3,12 +3,13 @@
 
 #include "op-attrs/computation_graph_op_attrs.dtg.h"
 #include "op-attrs/pcg_operator_attrs.dtg.h"
-#include "utils/record_formatter.h"
 
 namespace FlexFlow {
 
 OperatorType get_op_type(ComputationGraphOpAttrs const &);
-RecordFormatter as_dot(ComputationGraphOpAttrs const &);
+
+nlohmann::json cg_op_attrs_as_dot_json(ComputationGraphOpAttrs const &);
+
 std::optional<ComputationGraphOpAttrs>
     compgraph_op_attrs_from_pcg_op_attrs(PCGOperatorAttrs const &);
 

diff --git a/lib/op-attrs/include/op-attrs/ops/broadcast.h b/lib/op-attrs/include/op-attrs/ops/broadcast.h
@@ -9,8 +9,6 @@
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(BroadcastAttrs const &);
-
 tl::expected<TensorShape, std::string> get_output_shape(BroadcastAttrs const &,
                                                         TensorShape const &);
 ParallelTensorShape get_output_shape(BroadcastAttrs const &,

diff --git a/lib/op-attrs/include/op-attrs/ops/cast.h b/lib/op-attrs/include/op-attrs/ops/cast.h
@@ -9,8 +9,6 @@
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(CastAttrs const &);
-
 tl::expected<TensorShape, std::string> get_output_shape(CastAttrs const &,
                                                         TensorShape const &);
 

diff --git a/lib/op-attrs/include/op-attrs/ops/combine.h b/lib/op-attrs/include/op-attrs/ops/combine.h
@@ -3,13 +3,10 @@
 
 #include "op-attrs/ops/combine_attrs.dtg.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
-#include "utils/record_formatter.h"
 #include <tl/expected.hpp>
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(CombineAttrs const &);
-
 tl::expected<ParallelTensorShape, std::string>
     get_output_shape(CombineAttrs const &, ParallelTensorShape const &);
 

diff --git a/lib/op-attrs/include/op-attrs/ops/embedding.h b/lib/op-attrs/include/op-attrs/ops/embedding.h
@@ -11,8 +11,6 @@
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(EmbeddingAttrs const &);
-
 tl::expected<TensorShape, std::string> get_output_shape(EmbeddingAttrs const &,
                                                         TensorShape const &);
 tl::expected<TensorShape, std::string> get_weights_shape(EmbeddingAttrs const &,

diff --git a/lib/op-attrs/include/op-attrs/ops/linear.h b/lib/op-attrs/include/op-attrs/ops/linear.h
@@ -20,8 +20,6 @@ namespace FlexFlow {
 std::unordered_map<TensorSlotName, IncomingTensorRole>
     get_linear_incoming_tensor_roles(LinearAttrs const &);
 
-RecordFormatter as_dot(LinearAttrs const &);
-
 tl::expected<TensorShape, std::string>
     get_projection_shape(LinearAttrs const &attrs, TensorShape const &input);
 tl::expected<TensorShape, std::string> get_bias_shape(LinearAttrs const &attrs,

diff --git a/lib/op-attrs/include/op-attrs/ops/reduction.h b/lib/op-attrs/include/op-attrs/ops/reduction.h
@@ -3,13 +3,10 @@
 
 #include "op-attrs/ops/reduction_attrs.dtg.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
-#include "utils/record_formatter.h"
 #include <tl/expected.hpp>
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(ReductionAttrs const &);
-
 tl::expected<ParallelTensorShape, std::string>
     get_output_shape(ReductionAttrs const &attrs,
                      ParallelTensorShape const &input_shape);

diff --git a/lib/op-attrs/include/op-attrs/ops/repartition.h b/lib/op-attrs/include/op-attrs/ops/repartition.h
@@ -3,13 +3,10 @@
 
 #include "op-attrs/ops/repartition_attrs.dtg.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
-#include "utils/record_formatter.h"
 #include <tl/expected.hpp>
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(RepartitionAttrs const &);
-
 tl::expected<ParallelTensorShape, std::string>
     get_output_shape(RepartitionAttrs const &,
                      ParallelTensorShape const &input_shape);

diff --git a/lib/op-attrs/include/op-attrs/ops/replicate.h b/lib/op-attrs/include/op-attrs/ops/replicate.h
@@ -3,12 +3,9 @@
 
 #include "op-attrs/ops/replicate_attrs.dtg.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
-#include "utils/record_formatter.h"
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(ReplicateAttrs const &);
-
 ParallelTensorShape get_output_shape(ReplicateAttrs const &attrs,
                                      ParallelTensorShape const &input_shape);
 

diff --git a/lib/op-attrs/include/op-attrs/ops/weight.h b/lib/op-attrs/include/op-attrs/ops/weight.h
@@ -10,8 +10,6 @@
 
 namespace FlexFlow {
 
-RecordFormatter as_dot(WeightAttrs const &);
-
 TensorShape get_output_shape(WeightAttrs const &);
 ParallelTensorShape get_output_parallel_tensor_shape(WeightAttrs const &);
 

diff --git a/lib/op-attrs/include/op-attrs/pcg_operator_attrs.h b/lib/op-attrs/include/op-attrs/pcg_operator_attrs.h
@@ -4,14 +4,16 @@
 #include "op-attrs/computation_graph_op_attrs.dtg.h"
 #include "op-attrs/computation_graph_op_attrs.h"
 #include "op-attrs/pcg_operator_attrs.dtg.h"
+#include <nlohmann/json.hpp>
 
 namespace FlexFlow {
 
 bool is_parallel_op(PCGOperatorAttrs const &);
-OperatorType get_op_type(PCGOperatorAttrs const &);
+OperatorType pcg_op_attrs_get_op_type(PCGOperatorAttrs const &);
 PCGOperatorAttrs
     pcg_op_attrs_from_compgraph_op_attrs(ComputationGraphOpAttrs const &);
-RecordFormatter as_dot(PCGOperatorAttrs const &);
+
+nlohmann::json pcg_op_attrs_as_dot_json(PCGOperatorAttrs const &);
 
 } // namespace FlexFlow