Expose op folders publicly

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

@@ -19,8 +19,6 @@ struct ModelTrainingInstance {
                         tensor_guid_t const & label_tensor,
                         OptimizerAttrs const &);
 
-  void register_and_allocate_layers();
-  void allocate_optimizer_tensors();
   void execute_init();
   PerLayerElapsedTime execute_forward();
   PerLayerElapsedTime execute_backward();

lib/local-execution/src/model_training_instance.cc

@@ -14,16 +14,11 @@ ModelTrainingInstance::ModelTrainingInstance(Allocator const & allocator,
                                              tensor_guid_t const &label_tensor, 
                                              OptimizerAttrs const & optimizer_attrs) 
   : computation_graph(computation_graph), training_backing(allocator, computation_graph, tensor_backing_map, runtime_arg_config),
-  loss_attrs(loss_attrs), logit_tensor(logit_tensor), label_tensor(label_tensor), optimizer_attrs(optimizer_attrs) {}
+  loss_attrs(loss_attrs), logit_tensor(logit_tensor), label_tensor(label_tensor), optimizer_attrs(optimizer_attrs) {
 
-void ModelTrainingInstance::register_and_allocate_layers() {
+  // allocate each layer's tensors
   for (layer_guid_t const & node: topological_ordering(this->computation_graph)) {
     this->training_backing.register_and_allocate_layer(node);
-  }
-}
-
-void ModelTrainingInstance::allocate_optimizer_tensors() {
-  for (layer_guid_t const & node: topological_ordering(this->computation_graph)) {
     this->training_backing.allocate_layer_optimizer_tensors(node, this->optimizer_attrs);
   }
 }
@@ -61,4 +56,4 @@ void ModelTrainingInstance::execute_update() {
   this->optimizer_attrs = get_next_iteration_optimizer_attrs(this->optimizer_attrs);
 }
 
-}
+} // namespace FlexFlow

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "attention.h"
+#include "local-execution/ops/attention.h"
 #include "kernels/attention_kernels.h"
 #include "local-execution/op_task_signature.h"
 #include "op-attrs/ops/attention.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "batch_matmul.h"
+#include "local-execution/ops/batch_matmul.h"
 #include "kernels/batch_matmul_kernels.h"
 #include "local-execution/op_task_signature.h"
 #include "op-attrs/get_output_shapes.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "batch_norm.h"
