merge

.editorconfig

@@ -0,0 +1,22 @@
+root = true
+
+# Unix-style newlines with a newline ending every file
+[*]
+end_of_line = lf
+insert_final_newline = true
+
+[{CMakeLists.txt,*.cmake}]
+indent_style = space
+indent_size = 2
+
+[*.{cc,h,cu,cpp}]
+indent_style = space
+indent_size = 2
+
+[*.py]
+indent_style = space
+indent_size = 4
+
+[*.toml]
+indent_style = space
+indent_size = 2

.flake/pkgs/hpp2plantuml.nix

@@ -0,0 +1,11 @@
+{buildPythonPackage, fetchPypi}:
+
+buildPythonPackage rec {
+  pname = "hpp2plantuml";
+  version = "0.8.5";
+  format = "wheel";
+  src = fetchPypi {
+    inherit pname version format;
+    sha256 = "sha256-PfTJmBypI21AAK3sMojygQfrhnRqcMmVCW4dxGfDfQg=";
+  };
+}

.gitattributes

@@ -0,0 +1,2 @@
+*.dtg.cc linguist-generated=true
+*.dtg.h linguist-generated=true

.github/workflows/helpers/cmake_cuda.sh

@@ -23,6 +23,7 @@ IFS=" " read -r -a FLAGS <<< "$CMAKE_FLAGS"
         -DCMAKE_C_COMPILER_LAUNCHER=ccache \
         -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
         -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache \
+        -DFF_USE_CODE_COVERAGE=ON \
         "${FLAGS[@]}"
 
 # vim: set tabstop=2 shiftwidth=2 expandtab:

.github/workflows/per-lib-check.yml

@@ -72,13 +72,13 @@ jobs:
         run: |
           build_libs.sh kernels
 
-      - name: Build substitutions
-        run: |
-          build_libs.sh substitutions
+      # - name: Build substitutions
+      #   run: |
+      #     build_libs.sh substitutions
 
-      - name: Build compiler
-        run: |
-          build_libs.sh compiler
+      # - name: Build compiler
+      #   run: |
+      #     build_libs.sh compiler
 
       - name: Build substitution-generator
         run: |
@@ -88,14 +88,40 @@ jobs:
         run: |
           test_libs.sh utils
 
-      - name: Test substitutions
+      - name: Test op-attrs
         run: |
-          test_libs.sh substitutions
+          test_libs.sh op-attrs
 
-      - name: Test compiler
+      - name: Test pcg
         run: |
-          test_libs.sh compiler
+          test_libs.sh pcg
+
+      # - name: Test substitutions
+      #   run: |
+      #     test_libs.sh substitutions
+
+      # - name: Test compiler
+      #   run: |
+      #     test_libs.sh compiler
 
       - name: Test substitution-generator
         run: |
           test_libs.sh substitution-generator
+
+      - name: Generate code coverage
+        run: |
+          echo "gitwork: $GITHUB_WORKSPACE"
+          lcov --capture --directory . --output-file main_coverage.info
+          lcov --extract main_coverage.info "$GITHUB_WORKSPACE/lib/*" --output-file main_coverage.info
+          lcov --remove main_coverage.info "$GITHUB_WORKSPACE/lib/*.dtg.h" "$GITHUB_WORKSPACE/lib/*.dtg.cc"  --output-file main_coverage.info
+          lcov --list main_coverage.info
+      
+      - name: Upload code coverage
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: main_coverage.info
+          flags: unittests
+          name: codecov-umbrella
+          fail_ci_if_error: false
+          verbose: true

.proj.toml

@@ -6,11 +6,16 @@ fix_compile_commands = false
 
 build_targets = [
   "kernels",
+<<<<<<< HEAD
 ]
 test_targets = [
   # "utils-tests",
   # "substitutions-tests",
   # "compiler-tests",
+  "pcg",
+  # "substitutions",
+  # "compiler",
+  "substitution-generator",
 ]
 
 [cmake_flags_extra]
@@ -20,4 +25,4 @@ CMAKE_CUDA_ARCHITECTURES = "60"
 CMAKE_HIP_ARCHITECTURES = "gfx900"
 CMAKE_CXX_COMPILER = "hipcc"
 CMAKE_C_COMPILER = "hipcc"
-# FF_CUDA_ARCH = "60"
+# FF_CUDA_ARCH = "60"

CMakeLists.txt

