Add initial benchmark setup

.flake/pkgs/bencher-cli.nix

@@ -0,0 +1,39 @@
+{ stdenv
+, lib
+, fetchurl
+, alsaLib
+, openssl
+, zlib
+, pulseaudio
+, autoPatchelfHook
+}:
+
+stdenv.mkDerivation rec {
+  pname = "bencher-cli";
+  version = "0.4.33";
+
+  system = "x86_64-linux";
+
+  src = fetchurl {
+    url = "https://github.com/bencherdev/bencher/releases/download/v${version}/bencher-v${version}-linux-x86-64";
+    hash = "sha256-3q2ZGSqbcMaUcNMGJN+IsEP/+RlYHnsmiWdJ2oV2qmw=";
+  };
+
+  nativeBuildInputs = [
+    autoPatchelfHook
+  ];
+
+  dontUnpack = true;
+
+  installPhase = ''
+    ls
+    runHook preInstall
+    install -m755 -D $src $out/bin/bencher
+    runHook postInstall
+  '';
+
+  meta = with lib; {
+    homepage = "https://bencher.dev/";
+    platforms = platforms.linux;
+  };
+}

.flake/pkgs/cpptrace.nix

@@ -0,0 +1,48 @@
+{ fetchFromGitHub
+, stdenv
+, lib
+, cmake
+, pkg-config
+, zstd
+, libdwarf-lite
+, libunwind
+}:
+
+stdenv.mkDerivation rec {
+  pname = "cpptrace";
+  version = "0.7.5";
+
+  src = fetchFromGitHub {
+    owner = "jeremy-rifkin";
+    repo = "cpptrace";
+    rev = "v${version}";
+    sha256 = "sha256-2rDyH9vo47tbqqZrTupAOrMySj4IGKeWX8HBTGjFf+g=";
+  };
+
+  nativeBuildInputs = [
+    cmake
+    pkg-config
+  ];
+
+  buildInputs = [
+    zstd
+    libdwarf-lite
+    libunwind
+  ];
+
+  cmakeFlags = [
+    "-DBUILD_SHARED_LIBS=On" 
+    "-DCPPTRACE_USE_EXTERNAL_ZSTD=1"
+    "-DCPPTRACE_USE_EXTERNAL_LIBDWARF=1"
+    "-DCPPTRACE_STD_FORMAT=0"
+    "-DCPPTRACE_STATIC_DEFINE=0"
+    "-DCPPTRACE_UNWIND_WITH_LIBUNWIND=1"
+    "-DCPPTRACE_FIND_LIBDWARF_WITH_PKGCONFIG=1"
+  ];
+
+  meta = with lib; {
+    description = "Simple, portable, and self-contained stacktrace library for C++11 and newer";
+    homepage = "https://github.com/jeremy-rifkin/cpptrace";
+    license = licenses.mit;
+  };
+}

.flake/pkgs/libassert.nix

@@ -0,0 +1,39 @@
+{ fetchFromGitHub
+, stdenv
+, lib
+, cmake
+, cpptrace
+, zstd
+}:
+
+stdenv.mkDerivation rec {
+  pname = "libassert";
+  version = "2.1.4";
+
+  src = fetchFromGitHub {
+    owner = "jeremy-rifkin";
+    repo = "libassert";
+    rev = "v${version}";
+    sha256 = "sha256-Zkh6JjJqtOf91U01fpP4hKhhXfH7YGInodG8CZxHHXQ=";
+  };
+
+  nativeBuildInputs = [
+    cmake
+  ];
+
+  propagatedBuildInputs = [
+    cpptrace
+    zstd
+  ];
+
+  cmakeFlags = [
+    "-DBUILD_SHARED_LIBS=1"
+    "-DLIBASSERT_USE_EXTERNAL_CPPTRACE=1"
+  ];
+
+  meta = with lib; {
+    description = "The most over-engineered C++ assertion library";
+    homepage = "https://github.com/jeremy-rifkin/libassert";
+    license = licenses.mit;
+  };
+}

.flake/pkgs/libdwarf-lite.nix

@@ -0,0 +1,40 @@
+{
+  lib,
+  stdenv,
+  fetchFromGitHub,
+  cmake,
+}:
+
+stdenv.mkDerivation (finalAttrs: {
+  pname = "libdwarf-lite";
+  version = "0.11.1";
+
+  src = fetchFromGitHub {
+    owner = "jeremy-rifkin";
+    repo = "libdwarf-lite";
+    rev = "v${finalAttrs.version}";
+    hash = "sha256-qHikjAG5xuuHquqqKGuiDHXVZSlg/MbNp9JNSAKM/Hs=";
+  };
+
+  outputs = [
+    "dev"
+    "lib"
+    "out"
+  ];
+
+  nativeBuildInputs = [ cmake ];
+
+  cmakeFlags = [
+    (lib.cmakeBool "BUILD_DWARFDUMP" false)
+    (lib.cmakeBool "PIC_ALWAYS" true)
+  ];
+
+  meta = {
+    description = "Minimal libdwarf mirror for faster cloning and configuration";
+    homepage = "https://github.com/jeremy-rifkin/libdwarf-lite";
+    license = lib.licenses.lgpl21Only;
+    maintainers = [ ];
+    mainProgram = "libdwarf-lite";
+    platforms = lib.platforms.all;
+  };
+})

