Fix all broken tests

lib/local-execution/include/local-execution/config.h

@@ -16,7 +16,6 @@
 #ifndef _FLEXFLOW_LOCAL_EXECUTION_CONFIG_H_
 #define _FLEXFLOW_LOCAL_EXECUTION_CONFIG_H_
 
-#include "op-attrs/param_sync.dtg.h"
 #include "utils/fmt.h"
 #include "utils/visitable.h"
 #include <cstring>
@@ -35,12 +34,6 @@ enum class ComputationMode {
 #define MAP_TO_FB_MEMORY 0xABCD0000
 #define MAP_TO_ZC_MEMORY 0xABCE0000
 
-#ifdef FF_USE_NCCL
-constexpr ParamSync CHOSEN_SYNC_TYPE = ParamSync::NCCL;
-#else
-constexpr ParamSync CHOSEN_SYNC_TYPE = ParamSync::PS;
-#endif
-
 struct FFInitInfo : public use_visitable_cmp<FFInitInfo> {
   size_t workSpaceSize;
   bool allowTensorOpMathConversion;

lib/local-execution/src/local_cost_estimator.cc

@@ -4,9 +4,12 @@
 #include "local-execution/tracked_allocator.h"
 #include "op-attrs/computation_graph_op_attrs.h"
 #include "op-attrs/pcg_operator_attrs.h"
+#include "op-attrs/shape_inference.h"
+#include "pcg/computation_graph.h"
 #include "pcg/computation_graph_builder.h"
 #include "pcg/machine_view.dtg.h"
 #include "pcg/parallel_tensor_attrs.h"
+#include "utils/containers/get_only.h"
 #include "utils/containers/transform.h"
 
 namespace FlexFlow {
@@ -48,11 +51,10 @@ CostDetails LocalCostEstimator::estimate_cost(
   TensorBackingMap tensor_backing_map;
   std::vector<tensor_guid_t> input_tensor_ids;
 
-  ComputationGraphBuilder cg_builder;
+  ComputationGraph cg = make_empty_computation_graph();
   for (ParallelTensorShape const &input : inputs) {
     TensorShape tensor_shape = get_piece_shape(input);
-    tensor_guid_t tensor_id =
-        cg_builder.create_input(tensor_shape, CreateGrad::YES);
+    tensor_guid_t tensor_id = get_only(add_input_layer(cg, tensor_shape).outputs);
     GenericTensorAccessorW tensor_backing =
         allocator.allocate_tensor(tensor_shape);
     tensor_backing_map.insert({tensor_id, tensor_backing});
@@ -67,17 +69,29 @@ CostDetails LocalCostEstimator::estimate_cost(
       };
 
   // add operator to graph
+  std::vector<TensorShape> weight_shapes = get_weight_shapes(layer_attrs.op_attrs, 
+                                                             transform(inputs, get_piece_shape));
+
+  std::vector<tensor_guid_t> weight_tensor_ids = transform(weight_shapes, 
+                                                           [&](TensorShape const &tensor_shape) {
+                                                             LayerAttrs attrs = LayerAttrs{
+                                                               ComputationGraphOpAttrs{
+                                                                 WeightAttrs{
+                                                                   /*tensor_shape=*/tensor_shape,
+                                                                   /*initializer=*/InitializerAttrs{ZeroInitializerAttrs{}},
+                                                                 },
+                                                               },
+                                                               /*name=*/std::nullopt,
+                                                             };
+
+                                                             return get_only(add_layer(cg, attrs, /*inputs=*/{}, /*weights=*/{}).outputs);
+                                                           });
+
   std::vector<tensor_guid_t> output_tensor_ids =
-      cg_builder.add_layer(layer_attrs,
-                           input_tensor_ids,
-                           transform(get_vector_piece_attrs(weights),
-                                     [&](TensorAttrs const &a) {
-                                       return cg_builder.create_weight(a);
-                                     }),
-                           get_vector_piece_attrs(outputs));
+    add_layer(cg, layer_attrs, /*inputs=*/input_tensor_ids, /*weights=*/weight_tensor_ids).outputs;
 
   LocalTrainingBacking local_backing(allocator,
-                                     cg_builder.computation_graph,
+                                     cg,
                                      tensor_backing_map,
                                      this->runtime_arg_config);
 

lib/local-execution/src/local_slots_backing.cc

@@ -36,7 +36,7 @@ void LocalSlotsBacking::allocate_outgoing_tensors(
     }
 
     // gradient tensor allocation
-    if (tensor_attrs.create_gradients == CreateGrad::YES &&
+    if (tensor_attrs.create_grad == CreateGrad::YES &&
         !is_gradient_tensor_allocated(output_tensor)) {
       GenericTensorAccessorW gradient_tensor_backing =
           allocator.allocate_tensor(tensor_attrs.shape);

lib/local-execution/src/local_training_backing.cc

@@ -17,7 +17,7 @@ LocalTrainingBacking::LocalTrainingBacking(
 
   for (layer_guid_t const &node : topological_ordering(computation_graph)) {
     ComputationGraphOpAttrs attrs =
-        get_layer_attrs(computation_graph, node).attrs;
+        get_layer_attrs(computation_graph, node).op_attrs;
 
     // allocate outgoing tensors
     this->local_slots_backing.allocate_outgoing_tensors(
@@ -53,7 +53,7 @@ void LocalTrainingBacking::execute_init() {
        topological_ordering(this->computation_graph)) {
     if (this->task_registry.init_task_ids.at(operator_node).has_value()) {
       ComputationGraphOpAttrs attrs =
-          get_layer_attrs(this->computation_graph, operator_node).attrs;
+          get_layer_attrs(this->computation_graph, operator_node).op_attrs;
 
       OpTaskInvocation invocation = init(attrs);
       TaskArgumentAccessor accessor =
@@ -72,7 +72,7 @@ PerLayerElapsedTime LocalTrainingBacking::execute_forward() {
        topological_ordering(this->computation_graph)) {
     if (this->task_registry.forward_task_ids.at(operator_node).has_value()) {
       ComputationGraphOpAttrs attrs =
-          get_layer_attrs(this->computation_graph, operator_node).attrs;
+          get_layer_attrs(this->computation_graph, operator_node).op_attrs;
 
       OpTaskInvocation invocation = forward(attrs);
       TaskArgumentAccessor accessor =
@@ -91,7 +91,7 @@ PerLayerElapsedTime LocalTrainingBacking::execute_backward() {
        reversed(topological_ordering(this->computation_graph))) {
     if (this->task_registry.backward_task_ids.at(operator_node).has_value()) {
       ComputationGraphOpAttrs attrs =
-          get_layer_attrs(this->computation_graph, operator_node).attrs;
+          get_layer_attrs(this->computation_graph, operator_node).op_attrs;
 
       OpTaskInvocation invocation = backward(attrs);
       TaskArgumentAccessor accessor =

lib/local-execution/src/ops/batch_matmul.cc

@@ -16,7 +16,6 @@
 #include "batch_matmul.h"
 #include "kernels/batch_matmul_kernels.h"
 #include "local-execution/op_task_signature.h"
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/ops/batch_matmul.h"
 #include "utils/containers/transform.h"
 #include "utils/nonnegative_int/nonnegative_range.h"

lib/local-execution/src/ops/concat.cc

@@ -15,10 +15,8 @@
 
 #include "concat.h"
 #include "kernels/concat_kernels.h"
-
 #include "local-execution/op_task_signature.h"
 #include "local-execution/variadic_tensor_ref.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/hash-utils.h"
 
 namespace FlexFlow {

lib/local-execution/src/ops/conv_2d.cc

@@ -1,6 +1,5 @@
 #include "conv_2d.h"
 #include "kernels/conv_2d_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 
 namespace FlexFlow {
 

lib/local-execution/src/ops/dropout.cc

@@ -2,7 +2,6 @@
 #include "kernels/dropout_kernels.h"
 #include "local-execution/op_task_invocation.h"
 #include "local-execution/op_task_signature.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/hash-utils.h"
 
 namespace FlexFlow {

lib/local-execution/src/ops/element_binary.cc

@@ -1,7 +1,6 @@
 #include "element_binary.h"
 #include "kernels/element_binary_kernels.h"
 #include "local-execution/task_signature_impl.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/hash-utils.h"
 
 namespace FlexFlow {

lib/local-execution/src/ops/element_unary.cc

@@ -1,6 +1,5 @@
 #include "element_unary.h"
 #include "kernels/element_unary_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/parallel_tensor_shape.h"
 #include "utils/hash-utils.h"
 

lib/local-execution/src/ops/flat.cc

@@ -1,6 +1,5 @@
 #include "flat.h"
 #include "kernels/flat_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 
 namespace FlexFlow {
 

lib/local-execution/src/ops/gather.cc

@@ -15,7 +15,6 @@
 
 #include "gather.h"
 #include "kernels/gather_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/nonnegative_int/nonnegative_range.h"
 #include <optional>
 

lib/local-execution/src/ops/layer_norm.cc

@@ -15,7 +15,6 @@
 
 #include "layer_norm.h"
 #include "kernels/layer_norm_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/ops/layer_norm.h"
 #include "op-attrs/parallel_tensor_shape.h"
 #include "utils/exception.h"

lib/local-execution/src/ops/linear.cc

@@ -2,7 +2,6 @@
 #include "kernels/linear_kernels.h"
 #include "local-execution/task_argument_accessor.h"
 #include "op-attrs/ff_dim_t.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"
 

lib/local-execution/src/ops/pool_2d.cc

@@ -1,7 +1,6 @@
 #include "pool_2d.h"
 #include "kernels/pool_2d_kernels.h"
 
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/ops/pool_2d.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"

lib/local-execution/src/ops/reduce.cc

@@ -1,7 +1,6 @@
 #include "reduce.h"
 #include "kernels/reduce_kernels.h"
 
-#include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"
 #include "utils/type_traits_core.h"

lib/local-execution/src/ops/reduction.cc

@@ -15,7 +15,6 @@
 
 #include "reduction.h"
 #include "kernels/reduction_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"
 

lib/local-execution/src/ops/repartition.cc

@@ -15,7 +15,6 @@
 
 #include "repartition.h"
 #include "kernels/partition_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"
 

lib/local-execution/src/ops/replicate.cc

@@ -15,7 +15,6 @@
 
 #include "replicate.h"
 #include "kernels/replicate_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/parallel_tensor_shape.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"

lib/local-execution/src/ops/reshape.cc

@@ -15,7 +15,6 @@
 
 #include "reshape.h"
 #include "kernels/reshape_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 
 namespace FlexFlow {
 

lib/local-execution/src/ops/reverse.cc

@@ -16,7 +16,6 @@
 #include "reverse.h"
 #include "kernels/accessor.h"
 #include "kernels/reverse_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/nonnegative_int/nonnegative_range.h"
 
 namespace FlexFlow {

lib/local-execution/src/ops/softmax.cc

@@ -15,7 +15,6 @@
 
 #include "softmax.h"
 #include "kernels/softmax_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/parallel_tensor_shape.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"

lib/local-execution/src/ops/split.cc

@@ -16,7 +16,6 @@
 #include "split.h"
 #include "kernels/array_shape.h"
 #include "kernels/split_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"
 #include "utils/nonnegative_int/nonnegative_range.h"

lib/local-execution/src/ops/topk.cc

@@ -15,7 +15,6 @@
 
 #include "topk.h"
 #include "kernels/topk_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 
 namespace FlexFlow {

lib/local-execution/src/ops/transpose.cc

@@ -15,7 +15,6 @@
 
 #include "transpose.h"
 #include "kernels/transpose_kernels.h"
-#include "op-attrs/get_output_shapes.h"
 #include "op-attrs/ops/transpose.h"
 #include "utils/integer_conversions.h"
 

lib/local-execution/test/src/test_local_slots_backing.cc

@@ -187,7 +187,7 @@ TEST_SUITE(FF_TEST_SUITE) {
       };
       MultiHeadAttentionAttrs attrs =
           get_layer_attrs(cg_builder.computation_graph, layer_guid)
-              .attrs.get<MultiHeadAttentionAttrs>();
+              .op_attrs.get<MultiHeadAttentionAttrs>();
       OpTaskBinding binding = [&] {
         OpTaskBinding b;
         b.bind(QUERY, input_tensor(0));

lib/op-attrs/src/op-attrs/shape_inference.cc

@@ -239,6 +239,9 @@ std::vector<ParallelTensorShape>
       [&](Pool2DAttrs const &attrs) -> std::vector<ParallelTensorShape> {
         return {throw_if_unexpected(get_output_shape(attrs, get_only(input_shapes)))};
       },
+      [&](ReductionAttrs const &attrs) -> std::vector<ParallelTensorShape> {
+        return {throw_if_unexpected(get_output_shape(attrs, get_only(input_shapes)))};
+      },
       [&](RepartitionAttrs const &attrs) -> std::vector<ParallelTensorShape> {
         return {throw_if_unexpected(get_output_shape(attrs, get_only(input_shapes)))};
       },

lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_builder.h

@@ -135,14 +135,12 @@ struct ParallelComputationGraphBuilder {
                       nonnegative_int degree,
                       std::optional<std::string> const &name = std::nullopt);
 
+  ParallelTensorShape get_shape(parallel_tensor_guid_t const &) const;
 private:
   parallel_tensor_guid_t as_type(parallel_tensor_guid_t const &,
                                  DataType,
                                  std::string const &name);
 
-private:
-  ParallelTensorShape get_shape(parallel_tensor_guid_t const &) const;
-
 private:
   std::vector<parallel_tensor_guid_t>
       add_layer(ParallelLayerAttrs const &layer,

lib/pcg/test/src/pcg/computation_graph.cc

@@ -200,7 +200,7 @@ TEST_SUITE(FF_TEST_SUITE) {
       LayerAddedResult bias_weight_added = add_layer(cg, make_layer_attrs(bias_weight_attrs), {}, {});
       tensor_guid_t t_bias_weight = get_only(bias_weight_added.outputs);
 
-      LayerAddedResult linear_added = add_layer(cg, make_layer_attrs(linear_attrs), {}, {});
+      LayerAddedResult linear_added = add_layer(cg, make_layer_attrs(linear_attrs), {t_input}, {t_projection_weight, t_bias_weight});
 
       std::vector<tensor_guid_t> result = get_incoming_weights(cg, linear_added.layer);
       std::vector<tensor_guid_t> correct = {