Add initial benchmark setup (#1591)

* Add initial benchmark setup

* Add benchmark proj support

* Format

* Fix doctest test suite macros

* Address PR comments

* Fix proj

* Hopefully fix ci checks

* Fix build issue

* Update proj

.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;
+  };
+})

.github/workflows/tests.yml

@@ -39,17 +39,9 @@ jobs:
       - name: ccache
         uses: hendrikmuhs/ccache-action@v1.2
 
-      - name: Regenerate all dtgen files
+      - name: Run checks
         run: |
-          proj dtgen --force
-
-      - name: Run cmake
-        run: |
-          proj cmake --dtgen-skip
-
-      - name: Run build and tests
-        run: |
-          proj test --dtgen-skip -j$(nproc) --coverage --skip-gpu-tests
+          proj check cpu-ci
 
       - name: Upload code coverage
         uses: codecov/codecov-action@v4
@@ -104,14 +96,6 @@ jobs:
       - name: ccache
         uses: hendrikmuhs/ccache-action@v1.2
 
-      - name: regenerate all dtgen files
-        run: |
-          proj dtgen --force
-
-      - name: run cmake
-        run: |
-          proj cmake --dtgen-skip
-
-      - name: build and run gpu tests
+      - name: Run checks
         run: |
-          proj test --dtgen-skip -j$(nproc) --skip-build-cpu-tests
+          proj check gpu-ci

.gitignore

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

.proj.toml

@@ -3,31 +3,68 @@ testsuite_macro = "FF_TEST_SUITE"
 namespace_name = "FlexFlow"
 header_extension = ".h"
 
-build_targets = [
-  "utils",
-  "op-attrs",
-  "kernels",
-  "pcg",
-  "substitutions",
-  "compiler",
-  "substitution-generator",
-  "local-execution",
-  "models",
-  "export-model-arch",
-  "substitution-to-dot",
-]
-
-test_targets = [
-  "utils-tests",
-  "op-attrs-tests",
-  "kernels-tests",
-  "pcg-tests",
-  "substitutions-tests",
-  "compiler-tests",
-  "substitution-generator-tests",
-  "local-execution-tests",
-  "models-tests",
-]
+[targets.utils]
+type = "lib"
+tests = true
+benchmarks = true
+
+[targets.op-attrs]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.kernels]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.pcg]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.substitutions]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.compiler]
+type = "lib"
+tests = true
+benchmarks = true
+
+[targets.substitution-generator]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.local-execution]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.models]
+type = "lib"
+tests = true
+benchmarks = false
+
+[targets.export-model-arch]
+type = "bin"
+
+[targets.substitution-to-dot]
+type = "bin"
+
+# default_build_targets = [
+#   "utils",
+# ]
+
+# default_test_targets = [
+#   "utils",
+# ]
+
+# default_benchmark_targets = [
+#   "utils",
+# ]
 
 [cmake_flags_extra]
 FF_CUDA_ARCH = "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()
+