@@ -37,6 +37,7 @@ option(FF_USE_PREBUILT_LEGION "Enable use of Legion pre-compiled library, if ava
 option(FF_USE_ALL_PREBUILT_LIBRARIES "Enable use of all pre-compiled libraries, if available" OFF)
 option(FF_USE_PYTHON "Enable Python" ON)
 option(FF_BUILD_FROM_PYPI "Build from pypi" OFF)
+option(FF_USE_CODE_COVERAGE "Enable code coverage" OFF)
 
 set(FF_GASNET_CONDUITS aries udp mpi ibv ucx)
 set(FF_GASNET_CONDUIT "mpi" CACHE STRING "Select GASNet conduit ${FF_GASNET_CONDUITS}")
@@ -86,8 +87,10 @@ endif()
 # include(cuda)
 # include(cudnn)
 # include(nccl)
-include(CodeCoverage)
-append_coverage_compiler_flags()
+if (FF_USE_CODE_COVERAGE)
+  include(CodeCoverage)
+  append_coverage_compiler_flags()
+endif()
 # set_property(CACHE FF_GPU_BACKEND PROPERTY STRINGS ${FF_GPU_BACKENDS})
 
 include(json)

codecov.yml

@@ -0,0 +1,22 @@
+codecov:
+  branch: repo-refactor
+  notify:
+    require_ci_to_pass: false
+
+ignore:
+  - "**/*.dtg.h" 
+  - "**/*.dtg.cc"
+
+
+coverage:
+  status:
+    project:
+      default:
+        target: auto  # Automatically set target at 70% of the current project coverage
+        threshold: 0%  # Allows the coverage to drop by no more than 0% from the target
+        base: auto  # Picks the base of the pull request as a reference to compare against
+
+comment:
+  layout: "header, diff, flags, files"
+  behavior: default
+  require_changes: no

flake.nix

@@ -62,6 +62,7 @@
     {
       packages = {
         legion = pkgs.callPackage ./.flake/pkgs/legion.nix { };
+        hpp2plantuml = pkgs.python3Packages.callPackage ./.flake/pkgs/hpp2plantuml.nix { };
         rapidcheckFull = pkgs.symlinkJoin {
           name = "rapidcheckFull";
           paths = (with pkgs; [ rapidcheck.out rapidcheck.dev ]);
@@ -123,6 +124,7 @@
             ])
             (with self.packages.${system}; [
               legion
+              hpp2plantuml
               rapidcheckFull
               doctest
             ])
@@ -143,6 +145,10 @@
           inputsFrom = [ ci ];
           inherit (ci) CMAKE_FLAGS;
 
+          VIMPLUGINS = lib.strings.concatStringsSep "," [
+            "${proj-repo.packages.${system}.proj-nvim}"
+          ];
+
           buildInputs = builtins.concatLists [
             (with pkgs; [
               clang-tools
@@ -176,4 +182,4 @@
       };
     }
   );
-}
+}

lib/compiler/test/src/test_optimal_cost.cc

@@ -41,10 +41,9 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     MultiDiEdge e{n1, pcg.add_node_port(), n0, pcg.add_node_port()};
     pcg.add_edge(e);