+#include "local-execution/ops/batch_norm.h"
 #include "kernels/batch_norm_kernels.h"
 
 namespace FlexFlow {

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "cast.h"
+#include "local-execution/ops/cast.h"
 #include "kernels/cast_kernels.h"
 
 #include "local-execution/op_task_signature.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "combine.h"
+#include "local-execution/ops/combine.h"
 #include "kernels/combine_kernels.h"
 #include "local-execution/op_task_invocation.h"
 #include "utils/hash-utils.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "concat.h"
+#include "local-execution/ops/concat.h"
 #include "kernels/concat_kernels.h"
 
 #include "local-execution/op_task_signature.h"

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

@@ -1,4 +1,4 @@
-#include "conv_2d.h"
+#include "local-execution/ops/conv_2d.h"
 #include "kernels/conv_2d_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 

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

@@ -1,4 +1,4 @@
-#include "dropout.h"
+#include "local-execution/ops/dropout.h"
 #include "kernels/dropout_kernels.h"
 #include "local-execution/op_task_invocation.h"
 #include "local-execution/op_task_signature.h"

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

@@ -1,4 +1,4 @@
-#include "element_binary.h"
+#include "local-execution/ops/element_binary.h"
 #include "kernels/element_binary_kernels.h"
 #include "local-execution/task_signature_impl.h"
 #include "op-attrs/get_output_shapes.h"

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

@@ -1,4 +1,4 @@
-#include "element_unary.h"
+#include "local-execution/ops/element_unary.h"
 #include "kernels/element_unary_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "op-attrs/parallel_tensor_shape.h"

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

@@ -1,4 +1,4 @@
-#include "flat.h"
+#include "local-execution/ops/flat.h"
 #include "kernels/flat_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "gather.h"
+#include "local-execution/ops/gather.h"
 #include "kernels/gather_kernels.h"
 #include "local-execution/legion_tensor_shape.h"
 #include "op-attrs/get_output_shapes.h"

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

@@ -1,4 +1,4 @@
-#include "input.h"
+#include "local-execution/ops/input.h"
 
 namespace FlexFlow {
 

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "layer_norm.h"
+#include "local-execution/ops/layer_norm.h"
 #include "kernels/layer_norm_kernels.h"
 #include "local-execution/legion_tensor_shape.h"
 #include "op-attrs/get_output_shapes.h"

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

@@ -1,4 +1,4 @@
-#include "linear.h"
+#include "local-execution/ops/linear.h"
 #include "kernels/linear_kernels.h"
 #include "local-execution/task_argument_accessor.h"
 #include "op-attrs/ff_dim.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "noop.h"
+#include "local-execution/ops/noop.h"
 
 namespace FlexFlow {
 

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

@@ -1,4 +1,4 @@
-#include "pool_2d.h"
+#include "local-execution/ops/pool_2d.h"
 #include "kernels/pool_2d_kernels.h"
 
 #include "op-attrs/get_output_shapes.h"

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

@@ -1,4 +1,4 @@
-#include "reduce.h"
+#include "local-execution/ops/reduce.h"
 #include "kernels/reduce_kernels.h"
 
 #include "op-attrs/get_output_shapes.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "reduction.h"
+#include "local-execution/ops/reduction.h"
 #include "kernels/reduction_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "repartition.h"
+#include "local-execution/ops/repartition.h"
 #include "kernels/partition_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "replicate.h"
+#include "local-execution/ops/replicate.h"
 #include "kernels/replicate_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "op-attrs/parallel_tensor_shape.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "reshape.h"
+#include "local-execution/ops/reshape.h"
 #include "kernels/reshape_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "reverse.h"
+#include "local-execution/ops/reverse.h"
 #include "kernels/accessor.h"
 #include "kernels/reverse_kernels.h"
 #include "op-attrs/get_output_shapes.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "softmax.h"
+#include "local-execution/ops/softmax.h"
 #include "kernels/softmax_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "op-attrs/parallel_tensor_shape.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "split.h"
+#include "local-execution/ops/split.h"
 #include "kernels/array_shape.h"
 #include "kernels/split_kernels.h"
 #include "op-attrs/get_output_shapes.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "topk.h"
+#include "local-execution/ops/topk.h"
 #include "kernels/topk_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"

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

@@ -13,7 +13,7 @@
  * limitations under the License.
  */
 
-#include "transpose.h"
+#include "local-execution/ops/transpose.h"
 #include "kernels/transpose_kernels.h"
 #include "op-attrs/get_output_shapes.h"
 #include "op-attrs/ops/transpose.h"

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

@@ -1,4 +1,4 @@
-#include "weight.h"
+#include "local-execution/ops/weight.h"
 
 namespace FlexFlow {
 

lib/local-execution/src/task_signature_impl.cc

@@ -1,33 +1,33 @@
 #include "local-execution/task_signature_impl.h"
-#include "ops/attention.h"
-#include "ops/batch_matmul.h"
-#include "ops/batch_norm.h"
-#include "ops/cast.h"
-#include "ops/combine.h"
-#include "ops/concat.h"
-#include "ops/conv_2d.h"
-#include "ops/dropout.h"
-#include "ops/element_binary.h"
-#include "ops/element_unary.h"
-#include "ops/embedding.h"
-#include "ops/flat.h"
-#include "ops/gather.h"
-#include "ops/input.h"
-#include "ops/layer_norm.h"
-#include "ops/linear.h"
-#include "ops/noop.h"
-#include "ops/pool_2d.h"
-#include "ops/reduce.h"
-#include "ops/reduction.h"
-#include "ops/repartition.h"
-#include "ops/replicate.h"
-#include "ops/reshape.h"
-#include "ops/reverse.h"
-#include "ops/softmax.h"
-#include "ops/split.h"
-#include "ops/topk.h"
-#include "ops/transpose.h"
-#include "ops/weight.h"
+#include "local-execution/ops/attention.h"
+#include "local-execution/ops/batch_matmul.h"
+#include "local-execution/ops/batch_norm.h"
+#include "local-execution/ops/cast.h"
+#include "local-execution/ops/combine.h"
+#include "local-execution/ops/concat.h"
+#include "local-execution/ops/conv_2d.h"
+#include "local-execution/ops/dropout.h"
+#include "local-execution/ops/element_binary.h"
+#include "local-execution/ops/element_unary.h"
+#include "local-execution/ops/embedding.h"
+#include "local-execution/ops/flat.h"
+#include "local-execution/ops/gather.h"
+#include "local-execution/ops/input.h"
+#include "local-execution/ops/layer_norm.h"
+#include "local-execution/ops/linear.h"
+#include "local-execution/ops/noop.h"
+#include "local-execution/ops/pool_2d.h"
+#include "local-execution/ops/reduce.h"
+#include "local-execution/ops/reduction.h"
+#include "local-execution/ops/repartition.h"
+#include "local-execution/ops/replicate.h"
+#include "local-execution/ops/reshape.h"
+#include "local-execution/ops/reverse.h"
+#include "local-execution/ops/softmax.h"
+#include "local-execution/ops/split.h"
+#include "local-execution/ops/topk.h"
+#include "local-execution/ops/transpose.h"
+#include "local-execution/ops/weight.h"
 #include "utils/overload.h"
 
 namespace FlexFlow {

lib/op-attrs/include/op-attrs/operator_attrs.h

@@ -3,35 +3,35 @@
 
 #include "op-attrs/ops/core.h"
 #include "op-attrs/pcg_operator_attrs.dtg.h"
-#include "ops/attention.h"
-#include "ops/batch_matmul.h"
-#include "ops/batch_norm.h"
-#include "ops/broadcast.h"
-#include "ops/cast.h"
-#include "ops/combine.h"
-#include "ops/concat.h"
-#include "ops/conv_2d.h"
-#include "ops/dropout.h"
-#include "ops/element_binary.h"
-#include "ops/element_unary.h"
-#include "ops/embedding.h"
-#include "ops/flat.h"
-#include "ops/gather.h"
-#include "ops/input.h"
-#include "ops/layer_norm.h"
-#include "ops/linear.h"
-#include "ops/noop.h"
-#include "ops/pool_2d.h"
-#include "ops/reduce.h"
-#include "ops/reduction.h"
-#include "ops/repartition.h"
-#include "ops/replicate.h"
-#include "ops/reshape.h"
-#include "ops/reverse.h"
-#include "ops/softmax.h"
-#include "ops/split.h"
-#include "ops/topk.h"
-#include "ops/transpose.h"
+#include "local-execution/ops/attention.h"
+#include "local-execution/ops/batch_matmul.h"
+#include "local-execution/ops/batch_norm.h"
+#include "local-execution/ops/broadcast.h"
+#include "local-execution/ops/cast.h"
+#include "local-execution/ops/combine.h"
+#include "local-execution/ops/concat.h"
+#include "local-execution/ops/conv_2d.h"
+#include "local-execution/ops/dropout.h"
+#include "local-execution/ops/element_binary.h"
+#include "local-execution/ops/element_unary.h"
+#include "local-execution/ops/embedding.h"
+#include "local-execution/ops/flat.h"
+#include "local-execution/ops/gather.h"
+#include "local-execution/ops/input.h"
+#include "local-execution/ops/layer_norm.h"
+#include "local-execution/ops/linear.h"
+#include "local-execution/ops/noop.h"
+#include "local-execution/ops/pool_2d.h"
+#include "local-execution/ops/reduce.h"
+#include "local-execution/ops/reduction.h"
+#include "local-execution/ops/repartition.h"
+#include "local-execution/ops/replicate.h"
+#include "local-execution/ops/reshape.h"
+#include "local-execution/ops/reverse.h"
+#include "local-execution/ops/softmax.h"
+#include "local-execution/ops/split.h"
+#include "local-execution/ops/topk.h"
+#include "local-execution/ops/transpose.h"
 #include "utils/record_formatter.h"
 #include "utils/variant.h"
 #include <variant>