Local execution tests

.github/workflows/helpers/test_libs.sh

@@ -7,7 +7,7 @@ DIR="$(realpath -- "$(dirname "${BASH_SOURCE[0]}")")"
 REPO="$(realpath -- "$DIR/../../../")"
 
 TEST_LIBS=("${@/%/-tests}")
-REGEX="^$(IFS='|'; echo "${TEST_LIBS[*]}")\$"
+REGEX="^($(IFS='|'; echo "${TEST_LIBS[*]}"))\$"
 
 cd "$REPO/build-ci"
 make -j $(( $(nproc) < 2 ? 1 : $(nproc)-1 )) "${TEST_LIBS[@]}"

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

@@ -116,6 +116,10 @@ jobs:
         run: |
           test_libs.sh substitution-generator
 
+      - name: Test local-execution
+        run: |
+          test_libs.sh local-execution
+
       - name: Generate code coverage
         run: |
           echo "gitwork: $GITHUB_WORKSPACE"

.proj.toml

@@ -11,16 +11,18 @@ build_targets = [
   "substitutions",
   "compiler",
   "substitution-generator",
-  "local-execution", 
+  "local-execution",
 ]
 
 test_targets = [
+  # "kernels-tests",
   "utils-tests",
   "op-attrs-tests",
   "pcg-tests",
   "substitutions-tests",
   "compiler-tests",
   "substitution-generator-tests",
+  "local-execution-tests"
 ]
 
 [cmake_flags_extra]

cmake/flexflow-utils.cmake

@@ -118,7 +118,7 @@ function(ff_add_test_executable)
     ${FF_TEST_EXEC_NAME}
     ${FF_TEST_EXEC_DEPS})
 
-  target_compile_definitions(${FF_TEST_EXEC_NAME} PRIVATE FF_TEST_SUITE="${FF_TEST_EXEC_NAME}")
+  target_compile_definitions(${FF_TEST_EXEC_NAME} PRIVATE FF_TEST_SUITE="${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})

lib/kernels/CMakeLists.txt

@@ -40,4 +40,4 @@ set_target_properties(
   CUDA_STANDARD 17
 )
 
-add_subdirectory(test)
+add_subdirectory(test)

lib/kernels/include/kernels/accessor.h

@@ -145,13 +145,30 @@ std::vector<real_type<DT> const *>
 GenericTensorAccessorR read_only_accessor_from_write_accessor(
     GenericTensorAccessorW const &write_accessor);
 
+bool is_shape_and_dtype_equal(GenericTensorAccessorW const &acc1,
+                              GenericTensorAccessorW const &acc2);
+
+bool shape_and_dtype_matches(GenericTensorAccessorW const &accessor,
+                             ArrayShape const &expected_shape,
+                             DataType const &expected_dtype);
+
+bool shape_and_dtype_matches(GenericTensorAccessorR const &accessor,
+                             ArrayShape const &expected_shape,
+                             DataType const &expected_dtype);
+
+std::pair<ArrayShape, DataType>
+    get_shape_and_datatype(GenericTensorAccessorR const &accessor);
+std::pair<ArrayShape, DataType>
+    get_shape_and_datatype(GenericTensorAccessorW const &accessor);
+
 } // namespace FlexFlow
 
 namespace FlexFlow {
 static_assert(is_well_behaved_value_type_no_hash<GenericTensorAccessorR>::value,
               "");
 static_assert(is_well_behaved_value_type_no_hash<GenericTensorAccessorW>::value,
               "");
+
 } // namespace FlexFlow
 
 #endif

lib/kernels/include/kernels/array_shape.h

