CPU Kernel Tests

.flake/pkgs/fccf/default.nix

@@ -0,0 +1,54 @@
+{ fetchFromGitHub
+, stdenv
+, cmake
+, pkg-config
+, libclang
+, libllvm
+, lib
+, zlib
+, argparse
+, nlohmann_json
+, fmt
+}:
+
+stdenv.mkDerivation rec {
+  pname = "fccf";
+  version = "03d373fc65e2d7ceeac441ba4bbddfdc25618dff";
+
+  src = fetchFromGitHub {
+    owner = "p-ranav";
+    repo = "fccf";
+    rev = version;
+    sha256 = "sha256-3NdPon5ZfjoGFFgBlb0rzRnfWgSopvAc5Gls2NWHaOE=";
+  };
+
+  nativeBuildInputs = [
+    cmake
+    pkg-config
+  ];
+
+  buildInputs = [
+    libclang
+    libllvm
+    zlib
+    argparse
+    nlohmann_json
+    fmt
+  ];
+
+  patches = [
+    ./json-package-name.patch
+    ./fix-argparse-include.patch
+  ];
+
+  cmakeFlags = [
+    "-DCMAKE_BUILD_TYPE=Release"
+    "-DFETCHCONTENT_TRY_FIND_PACKAGE_MODE=ALWAYS"
+  ];
+
+  meta = with lib; {
+    description = "A command-line tool that quickly searches through C/C++ source code in a directory based on a search string and prints relevant code snippets that match the query";
+    homepage = "https://github.com/p-ranav/fccf";
+    license = licenses.mit;
+  };
+}

.flake/pkgs/fccf/fix-argparse-include.patch

@@ -0,0 +1,13 @@
+diff --git a/source/main.cpp b/source/main.cpp
+index 7e131d3..6c05d89 100644
+--- a/source/main.cpp
++++ b/source/main.cpp
+@@ -6,7 +6,7 @@
+ #include <string_view>
+ #include <vector>
+
+-#include <argparse.hpp>
++#include <argparse/argparse.hpp>
+ #include <nlohmann/json.hpp>
+ #include "searcher.hpp"
+ #include <unistd.h>

.flake/pkgs/fccf/json-package-name.patch

@@ -0,0 +1,12 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 20bcbbf..923075f 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -48,6 +48,7 @@ FetchContent_MakeAvailable(fmt)
+
+ FetchContent_Declare(json
+   URL https://github.com/nlohmann/json/releases/download/v3.10.5/json.tar.xz
++  FIND_PACKAGE_ARGS NAMES nlohmann_json
+ )
+ FetchContent_MakeAvailable(json)
+

.github/runs-on.yml

@@ -1,23 +1,4 @@
 images:
-  runs-on-gpu-pinned:
-    platform: "linux"
-    arch: "x64"
-    owner: "135269210855" # runs-on
-    # to find, go to 
-    # https://us-east-2.console.aws.amazon.com/ec2/home?region=us-east-2#Images:visibility=public-images;search=:runs-on;v=3;$case=tags:false%5C,client:false;$regex=tags:false%5C,client:false
-    name: "runs-on-v2.2-ubuntu22-gpu-x64-20250220122045"
-
-  runs-on-cpu-pinned:
-    platform: "linux"
-    arch: "x64"
-    owner: "135269210855" # runs-on
-    name: "runs-on-v2.2-ubuntu22-full-x64-20250220122045"
-
-  official-ubuntu-ami:
-    platform: "linux"
-    arch: "x64"
-    ami: "ami-0a60b027285c0d4c5"
-
   flexflow-gpu-ci:
     platform: "linux"
     arch: "x64"

.github/workflows/tests.yml

@@ -57,9 +57,9 @@ jobs:
     name: GPU unit tests
     needs: cpu-ci
     runs-on:
-      - runs-on
+      - runs-on=${{ github.run_id }}
       - family=g4dn.xlarge
-      - image=runs-on-gpu-pinned
+      - image=flexflow-gpu-ci
 
     strategy:
       max-parallel: 1

.proj.toml

@@ -2,57 +2,81 @@ project_name = "flexflow"
 testsuite_macro = "FF_TEST_SUITE"
 namespace_name = "FlexFlow"
 header_extension = ".h"
+cuda_launch_cmd = [
+  "nixGL",
+  "--",
+]
 
 [targets.utils]
 type = "lib"
-tests = true
-benchmarks = true
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = true
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.op-attrs]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.kernels]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = true
+has-cuda-benchmarks = false
 
 [targets.pcg]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.substitutions]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.compiler]
 type = "lib"
-tests = true
-benchmarks = true
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = true
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.substitution-generator]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.local-execution]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.models]
 type = "lib"
-tests = true
-benchmarks = false
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
 
 [targets.export-model-arch]
 type = "bin"
+cuda = false
 
 [targets.substitution-to-dot]
 type = "bin"
+cuda = false
 
 # default_build_targets = [
 #   "utils",

cmake/flexflow-utils.cmake

@@ -126,11 +126,16 @@ function(ff_add_test_executable)
     ${FF_TEST_EXEC_NAME}
     ${SRC})
 
+  target_include_directories(
+    ${FF_TEST_EXEC_NAME}
+    PRIVATE
+    ${FF_TEST_EXEC_PRIVATE_INCLUDE})
+
   target_link_libraries(
     ${FF_TEST_EXEC_NAME}
     ${FF_TEST_EXEC_DEPS})
 
-  target_compile_definitions(${FF_TEST_EXEC_NAME} PRIVATE FF_TEST_SUITE="${FF_TEST_EXEC_NAME}" FF_CUDA_TEST_SUITE="cuda-${FF_TEST_EXEC_NAME}")
+  target_compile_definitions(${FF_TEST_EXEC_NAME} PRIVATE FF_TEST_SUITE="cpu-${FF_TEST_EXEC_NAME}" FF_CUDA_TEST_SUITE="cuda-${FF_TEST_EXEC_NAME}")
 
   define_ff_vars(${FF_TEST_EXEC_NAME})
   ff_set_cxx_properties(${FF_TEST_EXEC_NAME})

flake.nix

@@ -59,6 +59,7 @@
         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 { };
+        fccf = pkgs.callPackage ./.flake/pkgs/fccf { };
         rapidcheckFull = pkgs.symlinkJoin {
           name = "rapidcheckFull";
           paths = (with pkgs; [ rapidcheck.out rapidcheck.dev ]);
@@ -162,6 +163,7 @@
               ruff
               jq
               gh
+              expect
             ])
             (with pkgs.python3Packages; [
               gitpython
@@ -179,6 +181,7 @@
             (with self.packages.${system}; [
               ffdb
               hpp2plantuml
+              fccf
             ])
           ];
         };

lib/kernels/CMakeLists.txt

@@ -7,8 +7,7 @@ file(GLOB_RECURSE SRC
      CONFIGURE_DEPENDS
      LIST_DIRECTORIES False
      src/*.cc
-     src/cuda/cuda_helper.cu
-     src/cuda/ops/*.cu
+     src/cuda/*.cu
      )
 
 add_library(
@@ -30,6 +29,7 @@ target_link_libraries(
   cudnn
   nccl
   utils
+  pcg
 )
 
 define_ff_vars(${project_target})