.gitignore

@@ -5,6 +5,9 @@
 *.dtg.cc
 *.dtg.h
 
+# nix build result
+/result
+
 # ctest-created files
 Testing/
 

.proj.toml

@@ -29,6 +29,11 @@ test_targets = [
   "models-tests",
 ]
 
+benchmark_targets = [
+  "utils-benchmarks",
+  "compiler-benchmarks",
+]
+
 [cmake_flags_extra]
 FF_CUDA_ARCH = "60"
 CMAKE_CUDA_ARCHITECTURES = "60"

CMakeLists.txt

@@ -56,6 +56,8 @@ option(FF_USE_EXTERNAL_SPDLOG "Use pre-installed spdlog" OFF)
 option(FF_USE_EXTERNAL_DOCTEST "Use pre-installed doctest" OFF)
 option(FF_USE_EXTERNAL_RAPIDCHECK "Use pre-installed rapidcheck" OFF)
 option(FF_USE_EXTERNAL_EXPECTED "Use pre-installed tl::expected" OFF)
+option(FF_USE_EXTERNAL_GBENCHMARK "Use pre-installed google benchmark" OFF)
+option(FF_USE_EXTERNAL_LIBASSERT "Use pre-installed libassert" OFF)
 
 option(FF_BUILD_RESNET "build resnet example" OFF)
 option(FF_BUILD_RESNEXT "build resnext example" OFF)
@@ -98,6 +100,8 @@ include(json)
 include(expected)
 include(spdlog)
 include(doctestlib) # named doctestlib to avoid a name collision with doctest.cmake in rapidcheck
+include(gbenchmark)
+include(libassert)
 include(visit_struct)
 include(CTest)
 include(fmt)

cmake/flexflow-utils.cmake

@@ -137,6 +137,45 @@ function(ff_add_test_executable)
   doctest_discover_tests(${FF_TEST_EXEC_NAME} ADD_LABELS 1)
 endfunction()
 