@@ -39,7 +39,8 @@ struct ArrayShape {
   legion_dim_t last_idx() const;
   legion_dim_t neg_idx(int) const;
 
-  std::optional<std::size_t> at_maybe(std::size_t) const;
+  std::optional<std::size_t> at_maybe(legion_dim_t) const;
+  std::optional<std::size_t> at_maybe(ff_dim_t) const;
 
   ArrayShape
       sub_shape(std::optional<std::variant<ff_dim_t, legion_dim_t>> start,
@@ -54,6 +55,9 @@ size_t get_volume(ArrayShape const &);
 
 TensorShape get_tensor_shape(ArrayShape const &, DataType);
 
+std::string format_as(ArrayShape const &);
+std::ostream &operator<<(std::ostream &, ArrayShape const &);
+
 } // namespace FlexFlow
 
 #endif

lib/kernels/include/kernels/attention_kernels.h

@@ -25,6 +25,25 @@ struct MHAPerDeviceState {
   int *hiWinIdx;
   void *reserveSpace;
   Allocator allocator;
+
+  bool operator==(MHAPerDeviceState const &other) const;
+  bool operator!=(MHAPerDeviceState const &other) const;
+
+private:
+  std::tuple<decltype(handle) const &,
+             decltype(weightSize) const &,
+             decltype(reserveSpaceSize) const &,
+             decltype(attnDesc) const &,
+             decltype(qDesc) const &,
+             decltype(kDesc) const &,
+             decltype(vDesc) const &,
+             decltype(oDesc) const &,
+             decltype(devQoSeqArray) const &,
+             decltype(devKvSeqArray) const &,
+             decltype(loWinIdx) const &,
+             decltype(hiWinIdx) const &,
+             decltype(reserveSpace) const &>
+      tie() const;
 };
 
 FF_VISITABLE_STRUCT_NO_EQ(MHAPerDeviceState,
@@ -43,6 +62,9 @@ FF_VISITABLE_STRUCT_NO_EQ(MHAPerDeviceState,
                           reserveSpace,
                           allocator);
 
+std::string format_as(MHAPerDeviceState const &x);
+std::ostream &operator<<(std::ostream &s, MHAPerDeviceState const &x);
+
 namespace Kernels {
 namespace MultiHeadAttention {
 

lib/kernels/include/kernels/device.h

@@ -95,11 +95,13 @@ using coord_t = long long;
     exit(1);                                                                   \
   } while (0)
 
+char const *getCudaErrorString(cudaError_t status);
+
 #define checkCUDA(status)                                                      \
   do {                                                                         \
     std::stringstream _error;                                                  \
     if (status != 0) {                                                         \
-      _error << "CUDA failure: " << cudaGetErrorString(status) << " ("         \
+      _error << "CUDA failure: " << getCudaErrorString(status) << " ("         \
              << status << ")";                                                 \
       FatalError(_error.str());                                                \
     }                                                                          \

lib/kernels/include/kernels/ff_handle.h

@@ -40,6 +40,9 @@ FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(PerDeviceFFHandle,
                                              allowTensorOpMathConversion);
 #endif
 
+std::string format_as(PerDeviceFFHandle const &x);
+std::ostream &operator<<(std::ostream &s, PerDeviceFFHandle const &x);
+
 } // namespace FlexFlow
 
 #endif

lib/kernels/include/kernels/legion_dim.h

@@ -15,6 +15,23 @@ using LegionOrdered = DimOrdered<legion_dim_t, T>;
 
 using LegionTensorDims = LegionOrdered<size_t>;
 
+template <typename T>
+FFOrdered<T>
+    ff_ordered_from_legion_ordered(LegionOrdered<T> const &legion_ordered) {
+  return FFOrdered<T>(legion_ordered.rbegin(), legion_ordered.rend());
+}
+
+template <typename T>
+std::string format_as(LegionOrdered<T> const &v) {
+  std::vector<T> as_vec(v.cbegin(), v.cend());
+  return fmt::format("<legion_ordered {}>", as_vec);
+}
+
+template <typename T>
+std::ostream &operator<<(std::ostream &s, LegionOrdered<T> const &v) {
+  return (s << fmt::to_string(v));
+}
+
 } // namespace FlexFlow
 
 #endif

lib/kernels/include/kernels/profiling.h

@@ -2,20 +2,11 @@
 #define _FLEXFLOW_KERNELS_PROFILING_H
 
 #include "device.h"
+#include "kernels/profiling_settings.dtg.h"
 #include "utils/visitable.h"
 
 namespace FlexFlow {
 
-struct ProfilingSettings : public use_visitable_cmp<ProfilingSettings> {
-public:
-  ProfilingSettings() = delete;
-  ProfilingSettings(int warmup_iters, int measure_iters);
-
-public:
-  int warmup_iters;
-  int measure_iters;
-};
-
 template <typename F, typename... Ts>
 std::optional<float>
     profiling_wrapper(F const &f, bool enable_profiling, Ts &&...ts) {
@@ -59,7 +50,4 @@ std::optional<float> profiling_wrapper(F const &f,
 
 } // namespace FlexFlow
 
-VISITABLE_STRUCT(::FlexFlow::ProfilingSettings, warmup_iters, measure_iters);
-MAKE_VISIT_HASHABLE(::FlexFlow::ProfilingSettings);
-
 #endif

lib/kernels/include/kernels/profiling_settings.struct.toml

@@ -0,0 +1,18 @@
+namespace = "FlexFlow"
+name = "ProfilingSettings"
+
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+]
+
+[[fields]]
+name = "warmup_iters"
+type = "int"
+
+[[fields]]
+name = "measure_iters"
+type = "int"

lib/kernels/src/accessor.cc

@@ -138,4 +138,33 @@ GenericTensorAccessorR read_only_accessor_from_write_accessor(
       writable.data_type, writable.shape, req<void const *>(writable.ptr)};
 }
 
+bool is_shape_and_dtype_equal(GenericTensorAccessorW const &acc1,
+                              GenericTensorAccessorW const &acc2) {
+  return acc1.shape == acc2.shape && acc1.data_type == acc2.data_type;
+}
+
+bool shape_and_dtype_matches(GenericTensorAccessorW const &accessor,
+                             ArrayShape const &expected_shape,
+                             DataType const &expected_dtype) {
+  return accessor.shape == expected_shape &&
+         accessor.data_type == expected_dtype;
+}
+
+bool shape_and_dtype_matches(GenericTensorAccessorR const &accessor,
+                             ArrayShape const &expected_shape,
+                             DataType const &expected_dtype) {
+  return accessor.shape == expected_shape &&
+         accessor.data_type == expected_dtype;
+}
+
+std::pair<ArrayShape, DataType>
+    get_shape_and_datatype(GenericTensorAccessorR const &accessor) {
+  return std::make_pair(accessor.shape, accessor.data_type);
+}
+
+std::pair<ArrayShape, DataType>
+    get_shape_and_datatype(GenericTensorAccessorW const &accessor) {
+  return std::make_pair(accessor.shape, accessor.data_type);
+}
+
 } // namespace FlexFlow

lib/kernels/src/array_shape.cc

@@ -39,7 +39,15 @@ std::size_t ArrayShape::num_elements() const {
 }
 
 std::size_t ArrayShape::operator[](legion_dim_t idx) const {
-  return dims[idx];
+  return dims.at(idx);
+}
+
+std::size_t ArrayShape::at(legion_dim_t idx) const {
+  return dims.at(idx);
+}
+
+std::size_t ArrayShape::at(ff_dim_t idx) const {
+  return dims.at(legion_dim_from_ff_dim(idx, this->num_dims()));
 }
 
 ArrayShape ArrayShape::sub_shape(
@@ -48,16 +56,37 @@ ArrayShape ArrayShape::sub_shape(
   NOT_IMPLEMENTED();
 }
 
-std::optional<std::size_t> ArrayShape::at_maybe(std::size_t index) const {
-  if (index < dims.size()) {
-    return dims.at(legion_dim_t(index));
+std::optional<std::size_t> ArrayShape::at_maybe(legion_dim_t index) const {
+  if (index.value < dims.size()) {
+    return dims.at(index);
   } else {
     return std::nullopt;
   }
 }
 
+std::optional<std::size_t> ArrayShape::at_maybe(ff_dim_t index) const {
+  return this->at_maybe(legion_dim_from_ff_dim(index, this->num_dims()));
+}
+
 size_t get_volume(ArrayShape const &shape) {
   return shape.get_volume();
 }
 
+TensorShape get_tensor_shape(ArrayShape const &shape, DataType dtype) {
+  return TensorShape{TensorDims{ff_ordered_from_legion_ordered(shape.dims)},
+                     dtype};
+}
+
+std::string format_as(ArrayShape const &x) {
+  std::ostringstream oss;
+  oss << "<ArrayShape";
+  oss << " dims=" << x.dims;
+  oss << ">";
+  return oss.str();
+}
+
+std::ostream &operator<<(std::ostream &s, ArrayShape const &x) {
+  return (s << fmt::to_string(x));
+}
+
 } // namespace FlexFlow

lib/kernels/src/cuda/cuda_helper.cu

@@ -220,25 +220,14 @@ __host__ void
 ffStatus_t
     cudnnSetTensorDescriptorFromArrayShape(cudnnTensorDescriptor_t tensor,
                                            ArrayShape const &shape) {
-  std::vector<std::size_t> reversed_dims(shape.dims.begin(), shape.dims.end());
-  reversed(reversed_dims);
-  ArrayShape flipped(reversed_dims);
-
-  if (flipped.get_dim() == 5) {
-    assert(flipped[legion_dim_t(0)] == 1);
-    flipped = flipped.sub_shape(legion_dim_t(1), std::nullopt);
-  }
-
-  assert(flipped.get_dim() > 0);
-  assert(flipped.get_dim() < 4);
-
-  return cudnnSetTensor4dDescriptor(tensor,
-                                    CUDNN_TENSOR_NCHW,
-                                    CUDNN_DATA_FLOAT,
-                                    flipped.at_maybe(0).value_or(1),
-                                    flipped.at_maybe(1).value_or(2),
-                                    flipped.at_maybe(2).value_or(3),
-                                    flipped.at_maybe(3).value_or(3));
+  return cudnnSetTensor4dDescriptor(
+      tensor,
+      CUDNN_TENSOR_NCHW,
+      CUDNN_DATA_FLOAT,
+      shape.at_maybe(legion_dim_t{0}).value_or(1),
+      shape.at_maybe(legion_dim_t{1}).value_or(1),
+      shape.at_maybe(legion_dim_t{2}).value_or(1),
+      shape.at_maybe(legion_dim_t{3}).value_or(1));
 }
 
 cudnnDataType_t ff_to_cudnn_datatype(DataType type) {