-    pcg.add_output(e,
-                   ParallelTensor(ParallelTensorDims({2, 1}),
-                                  DataType::FLOAT,
-                                  CreateGrad::YES));
+    ParallelDim dim = {2, 1, false};
+    ParallelTensorDims dims = {FFOrdered<ParallelDim>{dim}};
+    pcg.add_output(e, ParallelTensor(dims, DataType::FLOAT, CreateGrad::YES));
 
     auto test_allowed_machine_views = [](Operator const &,
                                          MachineSpecification const &) {

lib/kernels/include/kernels/array_shape.h

@@ -6,6 +6,7 @@
 #include "utils/stack_vector.h"
 #include "utils/visitable.h"
 #include <cstddef>
+#include <optional>
 #include <vector>
 
 namespace FlexFlow {
@@ -41,8 +42,10 @@ struct ArrayShape {
   std::optional<std::size_t> at_maybe(std::size_t) const;
 
   ArrayShape reversed_dim_order() const;
-  ArrayShape sub_shape(std::optional<legion_dim_t> start,
-                       std::optional<legion_dim_t> end);
+
+  ArrayShape
+      sub_shape(std::optional<std::variant<ff_dim_t, legion_dim_t>> start,
+                std::optional<std::variant<ff_dim_t, legion_dim_t>> end) const;
 
 public:
   LegionTensorDims dims;

lib/kernels/include/kernels/cast_kernels.h

@@ -4,7 +4,7 @@
 #include "device.h"
 #include "kernels/accessor.h"
 #include "kernels/ff_handle.h"
-#include "op-attrs/activation.h"
+#include "op-attrs/activation.dtg.h"
 
 namespace FlexFlow {
 namespace Kernels {

lib/kernels/include/kernels/conv_2d_kernels.h

@@ -4,7 +4,7 @@
 #include "device.h"
 #include "kernels/accessor.h"
 #include "kernels/ff_handle.h"
-#include "op-attrs/activation.h"
+#include "op-attrs/activation.dtg.h"
 #include "utils/visitable.h"
 
 namespace FlexFlow {

lib/kernels/include/kernels/element_binary_kernels.h

@@ -5,7 +5,7 @@
 #include "ff_handle.h"
 #include "kernels/array_shape.h"
 #include "op-attrs/datatype.h"
-#include "op-attrs/op.h"
+#include "op-attrs/operator_type.h"
 
 namespace FlexFlow {
 

lib/kernels/include/kernels/element_unary_kernels.h

@@ -9,9 +9,6 @@
 
 namespace FlexFlow {
 
-using ElementUnaryUnifiedAttrs =
-    std::variant<ElementUnaryAttrs, ElementScalarUnaryAttrs>;
-
 struct ElementUnaryPerDeviceState {
   ffTensorDescriptor_t inputTensor, outputTensor;
   req<ffActivationDescriptor_t> actiDesc;
@@ -27,18 +24,34 @@ namespace ElementUnary {
 
 ElementUnaryPerDeviceState init_kernel(ArrayShape const &input_shape,
                                        ArrayShape const &output_shape,
-                                       ElementUnaryUnifiedAttrs const &attrs);
+                                       ElementUnaryAttrs const &attrs);
+
+void forward_kernel(ffStream_t stream,
+                    ElementUnaryPerDeviceState const &device_state,
+                    ElementUnaryAttrs const &attrs,
+                    PerDeviceFFHandle &handle,
+                    GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output);
 
 void forward_kernel(ffStream_t stream,
                     ElementUnaryPerDeviceState const &device_state,
-                    ElementUnaryUnifiedAttrs const &attrs,
+                    ElementScalarUnaryAttrs const &attrs,
                     PerDeviceFFHandle &handle,
                     GenericTensorAccessorR const &input,
                     GenericTensorAccessorW const &output);
 
 void backward_kernel(ffStream_t stream,
                      ElementUnaryPerDeviceState const &device_state,
-                     ElementUnaryUnifiedAttrs const &attrs,
+                     ElementUnaryAttrs const &attrs,
+                     PerDeviceFFHandle &handle,
+                     GenericTensorAccessorR const &input,
+                     GenericTensorAccessorW const &input_grad,
+                     GenericTensorAccessorR const &output,
+                     GenericTensorAccessorR const &output_grad);
+
+void backward_kernel(ffStream_t stream,
+                     ElementUnaryPerDeviceState const &device_state,
+                     ElementScalarUnaryAttrs const &attrs,
                      PerDeviceFFHandle &handle,
                      GenericTensorAccessorR const &input,
                      GenericTensorAccessorW const &input_grad,

lib/kernels/include/kernels/gather_kernels.h

@@ -2,36 +2,34 @@
 #define _FLEXFLOW_OPS_KERNELS_GATHER_KERNELS_H
 
 #include "accessor.h"
-#include "device.h"
+#include "kernels/device.h"
 
 namespace FlexFlow {
 
 struct GatherPerDeviceState {
-  int legion_dim;
-  req<DataType> index_data_type;
+  PerDeviceFFHandle handle;
+  legion_dim_t legion_dim;
 };
+
 FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(GatherPerDeviceState,
-                                             legion_dim,
-                                             index_data_type);
+                                             handle,
+                                             legion_dim);
 
 namespace Kernels {
 namespace Gather {
+
 void forward_kernel(ffStream_t stream,
                     GatherPerDeviceState const &m,
                     GenericTensorAccessorR const &input,
                     GenericTensorAccessorR const &index,
-                    GenericTensorAccessorW const &output,
-                    size_t stride,
-                    size_t input_dim_size,
-                    size_t output_dim_size);
+                    GenericTensorAccessorW const &output);
+
 void backward_kernel(ffStream_t stream,
                      GatherPerDeviceState const &m,
                      GenericTensorAccessorR const &output_grad,
                      GenericTensorAccessorR const &index,
-                     GenericTensorAccessorW const &input_grad,
-                     size_t stride,
-                     size_t input_dim_size,
-                     size_t output_dim_size);
+                     GenericTensorAccessorW const &input_grad);
+
 } // namespace Gather
 } // namespace Kernels
 } // namespace FlexFlow