+function(ff_add_benchmark_executable)
+  ff_parse_args(
+    PREFIX 
+      FF_BENCHMARK_EXEC
+    ARGS
+      NAME
+    VARIADIC_ARGS
+      SRC_PATTERNS
+      PRIVATE_INCLUDE
+      DEPS
+    PARSE
+      ${ARGN}
+  )
+
+  if(NOT FF_BENCHMARK_EXEC_NAME MATCHES "^(.*)-benchmarks\$")
+    message(FATAL_ERROR "Benchmark name ${FF_BENCHMARK_EXEC_NAME} must end with -benchmarks")
+  endif()
+  set(FF_LIBRARY_NAME ${CMAKE_MATCH_1})
+
+  project(${FF_BENCHMARK_EXEC_NAME})
+  file(GLOB_RECURSE SRC
+       CONFIGURE_DEPENDS
+       LIST_DIRECTORIES False
+       ${FF_BENCHMARK_EXEC_SRC_PATTERNS})
+
+  add_executable(
+    ${FF_BENCHMARK_EXEC_NAME}
+    ${SRC})
+
+  target_link_libraries(
+    ${FF_BENCHMARK_EXEC_NAME}
+    ${FF_BENCHMARK_EXEC_DEPS}
+    gbenchmark
+    gbenchmark-main)
+
+  define_ff_vars(${FF_BENCHMARK_EXEC_NAME})
+  ff_set_cxx_properties(${FF_BENCHMARK_EXEC_NAME})
+endfunction()
+
 function(ff_add_executable)
   ff_parse_args(
     PREFIX 

cmake/gbenchmark.cmake

@@ -0,0 +1,10 @@
+include(aliasing)
+
+if (FF_USE_EXTERNAL_GBENCHMARK)
+  find_package(benchmark REQUIRED)
+  alias_library(gbenchmark benchmark::benchmark)
+  alias_library(gbenchmark-main benchmark::benchmark_main)
+else()
+  message(FATAL_ERROR "Currently FF_USE_EXTERNAL_GBENCHMARK is required")
+endif()
+

cmake/libassert.cmake

@@ -0,0 +1,9 @@
+include(aliasing)
+
+if(FF_USE_EXTERNAL_LIBASSERT)
+  find_package(libassert REQUIRED)
+else()
+  message(FATAL_ERROR "Currently FF_USE_EXTERNAL_LIBASSERT is required")
+endif()
+
+alias_library(libassert libassert::assert)

flake.nix

@@ -51,8 +51,12 @@
       proj = proj-repo.packages.${system}.proj;
     in 
     {
-      packages = {
+      packages = rec {
+        libdwarf-lite = pkgs.callPackage ./.flake/pkgs/libdwarf-lite.nix { };
+        cpptrace = pkgs.callPackage ./.flake/pkgs/cpptrace.nix { inherit libdwarf-lite; };
+        libassert = pkgs.callPackage ./.flake/pkgs/libassert.nix { inherit cpptrace; };
         legion = pkgs.callPackage ./.flake/pkgs/legion.nix { };
+        bencher-cli = pkgs.callPackage ./.flake/pkgs/bencher-cli.nix { };
         ffdb = pkgs.callPackage ./.flake/pkgs/ffdb { inherit proj; };
         hpp2plantuml = pkgs.python3Packages.callPackage ./.flake/pkgs/hpp2plantuml.nix { };
         rapidcheckFull = pkgs.symlinkJoin {
@@ -86,6 +90,8 @@
                                 -DFF_USE_EXTERNAL_DOCTEST=ON \
                                 -DFF_USE_EXTERNAL_RAPIDCHECK=ON \
                                 -DFF_USE_EXTERNAL_EXPECTED=ON \
+                                -DFF_USE_EXTERNAL_GBENCHMARK=ON \
+                                -DFF_USE_EXTERNAL_LIBASSERT=ON \
                                 -DFF_USE_EXTERNAL_RANGEV3=ON \
                                 -DFF_USE_EXTERNAL_BOOST_PREPROCESSOR=ON \
                                 -DFF_USE_EXTERNAL_TYPE_INDEX=ON"
@@ -113,11 +119,13 @@
               doxygen
               lcov # for code coverage
               compdb
+              gbenchmark
             ])
             (with proj-repo.packages.${system}; [
               proj
             ])
             (with self.packages.${system}; [
+              libassert
               legion
               rapidcheckFull
               doctest

lib/compiler/CMakeLists.txt

@@ -18,3 +18,4 @@ ff_add_library(
 
 add_subdirectory(ffi)
 add_subdirectory(test)
+add_subdirectory(benchmark)

lib/compiler/benchmark/CMakeLists.txt

@@ -0,0 +1,12 @@
+ff_add_benchmark_executable(
+  NAME
+    compiler-benchmarks
+  SRC_PATTERNS
+    src/*.cc
+  PRIVATE_INCLUDE
+    src/
+  DEPS
+    utils
+    compiler
+    models
+)

lib/compiler/benchmark/src/compiler/series_parallel/computation_graph/get_computation_graph_series_parallel_decomposition.cc

@@ -0,0 +1,39 @@
+#include "compiler/series_parallel/computation_graph/get_computation_graph_series_parallel_decomposition.h"
+#include "models/bert/bert.h"
+#include "models/candle_uno/candle_uno.h"
+#include "models/inception_v3/inception_v3.h"
+#include "models/split_test/split_test.h"
+#include "models/transformer/transformer.h"
+#include <benchmark/benchmark.h>
+
+using namespace ::FlexFlow;
+
+static void benchmark_get_computation_graph_series_parallel_decomposition(
+    benchmark::State &state, ComputationGraph const &cg) {
+  for (auto _ : state) {
+    get_computation_graph_series_parallel_decomposition(cg);
+  }
+}
+
+BENCHMARK_CAPTURE(benchmark_get_computation_graph_series_parallel_decomposition,
+                  split_test,
+                  get_split_test_computation_graph(/*batch_size=*/8_n));
+
+BENCHMARK_CAPTURE(
+    benchmark_get_computation_graph_series_parallel_decomposition,
+    transformer,
+    get_transformer_computation_graph(get_default_transformer_config()));
+
+BENCHMARK_CAPTURE(benchmark_get_computation_graph_series_parallel_decomposition,
+                  bert,
+                  get_bert_computation_graph(get_default_bert_config()));
+
+BENCHMARK_CAPTURE(
+    benchmark_get_computation_graph_series_parallel_decomposition,
+    candle_uno,
+    get_candle_uno_computation_graph(get_default_candle_uno_config()));
+
+BENCHMARK_CAPTURE(benchmark_get_computation_graph_series_parallel_decomposition,
+                  inception_v3,
+                  get_inception_v3_computation_graph(
+                      get_default_inception_v3_training_config()));

lib/utils/CMakeLists.txt

@@ -13,7 +13,10 @@ ff_add_library(
     fmt
     json
     cuda
+    rapidcheck
+    libassert
 )      
 
 add_subdirectory(ffi)
 add_subdirectory(test)
+add_subdirectory(benchmark)