diff --git a/cmake/nccl.cmake b/cmake/nccl.cmake index 5e1ab05591..83d5d0d1f7 100644 --- a/cmake/nccl.cmake +++ b/cmake/nccl.cmake @@ -89,7 +89,7 @@ else() BUILD_BYPRODUCTS ${CMAKE_BINARY_DIR}/deps/nccl/lib/libnccl${LIBEXT} INSTALL_COMMAND "" CONFIGURE_COMMAND "" - BUILD_COMMAND make src.build "${NCCL_BUILD_NVCC_GENCODE}" "CUDA_HOME=${CUDA_TOOLKIT_ROOT_DIR}" "BUILDDIR=${CMAKE_BINARY_DIR}/deps/nccl" "CXX=${CMAKE_CXX_COMPILER}" CC="${CMAKE_CC_COMPILER}" + BUILD_COMMAND make src.build "${NCCL_BUILD_NVCC_GENCODE}" "CUDA_HOME=${CUDA_TOOLKIT_ROOT_DIR}" "BUILDDIR=${CMAKE_BINARY_DIR}/deps/nccl" "CXX=${CMAKE_CXX_COMPILER}" CC="${CMAKE_CC_COMPILER}" CXXFLAGS="-w" BUILD_IN_SOURCE 1 ) diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index f5d4c788af..f7c166f0dd 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -1,8 +1,8 @@ add_subdirectory(pcg) add_subdirectory(compiler) -# add_subdirectory(runtime) +add_subdirectory(runtime) add_subdirectory(op-attrs) add_subdirectory(kernels) add_subdirectory(utils) -# add_subdirectory(ffi) +add_subdirectory(ffi) add_subdirectory(substitutions) diff --git a/lib/compiler/ffi/include/flexflow/compiler.h b/lib/compiler/ffi/include/flexflow/compiler.h index b8c79aa584..14bd66150f 100644 --- a/lib/compiler/ffi/include/flexflow/compiler.h +++ b/lib/compiler/ffi/include/flexflow/compiler.h @@ -7,18 +7,35 @@ FLEXFLOW_FFI_BEGIN() -typedef enum { - FLEXFLOW_COMPILER_STATUS_OK, - FLEXFLOW_COMPILER_ERROR_UNKNOWN -} flexflow_compiler_error_t; - typedef enum { FLEXFLOW_SEARCH_ALGORITHM_DATA_PARALLEL } flexflow_search_algorithm_t; FF_NEW_OPAQUE_TYPE(flexflow_search_algorithm_config_t); FF_NEW_OPAQUE_TYPE(flexflow_search_result_t); -FF_NEW_OPAQUE_TYPE(fleflow_cost_estimator_t); +FF_NEW_OPAQUE_TYPE(flexflow_cost_estimator_t); + +// Error functions +typedef enum { + FLEXFLOW_COMPILER_STATUS_OK, + FLEXFLOW_COMPILER_ERROR_UNKNOWN +} flexflow_compiler_error_code_t; + +FF_NEW_OPAQUE_TYPE(flexflow_compiler_error_t); + +flexflow_error_t flexflow_compiler_error_wrap(flexflow_compiler_error_t); +flexflow_error_t flexflow_compiler_error_unwrap(flexflow_error_t, + flexflow_compiler_error_t *); +flexflow_error_t flexflow_compiler_error_is_ok(flexflow_compiler_error_t, + bool *); +flexflow_error_t flexflow_compiler_error_get_string(flexflow_compiler_error_t, + char **message); +flexflow_error_t + flexflow_compiler_error_get_error_code(flexflow_compiler_error_t, + flexflow_compiler_error_code_t *out); +flexflow_error_t flexflow_compiler_error_destroy(flexflow_compiler_error_t); + +// flexflow_error_t flexflow_computation_graph_optimize(flexflow_computation_graph_t, diff --git a/lib/compiler/ffi/internal/internal/compiler.h b/lib/compiler/ffi/internal/internal/compiler.h new file mode 100644 index 0000000000..f881844415 --- /dev/null +++ b/lib/compiler/ffi/internal/internal/compiler.h @@ -0,0 +1,6 @@ +#ifndef _FLEXFLOW_COMPILER_FFI_INTERNAL_INTERNAL_COMPILER_H +#define _FLEXFLOW_COMPILER_FFI_INTERNAL_INTERNAL_COMPILER_H + +#include "flexflow/compiler.h" + +#endif diff --git a/lib/ffi/CMakeLists.txt b/lib/ffi/CMakeLists.txt index cde77698f7..e441e6b84e 100644 --- a/lib/ffi/CMakeLists.txt +++ b/lib/ffi/CMakeLists.txt @@ -26,6 +26,8 @@ target_link_libraries( substitutions-ffi pcg-ffi op-attrs-ffi + PRIVATE + pcg-ffi-internal ) ff_set_cxx_properties(${target}) diff --git a/lib/ffi/include/flexflow/flexflow.h b/lib/ffi/include/flexflow/flexflow.h index ade96fff76..5b5eb1c473 100644 --- a/lib/ffi/include/flexflow/flexflow.h +++ b/lib/ffi/include/flexflow/flexflow.h @@ -5,13 +5,22 @@ #include "flexflow/op-attrs.h" #include "flexflow/pcg.h" #include "flexflow/runtime.h" +#include "flexflow/utils.h" #include +#include + +FLEXFLOW_FFI_BEGIN(); #define CHECK_FLEXFLOW(status) \ do { \ - if (flexflow_status_is_ok(status)) { \ + bool is_ok; \ + flexflow_status_is_ok(status, &is_ok); \ + if (is_ok) { \ + char *error_msg; \ + assert(flexflow_status_is_ok( \ + flexflow_get_error_string(status, &err_msg), &is_ok)); \ fprintf(stderr, \ - "FlexFlow encountered an errorat %s:%d : %s\n", \ + "FlexFlow encountered an error at %s:%d : %s\n", \ __FILE__, \ __LINE__, \ flexflow_get_error_string(status)); \ @@ -19,8 +28,11 @@ } \ } while (0) -bool flexflow_status_is_ok(flexflow_error_t); -char *flexflow_get_error_string(flexflow_error_t); -int flexflow_get_error_return_code(flexflow_error_t); +flexflow_error_t flexflow_status_is_ok(flexflow_error_t, bool *); +flexflow_error_t flexflow_get_error_string(flexflow_error_t, char *); + +flexflow_error_t flexflow_error_destroy(flexflow_error_t); + +FLEXFLOW_FFI_END(); #endif diff --git a/lib/ffi/src/ffi.cc b/lib/ffi/src/ffi.cc index a15fa598f0..f6cf67cf2c 100644 --- a/lib/ffi/src/ffi.cc +++ b/lib/ffi/src/ffi.cc @@ -1 +1,69 @@ #include "flexflow/flexflow.h" +#include "flexflow/op-attrs.h" +#include "flexflow/runtime.h" +#include "flexflow/utils.h" +#include "internal/pcg.h" + +flexflow_error_t flexflow_status_is_ok(flexflow_error_t e, bool *result) { + switch (e.error_source) { + case FLEXFLOW_ERROR_SOURCE_RUNTIME: { + flexflow_runtime_error_t err; + CHECK_FLEXFLOW(flexflow_runtime_error_unwrap(e, &err)); + CHECK_FLEXFLOW(flexflow_runtime_error_is_ok(err, result)); + } + case FLEXFLOW_ERROR_SOURCE_PCG: { + flexflow_pcg_error_t err; + CHECK_FLEXFLOW(flexflow_pcg_error_unwrap(e, &err)); + CHECK_FLEXFLOW(flexflow_pcg_error_is_ok(err, result)); + } + case FLEXFLOW_ERROR_SOURCE_COMPILER: { + flexflow_compiler_error_t err; + CHECK_FLEXFLOW(flexflow_compiler_error_unwrap(e, &err)); + CHECK_FLEXFLOW(flexflow_compiler_error_is_ok(err, result)); + } + case FLEXFLOW_ERROR_SOURCE_OPATTRS: { + flexflow_opattrs_error_t err; + CHECK_FLEXFLOW(flexflow_opattrs_error_unwrap(e, &err)); + CHECK_FLEXFLOW(flexflow_opattrs_error_is_ok(err, result)); + } + case FLEXFLOW_ERROR_SOURCE_UTILS: { + flexflow_utils_error_t err; + CHECK_FLEXFLOW(flexflow_utils_error_unwrap(e, &err)); + CHECK_FLEXFLOW(flexflow_utils_error_is_ok(err, result)); + } + default: + return flexflow_utils_error_create(FLEXFLOW_UTILS_INVALID_ERROR_SOURCE); + }; + + return flexflow_utils_error_create(FLEXFLOW_UTILS_STATUS_OK); +} + +flexflow_error_t flexflow_error_destroy(flexflow_error_t e) { + switch (e.error_source) { + case FLEXFLOW_ERROR_SOURCE_RUNTIME: { + flexflow_runtime_error_t err; + CHECK_FLEXFLOW(flexflow_runtime_error_unwrap(e, &err)); + return flexflow_runtime_error_destroy(err); + } + case FLEXFLOW_ERROR_SOURCE_PCG: { + flexflow_pcg_error_t err; + CHECK_FLEXFLOW(flexflow_pcg_error_unwrap(e, &err)); + return flexflow_pcg_error_destroy(err); + } + case FLEXFLOW_ERROR_SOURCE_COMPILER: { + flexflow_compiler_error_t err; + CHECK_FLEXFLOW(flexflow_compiler_error_unwrap(e, &err)); + return flexflow_compiler_error_destroy(err); + } + case FLEXFLOW_ERROR_SOURCE_OPATTRS: { + flexflow_opattrs_error_t err; + CHECK_FLEXFLOW(flexflow_opattrs_error_unwrap(e, &err)); + return flexflow_opattrs_error_destroy(err); + } + case FLEXFLOW_ERROR_SOURCE_UTILS: { + return flexflow_utils_error_create(FLEXFLOW_UTILS_STATUS_OK); + } + default: + return flexflow_utils_error_create(FLEXFLOW_UTILS_INVALID_ERROR_SOURCE); + } +} diff --git a/lib/op-attrs/ffi/CMakeLists.txt b/lib/op-attrs/ffi/CMakeLists.txt index a591dfe2e3..83fab85271 100644 --- a/lib/op-attrs/ffi/CMakeLists.txt +++ b/lib/op-attrs/ffi/CMakeLists.txt @@ -1,4 +1,5 @@ set(target op-attrs-ffi) +set(internal_target internal-${target}) project(${target}) file(GLOB_RECURSE SRC @@ -16,6 +17,7 @@ target_include_directories( include/ PRIVATE src/ + internal/ ) target_link_libraries( ${target} @@ -23,6 +25,22 @@ target_link_libraries( utils-ffi PRIVATE op-attrs + internal-utils-ffi ) ff_set_cxx_properties(${target}) + +add_library( + ${internal_target} + INTERFACE +) +target_link_libraries( + ${internal_target} + INTERFACE + ${target} +) +target_include_directories( + ${internal_target} + INTERFACE + internal/ +) diff --git a/lib/op-attrs/ffi/include/flexflow/op-attrs.h b/lib/op-attrs/ffi/include/flexflow/op-attrs.h index 1182ee67db..08b8e26f83 100644 --- a/lib/op-attrs/ffi/include/flexflow/op-attrs.h +++ b/lib/op-attrs/ffi/include/flexflow/op-attrs.h @@ -5,6 +5,30 @@ FLEXFLOW_FFI_BEGIN() +// Error handling + +typedef enum { + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_PARAM_SYNC_VALUE, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_DATATYPE_VALUE, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_ACTIVATION_VALUE, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_POOL_OP_VALUE, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_AGGREGATE_OP_VALUE, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_OP_TYPE_VALUE +} flexflow_opattrs_error_code_t; + +FF_NEW_OPAQUE_TYPE(flexflow_opattrs_error_t); +flexflow_error_t flexflow_opattrs_error_wrap(flexflow_opattrs_error_t); +flexflow_error_t flexflow_opattrs_error_unwrap(flexflow_error_t, + flexflow_opattrs_error_t *); +flexflow_error_t flexflow_opattrs_error_is_ok(flexflow_opattrs_error_t, bool *); +flexflow_error_t flexflow_opattrs_error_get_string(flexflow_opattrs_error_t, + char **); +flexflow_error_t flexflow_opattrs_error_destroy(flexflow_opattrs_error_t); + +// + +FF_NEW_OPAQUE_TYPE(flexflow_regularizer_attrs_t); + typedef enum { FLEXFLOW_DATATYPE_BOOL, FLEXFLOW_DATATYPE_INT32, @@ -29,7 +53,8 @@ typedef enum { typedef enum { FLEXFLOW_PARAM_SYNC_PARAMETER_SERVER, - FLEXFLOW_PARAM_SYNC_NCCL + FLEXFLOW_PARAM_SYNC_NCCL, + FLEXFLOW_PARAM_SYNC_NONE } flexflow_param_sync_t; typedef enum { @@ -37,11 +62,6 @@ typedef enum { FLEXFLOW_AGGREGATE_OP_AVG, } flexflow_aggregate_op_t; -typedef enum { - FLEXFLOW_OPATTRS_STATUS_OK, - FLEXFLOW_OPATTRS_ERROR_UNKNOWN -} flexflow_opattrs_error_t; - typedef enum { // does _not_ have to stay synchronized with op-attrs/op.h FLEXFLOW_OP_TYPE_NOOP, FLEXFLOW_OP_TYPE_INPUT, diff --git a/lib/op-attrs/ffi/internal/internal/op-attrs.h b/lib/op-attrs/ffi/internal/internal/op-attrs.h new file mode 100644 index 0000000000..df0d6ce61c --- /dev/null +++ b/lib/op-attrs/ffi/internal/internal/op-attrs.h @@ -0,0 +1,36 @@ +#ifndef _FLEXFLOW_OPATTRS_FFI_INTERNAL_INTERNAL_OPATTRS_H +#define _FLEXFLOW_OPATTRS_FFI_INTERNAL_INTERNAL_OPATTRS_H + +#include "flexflow/op-attrs.h" +#include "internal/opaque.h" +#include "op-attrs/activation.h" +#include "op-attrs/datatype.h" +#include "op-attrs/op.h" +#include "op-attrs/ops/embedding.h" +#include "op-attrs/ops/linear.h" +#include "op-attrs/ops/pool_2d.h" +#include "op-attrs/param_sync.h" + +using namespace FlexFlow; + +REGISTER_OPAQUE(flexflow_regularizer_attrs_t, optional); + +optional to_internal(flexflow_param_sync_t); +flexflow_param_sync_t to_external(optional); + +DataType to_internal(flexflow_datatype_t); +flexflow_datatype_t to_external(DataType); + +optional to_internal(flexflow_activation_t); +flexflow_activation_t to_external(optional); + +PoolOp to_internal(flexflow_pool_op_t e); +flexflow_pool_op_t to_external(PoolOp i); + +AggregateOp to_internal(flexflow_aggregate_op_t e); +flexflow_aggregate_op_t to_external(AggregateOp i); + +OperatorType to_internal(flexflow_op_type_t e); +flexflow_op_type_t to_external(OperatorType i); + +#endif diff --git a/lib/op-attrs/ffi/src/op-attrs.cc b/lib/op-attrs/ffi/src/op-attrs.cc index 265601b36b..828574dc25 100644 --- a/lib/op-attrs/ffi/src/op-attrs.cc +++ b/lib/op-attrs/ffi/src/op-attrs.cc @@ -1 +1,184 @@ #include "flexflow/op-attrs.h" +#include "internal/enums.h" +#include "internal/error.h" +#include "internal/op-attrs.h" +#include "op-attrs/op.h" +#include "op-attrs/ops/embedding.h" +#include "utils/bidict.h" + +flexflow_utils_exception_t + make_opattrs_exception(flexflow_opattrs_error_code_t); + +REGISTER_FFI_ENUM(flexflow_param_sync_t, + ParamSync, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_PARAM_SYNC_VALUE, + {{FLEXFLOW_PARAM_SYNC_PARAMETER_SERVER, ParamSync::PS}, + {FLEXFLOW_PARAM_SYNC_NCCL, ParamSync::NCCL}}); + +REGISTER_FFI_ENUM(flexflow_datatype_t, + DataType, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_DATATYPE_VALUE, + {{FLEXFLOW_DATATYPE_BOOL, DataType::BOOL}, + {FLEXFLOW_DATATYPE_INT32, DataType::INT32}, + {FLEXFLOW_DATATYPE_INT64, DataType::INT64}, + {FLEXFLOW_DATATYPE_HALF, DataType::HALF}, + {FLEXFLOW_DATATYPE_FLOAT, DataType::FLOAT}, + {FLEXFLOW_DATATYPE_DOUBLE, DataType::DOUBLE}}); + +REGISTER_FFI_ENUM(flexflow_activation_t, + optional, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_ACTIVATION_VALUE, + {{FLEXFLOW_ACTIVATION_RELU, Activation::RELU}, + {FLEXFLOW_ACTIVATION_SIGMOID, Activation::SIGMOID}, + {FLEXFLOW_ACTIVATION_TANH, Activation::TANH}, + {FLEXFLOW_ACTIVATION_GELU, Activation::GELU}, + {FLEXFLOW_ACTIVATION_NONE, nullopt}}); + +REGISTER_FFI_ENUM(flexflow_pool_op_t, + PoolOp, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_POOL_OP_VALUE, + {{FLEXFLOW_POOL_OP_MAX, PoolOp::MAX}, + {FLEXFLOW_POOL_OP_AVG, PoolOp::AVG}}); + +REGISTER_FFI_ENUM(flexflow_aggregate_op_t, + AggregateOp, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_AGGREGATE_OP_VALUE, + {{FLEXFLOW_AGGREGATE_OP_SUM, AggregateOp::SUM}, + {FLEXFLOW_AGGREGATE_OP_AVG, AggregateOp::AVG}}); + +REGISTER_FFI_ENUM(flexflow_op_type_t, + OperatorType, + FLEXFLOW_OPATTRS_ERROR_CODE_INVALID_OP_TYPE_VALUE, + { + {FLEXFLOW_OP_TYPE_NOOP, Op::NOOP}, + {FLEXFLOW_OP_TYPE_INPUT, Op::INPUT}, + {FLEXFLOW_OP_TYPE_WEIGHT, Op::WEIGHT}, + {FLEXFLOW_OP_TYPE_CONV2D, Op::CONV2D}, + {FLEXFLOW_OP_TYPE_DROPOUT, Op::DROPOUT}, + {FLEXFLOW_OP_TYPE_LINEAR, Op::LINEAR}, + {FLEXFLOW_OP_TYPE_BATCHMATMUL, Op::BATCHMATMUL}, + {FLEXFLOW_OP_TYPE_POOL2D, Op::POOL2D}, + {FLEXFLOW_OP_TYPE_SCALAR_MULTIPLY, Op::SCALAR_MULTIPLY}, + {FLEXFLOW_OP_TYPE_SCALAR_ADD, Op::SCALAR_ADD}, + {FLEXFLOW_OP_TYPE_SCALAR_FLOOR_DIV, Op::SCALAR_FLOOR_DIV}, + {FLEXFLOW_OP_TYPE_SCALAR_TRUE_DIV, Op::SCALAR_TRUE_DIV}, + {FLEXFLOW_OP_TYPE_SCALAR_SUB, Op::SCALAR_SUB}, + {FLEXFLOW_OP_TYPE_RELU, Op::RELU}, + {FLEXFLOW_OP_TYPE_IDENTITY, Op::IDENTITY}, + {FLEXFLOW_OP_TYPE_SIGMOID, Op::SIGMOID}, + {FLEXFLOW_OP_TYPE_TANH, Op::TANH}, + {FLEXFLOW_OP_TYPE_ELU, Op::ELU}, + {FLEXFLOW_OP_TYPE_FLAT, Op::FLAT}, + {FLEXFLOW_OP_TYPE_SOFTMAX, Op::SOFTMAX}, + {FLEXFLOW_OP_TYPE_BATCHNORM, Op::BATCHNORM}, + {FLEXFLOW_OP_TYPE_CONCAT, Op::CONCAT}, + {FLEXFLOW_OP_TYPE_SPLIT, Op::SPLIT}, + {FLEXFLOW_OP_TYPE_EMBEDDING, Op::EMBEDDING}, + {FLEXFLOW_OP_TYPE_GROUP_BY, Op::GROUP_BY}, + {FLEXFLOW_OP_TYPE_CACHE, Op::CACHE}, + {FLEXFLOW_OP_TYPE_AGGREGATE, Op::AGGREGATE}, + {FLEXFLOW_OP_TYPE_AGG_SPEC, Op::AGG_SPEC}, + {FLEXFLOW_OP_TYPE_RESHAPE, Op::RESHAPE}, + {FLEXFLOW_OP_TYPE_REVERSE, Op::REVERSE}, + {FLEXFLOW_OP_TYPE_TRANSPOSE, Op::TRANSPOSE}, + {FLEXFLOW_OP_TYPE_EW_ADD, Op::EW_ADD}, + {FLEXFLOW_OP_TYPE_EW_MUL, Op::EW_MUL}, + {FLEXFLOW_OP_TYPE_MATMUL, Op::MATMUL}, + {FLEXFLOW_OP_TYPE_MUL, Op::MUL}, + {FLEXFLOW_OP_TYPE_ENLARGE, Op::ENLARGE}, + {FLEXFLOW_OP_TYPE_SQUEEZE, Op::SQUEEZE}, + {FLEXFLOW_OP_TYPE_UNSQUEEZE, Op::UNSQUEEZE}, + {FLEXFLOW_OP_TYPE_EW_SUB, Op::EW_SUB}, + {FLEXFLOW_OP_TYPE_EW_DIV, Op::EW_DIV}, + {FLEXFLOW_OP_TYPE_EW_EQUAL, Op::EW_EQUAL}, + {FLEXFLOW_OP_TYPE_EW_GREATER, Op::EW_GREATER}, + {FLEXFLOW_OP_TYPE_EW_LESS, Op::EW_LESS}, + {FLEXFLOW_OP_TYPE_EW_MAX, Op::EW_MAX}, + {FLEXFLOW_OP_TYPE_EW_MIN, Op::EW_MIN}, + {FLEXFLOW_OP_TYPE_REDUCE_ARGMAX, Op::REDUCE_ARGMAX}, + {FLEXFLOW_OP_TYPE_REDUCE_ARGMIN, Op::REDUCE_ARGMIN}, + {FLEXFLOW_OP_TYPE_REDUCE_MAX, Op::REDUCE_MAX}, + {FLEXFLOW_OP_TYPE_REDUCE_MEAN, Op::REDUCE_MEAN}, + {FLEXFLOW_OP_TYPE_REDUCE_MIN, Op::REDUCE_MIN}, + {FLEXFLOW_OP_TYPE_REDUCE_PROD, Op::REDUCE_PROD}, + {FLEXFLOW_OP_TYPE_REDUCE_SUM, Op::REDUCE_SUM}, + {FLEXFLOW_OP_TYPE_PAD, Op::PAD}, + {FLEXFLOW_OP_TYPE_SHAPE, Op::SHAPE}, + {FLEXFLOW_OP_TYPE_SIZE, Op::SIZE}, + {FLEXFLOW_OP_TYPE_TOPK, Op::TOPK}, + {FLEXFLOW_OP_TYPE_WHERE, Op::WHERE}, + {FLEXFLOW_OP_TYPE_CEIL, Op::CEIL}, + {FLEXFLOW_OP_TYPE_CAST, Op::CAST}, + {FLEXFLOW_OP_TYPE_EXP, Op::EXP}, + {FLEXFLOW_OP_TYPE_ROUND, Op::ROUND}, + {FLEXFLOW_OP_TYPE_LOG, Op::LOG}, + {FLEXFLOW_OP_TYPE_LOGICAL_NOT, Op::LOGICAL_NOT}, + {FLEXFLOW_OP_TYPE_SQRT, Op::SQRT}, + {FLEXFLOW_OP_TYPE_SIN, Op::SIN}, + {FLEXFLOW_OP_TYPE_COS, Op::COS}, + {FLEXFLOW_OP_TYPE_LEAKYRELU, Op::LEAKYRELU}, + {FLEXFLOW_OP_TYPE_SLICE, Op::SLICE}, + {FLEXFLOW_OP_TYPE_RESIZE, Op::RESIZE}, + {FLEXFLOW_OP_TYPE_PRELU, Op::PRELU}, + {FLEXFLOW_OP_TYPE_GELU, Op::GELU}, + {FLEXFLOW_OP_TYPE_MULTIHEAD_ATTENTION, + Op::MULTIHEAD_ATTENTION}, + {FLEXFLOW_OP_TYPE_FUSED, Op::FUSED}, + {FLEXFLOW_OP_TYPE_RSQRT, Op::RSQRT}, + {FLEXFLOW_OP_TYPE_POW, Op::POW}, + {FLEXFLOW_OP_TYPE_MEAN, Op::MEAN}, + {FLEXFLOW_OP_TYPE_LAYERNORM, Op::LAYERNORM}, + {FLEXFLOW_OP_TYPE_GATHER, Op::GATHER}, + {FLEXFLOW_OP_TYPE_BROADCAST, Op::BROADCAST}, + {FLEXFLOW_OP_TYPE_REPARTITION, Op::REPARTITION}, + {FLEXFLOW_OP_TYPE_COMBINE, Op::COMBINE}, + {FLEXFLOW_OP_TYPE_REPLICATE, Op::REPLICATE}, + {FLEXFLOW_OP_TYPE_REDUCTION, Op::REDUCTION}, + {FLEXFLOW_OP_TYPE_BATCH, Op::BATCH}, + {FLEXFLOW_OP_TYPE_PIPELINE, Op::PIPELINE}, + {FLEXFLOW_OP_TYPE_FUSED_PARALLEL, Op::FUSED_PARALLEL}, + }); + +flexflow_error_t make_opattrs_error(flexflow_opattrs_error_code_t); + +ParamSync to_internal(flexflow_param_sync_t e) { + return to_internal_impl(e); +} +flexflow_param_sync_t to_external(ParamSync i) { + return to_external_impl(i); +} + +DataType to_internal(flexflow_datatype_t e) { + return to_internal_impl(e); +} +flexflow_datatype_t to_external(DataType i) { + return to_external_impl(i); +} + +optional to_internal(flexflow_activation_t e) { + return to_internal_impl(e); +} +flexflow_activation_t to_external(optional i) { + return to_external_impl(i); +} + +PoolOp to_internal(flexflow_pool_op_t e) { + return to_internal_impl(e); +} +flexflow_pool_op_t to_external(PoolOp i) { + return to_external_impl(i); +} + +AggregateOp to_internal(flexflow_aggregate_op_t e) { + return to_internal_impl(e); +} +flexflow_aggregate_op_t to_external(AggregateOp i) { + return to_external_impl(i); +} + +OperatorType to_internal(flexflow_op_type_t e) { + return to_internal_impl(e); +} +flexflow_op_type_t to_external(OperatorType i) { + return to_external_impl(i); +} diff --git a/lib/op-attrs/include/op-attrs/ff_dim.h b/lib/op-attrs/include/op-attrs/ff_dim.h index be1f148a70..298e134e08 100644 --- a/lib/op-attrs/include/op-attrs/ff_dim.h +++ b/lib/op-attrs/include/op-attrs/ff_dim.h @@ -2,6 +2,7 @@ #define _FLEXFLOW_OPATTRS_INCLUDE_FF_DIM_H #include "utils/strong_typedef.h" +#include "utils/type_traits.h" #include namespace FlexFlow { diff --git a/lib/op-attrs/include/op-attrs/tensor_shape.h b/lib/op-attrs/include/op-attrs/tensor_shape.h index fa34860817..7ad4d06dbd 100644 --- a/lib/op-attrs/include/op-attrs/tensor_shape.h +++ b/lib/op-attrs/include/op-attrs/tensor_shape.h @@ -11,7 +11,7 @@ namespace FlexFlow { using TensorDims = FFOrdered; -struct TensorShape : public use_visitable_cmp { +struct TensorShape { TensorShape() = delete; template @@ -25,10 +25,10 @@ struct TensorShape : public use_visitable_cmp { TensorDims dims; DataType data_type; }; +FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(TensorShape, dims, data_type); -} // namespace FlexFlow +DataType get_data_type(TensorShape const &); -VISITABLE_STRUCT(::FlexFlow::TensorShape, dims, data_type); -MAKE_VISIT_HASHABLE(::FlexFlow::TensorShape); +} // namespace FlexFlow #endif diff --git a/lib/pcg/ffi/CMakeLists.txt b/lib/pcg/ffi/CMakeLists.txt index 714bb0f081..a1af130e9c 100644 --- a/lib/pcg/ffi/CMakeLists.txt +++ b/lib/pcg/ffi/CMakeLists.txt @@ -1,4 +1,5 @@ set(target pcg-ffi) +set(internal_target pcg-ffi-internal) project(${target}) file(GLOB_RECURSE SRC @@ -15,15 +16,34 @@ target_include_directories( PUBLIC include/ PRIVATE + internal/ src/ ) target_link_libraries( ${target} - PUBLIC + PRIVATE pcg + internal-utils-ffi + internal-op-attrs-ffi + PUBLIC utils-ffi op-attrs-ffi ) +add_library( + ${internal_target} + INTERFACE +) +target_link_libraries( + ${internal_target} + INTERFACE + ${target} +) +target_include_directories( + ${internal_target} + INTERFACE + internal/ +) + define_ff_vars(${target}) ff_set_cxx_properties(${target}) diff --git a/lib/pcg/ffi/include/flexflow/pcg.h b/lib/pcg/ffi/include/flexflow/pcg.h index 98e3d8db8e..03dee5ec95 100644 --- a/lib/pcg/ffi/include/flexflow/pcg.h +++ b/lib/pcg/ffi/include/flexflow/pcg.h @@ -8,11 +8,6 @@ FLEXFLOW_FFI_BEGIN(); -typedef enum { - FLEXFLOW_PCG_STATUS_OK, - FLEXFLOW_PCG_ERROR_UNKNOWN, -} flexflow_pcg_error_t; - typedef enum { FLEXFLOW_DEVICE_TYPE_CPU, FLEXFLOW_DEVICE_TYPE_GPU, @@ -30,22 +25,49 @@ FF_NEW_OPAQUE_TYPE(flexflow_optimizer_t); FF_NEW_OPAQUE_TYPE(flexflow_machine_specification_t); FF_NEW_OPAQUE_TYPE(flexflow_model_compilation_input_t); FF_NEW_OPAQUE_TYPE(flexflow_model_compilation_result_t); +FF_NEW_OPAQUE_TYPE(flexflow_pcg_error_t); +FF_NEW_OPAQUE_TYPE(flexflow_tensor_list_t); +FF_NEW_OPAQUE_TYPE(flexflow_int_list_t); + +typedef enum { + FLEXFLOW_PCG_STATUS_INVALID_ERROR_CODE, + FLEXFLOW_PCG_STATUS_INVALID_FILE_PTR, + FLEXFLOW_PCG_STATUS_FILE_WRITE_FAILED, + FLEXFLOW_PCG_STATUS_FILE_READ_FAILED, + FLEXFLOW_PCG_STATUS_TENSOR_LIST_ACCESS_OUT_OF_BOUNDS, + FLEXFLOW_PCG_STATUS_NEGATIVE_ARRAY_LENGTH_FOUND, + FLEXFLOW_PCG_ERROR_UNKNOWN, +} flexflow_pcg_error_code_t; extern flexflow_initializer_t NO_INITIALIZER; +extern flexflow_regularizer_attrs_t NO_REGULARIZER; -flexflow_error_t make_pcg_error(flexflow_pcg_error_t); +flexflow_error_t flexflow_pcg_error_wrap(flexflow_pcg_error_t); +flexflow_error_t flexflow_pcg_error_unwrap(flexflow_error_t, + flexflow_pcg_error_t *); +flexflow_error_t flexflow_pcg_error_is_ok(flexflow_pcg_error_t, bool *); +flexflow_error_t flexflow_pcg_error_get_string(flexflow_pcg_error_t, char **); +flexflow_error_t flexflow_pcg_error_get_error_code(flexflow_pcg_error_t, + flexflow_pcg_error_code_t *); +flexflow_error_t flexflow_pcg_error_destroy(flexflow_pcg_error_t); -char *flexflow_pcg_get_error_string(flexflow_pcg_error_t); +flexflow_error_t flexflow_tensor_list_get_num_elements(flexflow_tensor_list_t, + size_t *out); +flexflow_error_t flexflow_tensor_list_get_element(flexflow_tensor_list_t, + size_t, + flexflow_tensor_t *out); +flexflow_error_t flexflow_tensor_list_destroy(flexflow_tensor_list_t); flexflow_error_t flexflow_computation_graph_create(flexflow_computation_graph_t *out); flexflow_error_t flexflow_computation_graph_destroy(flexflow_computation_graph_t); -flexflow_error_t flexflow_computation_graph_serialize_to_buf( - flexflow_computation_graph_t, void *buf, size_t buf_size); -flexflow_error_t flexflow_computation_graph_deserialize_from_buf( - flexflow_computation_graph_t, void *buf, size_t buf_size); +flexflow_error_t + flexflow_computation_graph_serialize_to_buffer(flexflow_computation_graph_t, + char **out); +flexflow_error_t flexflow_computation_graph_deserialize_from_buffer( + char *buf, flexflow_computation_graph_t *out); flexflow_error_t flexflow_computation_graph_serialize_to_file(flexflow_computation_graph_t, @@ -67,9 +89,9 @@ flexflow_error_t flexflow_tensor_get_sync_type(flexflow_tensor_t, flexflow_error_t flexflow_tensor_get_datatype(flexflow_tensor_t, flexflow_datatype_t *out); flexflow_error_t flexflow_tensor_get_num_dims(flexflow_tensor_t, int *out); -flexflow_error_t flexflow_tensor_get_dims(flexflow_tensor_t, int *out); -flexflow_error_t flexflow_tensor_destroy(flexflow_computation_graph_t, - flexflow_tensor_t); +flexflow_error_t flexflow_tensor_get_dims(flexflow_tensor_t, + flexflow_int_list_t *out); +flexflow_error_t flexflow_tensor_destroy(flexflow_tensor_t); flexflow_error_t flexflow_computation_graph_add_op_exp(flexflow_computation_graph_t, @@ -120,14 +142,14 @@ flexflow_error_t flexflow_error_t flexflow_computation_graph_add_op_pow(flexflow_computation_graph_t, flexflow_tensor_t, - flexflow_tensor_t *out, float exponent, + flexflow_tensor_t *out, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_scalar_multiply( flexflow_computation_graph_t, flexflow_tensor_t, - flexflow_tensor_t *out, float scalar, + flexflow_tensor_t *out, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_scalar_add(flexflow_computation_graph_t, @@ -203,12 +225,14 @@ flexflow_error_t flexflow_computation_graph_add_op_conv2d( bool use_bias = true, flexflow_initializer_t kernel_initializer = NO_INITIALIZER, flexflow_initializer_t bias_initializer = NO_INITIALIZER, + flexflow_regularizer_attrs_t kernel_regularizer = NO_REGULARIZER, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_dropout(flexflow_computation_graph_t, flexflow_tensor_t, float rate, - unsigned long long seed, + flexflow_tensor_t *out, + unsigned long long seed = 0, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_embedding( flexflow_computation_graph_t, @@ -216,22 +240,24 @@ flexflow_error_t flexflow_computation_graph_add_op_embedding( int num_entries, int out_dim, flexflow_aggregate_op_t aggr_op, - flexflow_datatype_t output_type, + flexflow_tensor_t *out, + flexflow_datatype_t output_type = FLEXFLOW_DATATYPE_FLOAT, flexflow_initializer_t initializer = NO_INITIALIZER, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_gather(flexflow_computation_graph_t, - flexflow_tensor_t, - flexflow_tensor_t *outs, + flexflow_tensor_t input, + flexflow_tensor_t index, int dim, + flexflow_tensor_list_t *out, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_group_by(flexflow_computation_graph_t, flexflow_tensor_t data, flexflow_tensor_t assign, - flexflow_tensor_t *outs, int n, float alpha, + flexflow_tensor_list_t *out, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_cache( flexflow_computation_graph_t, @@ -251,9 +277,9 @@ flexflow_error_t flexflow_computation_graph_add_op_aggregate( flexflow_tensor_t true_gate_assign, flexflow_tensor_t full_gate_gradients, flexflow_tensor_t *exp_preds, - flexflow_tensor_t *out, int n, float lambda_bal, + flexflow_tensor_t *out, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_aggregate_spec( flexflow_computation_graph_t, @@ -280,6 +306,7 @@ flexflow_error_t flexflow_tensor_t, flexflow_tensor_t *out, int *axes, + int num_axes, bool elementwise_affine, float eps, char *name = NULL); @@ -326,6 +353,7 @@ flexflow_error_t flexflow_tensor_t, flexflow_tensor_t *out, int *dims, + int num_dims, bool keepdims, char *name = NULL); flexflow_error_t @@ -341,8 +369,9 @@ flexflow_error_t flexflow_error_t flexflow_computation_graph_add_op_split(flexflow_computation_graph_t, flexflow_tensor_t, - flexflow_tensor_t *outs, + flexflow_tensor_list_t *out, int *splits, + int num_splits, int axis, char *name = NULL); flexflow_error_t @@ -359,13 +388,16 @@ flexflow_error_t flexflow_error_t flexflow_computation_graph_add_op_transpose(flexflow_computation_graph_t, flexflow_tensor_t, + flexflow_tensor_t *out, int *permutation, + int num_permutation_values, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_reduce_sum(flexflow_computation_graph_t, flexflow_tensor_t, flexflow_tensor_t *out, int *axes, + int num_axes, bool keepdims = false, char *name = NULL); flexflow_error_t @@ -373,6 +405,7 @@ flexflow_error_t flexflow_tensor_t, flexflow_tensor_t *out, int *shape, + int num_shape_entries, char *name = NULL); flexflow_error_t flexflow_computation_graph_add_op_reverse(flexflow_computation_graph_t, @@ -383,15 +416,16 @@ flexflow_error_t flexflow_error_t flexflow_computation_graph_add_op_topk(flexflow_computation_graph_t, flexflow_tensor_t, - flexflow_tensor_t *outs, + flexflow_tensor_list_t *out, int k, bool sorted, char *name = NULL); -flexflow_error_t flexflow_computation_graph_add_multihead_attention( +flexflow_error_t flexflow_computation_graph_add_op_multihead_attention( flexflow_computation_graph_t, flexflow_tensor_t query, flexflow_tensor_t key, flexflow_tensor_t value, + flexflow_tensor_t *output, int embed_dim, int num_heads, int kdim = 0, @@ -399,6 +433,7 @@ flexflow_error_t flexflow_computation_graph_add_multihead_attention( float dropout = 0.0f, bool bias = true, bool add_bias_kv = false, + bool add_zero_attn = false, flexflow_initializer_t initializer = NO_INITIALIZER, char *name = NULL); diff --git a/lib/pcg/ffi/internal/internal/pcg.h b/lib/pcg/ffi/internal/internal/pcg.h new file mode 100644 index 0000000000..56e9c9a812 --- /dev/null +++ b/lib/pcg/ffi/internal/internal/pcg.h @@ -0,0 +1,5 @@ +#include "flexflow/pcg.h" + +/* struct flexflow_pcg_error_t { */ +/* flexflow_pcg_error_code_t err_code; */ +/* }; */ diff --git a/lib/pcg/ffi/src/pcg.cc b/lib/pcg/ffi/src/pcg.cc index 4e1bd55bc3..c958ed6c18 100644 --- a/lib/pcg/ffi/src/pcg.cc +++ b/lib/pcg/ffi/src/pcg.cc @@ -1,2 +1,936 @@ +#include "pcg.h" #include "flexflow/pcg.h" -#include "pcg/model_compilation.h" +#include "flexflow/utils.h" +#include "internal/error.h" +#include "internal/error_handling.h" +#include "internal/op-attrs.h" +#include "internal/opaque.h" +#include "pcg/computation_graph.h" +#include "pcg/computation_graph_builder.h" +#include "pcg/file_format/v1/graphs.h" +#include "utils/exception.h" + +using namespace FlexFlow; + +flexflow_error_t flexflow_pcg_error_wrap(flexflow_pcg_error_t e) { + return flexflow_error_wrap(FLEXFLOW_ERROR_SOURCE_PCG, *unwrap_opaque(e)); +} + +flexflow_error_t make_pcg_error(flexflow_pcg_error_code_t); +flexflow_ffi_exception_t make_pcg_exception(flexflow_pcg_error_code_t); + +flexflow_error_t flexflow_pcg_error_unwrap(flexflow_error_t err, + flexflow_pcg_error_t *out) { + return flexflow_error_unwrap(err, FLEXFLOW_ERROR_SOURCE_PCG, out); +} + +flexflow_error_t flexflow_pcg_error_is_ok(flexflow_pcg_error_t, bool *out) { + *out = false; + return status_ok(); +} + +template +std::vector make_vector(T const *ptr, int num_values) { + if (num_values < 0) { + throw make_pcg_exception(FLEXFLOW_PCG_STATUS_NEGATIVE_ARRAY_LENGTH_FOUND); + } + + return {ptr, ptr + num_values}; +} + +flexflow_error_t flexflow_pcg_error_get_string(flexflow_pcg_error_t err, + char **m_out) { + flexflow_pcg_error_code_t err_code; + RAISE_FLEXFLOW(flexflow_pcg_error_get_error_code(err, &err_code)); + + auto out = const_cast(m_out); + + switch (err_code) { + case FLEXFLOW_PCG_ERROR_UNKNOWN: { + *out = "Unknown error"; + break; + } + default: + return make_pcg_error(FLEXFLOW_PCG_STATUS_INVALID_ERROR_CODE); + } + + return status_ok(); +} + +flexflow_error_t + flexflow_pcg_error_get_error_code(flexflow_error_t err, + flexflow_pcg_error_code_t *out) { + flexflow_pcg_error_t opaque; + RAISE_FLEXFLOW(flexflow_pcg_error_unwrap(err, &opaque)); + internal_flexflow_pcg_error_t const *unwrapped = unwrap_opaque(opaque); + *out = unwrapped->err_code; + + return status_ok(); +} + +flexflow_error_t flexflow_pcg_error_destroy(flexflow_pcg_error_t) { + return status_ok(); +} + +flexflow_error_t + flexflow_computation_graph_create(flexflow_computation_graph_t *out) { + return handle_errors(out, [&] { return ComputationGraph{}; }); +} + +flexflow_error_t + flexflow_computation_graph_destroy(flexflow_computation_graph_t opaque) { + return handle_errors([&] {}); +} + +template +std::string get_v1_string(T const &t) { + return json{to_v1(t)}.dump(); +} + +template +T from_v1_string(char const *s) { + return json::parse(s) + .template get>()))>(); +} + +flexflow_error_t flexflow_computation_graph_serialize_to_buffer( + flexflow_computation_graph_t opaque, char **out) { + return handle_errors([&] { + ComputationGraph const *cg = unwrap_opaque(opaque); + + std::string json_str = get_v1_string(*cg); + + *out = new char[json_str.size() + 1]; + strncpy(*out, json_str.c_str(), json_str.size() + 1); + assert((*out)[json_str.size()] == '\x00'); + }); +} + +flexflow_error_t flexflow_computation_graph_deserialize_from_buffer( + char *buf, flexflow_computation_graph_t *out) { + return handle_errors(out, + [&] { return from_v1_string(buf); }); +} + +flexflow_error_t flexflow_computation_graph_serialize_to_file( + flexflow_computation_graph_t cg, FILE *f) { + return handle_errors([&] { + if (f == nullptr) { + throw make_pcg_exception(FLEXFLOW_PCG_STATUS_INVALID_FILE_PTR); + } + + std::string json_str = get_v1_string(*unwrap_opaque(cg)); + + size_t num_bytes_written = + fwrite(json_str.c_str(), sizeof(char), json_str.size(), f); + if (num_bytes_written < json_str.size()) { + throw make_pcg_exception(FLEXFLOW_PCG_STATUS_FILE_WRITE_FAILED); + } + }); +} + +flexflow_error_t flexflow_computation_graph_deserialize_from_file( + FILE *f, flexflow_computation_graph_t *out) { + return handle_errors(out, [&] { + size_t constexpr BUF_SIZE = 256; + + char buf[BUF_SIZE]; + std::ostringstream oss; + + while (true) { + size_t num_bytes_read = fread(buf, sizeof(char), BUF_SIZE, f); + oss.write(buf, num_bytes_read); + if (num_bytes_read < BUF_SIZE) { + if (feof(f)) { + break; + } else { + throw make_pcg_exception(FLEXFLOW_PCG_STATUS_FILE_READ_FAILED); + } + } + } + + return from_v1_string(oss.str().c_str()); + }); +} + +flexflow_error_t flexflow_tensor_create(flexflow_computation_graph_t cg, + int num_dims, + int *dims, + flexflow_datatype_t datatype, + bool create_grad, + flexflow_tensor_t *out) { + return handle_errors(out, [&] { + TensorDims ordered_dims{make_vector(dims, num_dims)}; + TensorShape shape = {ordered_dims, to_internal(datatype)}; + return create_tensor(deref_opaque(cg), shape, create_grad); + }); +} + +flexflow_error_t flexflow_tensor_get_create_grad(flexflow_tensor_t opaque, + bool *out) { + return handle_errors(out, + [&] { return c_deref_opaque(opaque).create_gradients; }); +} + +flexflow_error_t flexflow_tensor_get_initializer(flexflow_tensor_t opaque, + flexflow_initializer_t *out) { + return handle_errors(out, [&] { return c_deref_opaque(opaque).initializer; }); +} + +flexflow_error_t flexflow_tensor_get_sync_type(flexflow_tensor_t opaque, + flexflow_param_sync_t *out) { + return handle_errors(out, [&] { return c_deref_opaque(opaque).sync_type; }); +} + +flexflow_error_t flexflow_tensor_get_datatype(flexflow_tensor_t opaque, + flexflow_datatype_t *out) { + return handle_errors(out, [&] { return c_deref_opaque(opaque).data_type; }); +} + +flexflow_error_t flexflow_tensor_get_num_dims(flexflow_tensor_t opaque, + int *out) { + return handle_errors(out, [&] { return c_deref_opaque(opaque).num_dims(); }); +} + +flexflow_error_t flexflow_tensor_get_dims(flexflow_tensor_t opaque, int *out) { + return handle_errors(out, [&] { return c_deref_opaque(opaque).dims; }); +} + +flexflow_error_t flexflow_tensor_destroy(flexflow_tensor_t opaque) { + return handle_errors([&] { delete_opaque(opaque); }); +} + +optional maybe_string(char *); +using BinaryFunc = std::function const &)>; +using UnaryFunc = std::function)>; +using ScalarUnaryFunc = std::function)>; + +#define BINARY_OP(insert_func_layer) \ + [](ComputationGraph &cg, \ + Tensor const &lhs, \ + Tensor const &rhs, \ + optional const &name) -> Tensor { \ + return func(cg, lhs, rhs, name); \ + } + +#define BINARY_OP(insert_func_layer) \ + [](ComputationGraph &cg, Tensor const &t, optional const &name) \ + -> Tensor { return func(cg, t, name); } + +#define SCALAR_BINARY_OP(insert_func_layer) \ + [](ComputationGraph &cg, \ + Tensor const &t, \ + float scalar, \ + optional const &name) -> Tensor { \ + return func(cg, t, scalar, name); \ + } + +flexflow_error_t add_op(BinaryFunc const &f, + flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + char *name) { + return handle_errors(out, [&] { + return f(deref_opaque(cg), + c_deref_opaque(lhs), + c_deref_opaque(rhs), + maybe_string(name)); + }); +} + +flexflow_error_t add_op(UnaryFunc const &f, + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return handle_errors(out, [&] { + return f(deref_opaque(cg), c_deref_opaque(input), maybe_string(name)); + }); +} + +flexflow_error_t add_op(ScalarUnaryFunc const &f, + flexflow_computation_graph_t opaque_cg, + flexflow_tensor_t opaque_input, + float scalar, + flexflow_tensor_t *out, + char *raw_name) { + UnaryFunc ff = [&](ComputationGraph &cg, + Tensor const &input, + optional const &name) -> Tensor { + return f(cg, input, scalar, name); + }; + return add_op(ff, opaque_cg, opaque_input, out, raw_name); +} + +flexflow_error_t flexflow_computation_graph_add_op_exp( + flexflow_computation_graph_t opaque_cg, + flexflow_tensor_t opaque_input, + flexflow_tensor_t *out, + char *name) { + return add_op( + BINARY_OP(insert_exp_layer), opaque_cg, opaque_input, out, name); +} + +flexflow_error_t flexflow_computation_graph_add_op_add( + flexflow_computation_graph_t opaque_cg, + flexflow_tensor_t opaque_lhs, + flexflow_tensor_t opaque_rhs, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_add_layer), + opaque_cg, + opaque_lhs, + opaque_rhs, + out, + name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_subtract(flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_subtract_layer), cg, lhs, rhs, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_multiply(flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_multiply_layer), cg, lhs, rhs, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_divide(flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_divide_layer), cg, lhs, rhs, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_max(flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_max_layer), cg, lhs, rhs, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_min(flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_min_layer), cg, lhs, rhs, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_rsqrt(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_rsqrt_layer), cg, input, out, name); +} + +flexflow_error_t flexflow_computation_graph_add_op_scalar_multiply( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + float scalar, + flexflow_tensor_t *out, + char *name) { + return add_op(SCALAR_BINARY_OP(insert_scalar_multiply_layer), + cg, + input, + scalar, + out, + name); +} + +flexflow_error_t flexflow_computation_graph_add_op_scalar_add( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + float scalar, + flexflow_tensor_t *out, + char *name) { + return add_op( + SCALAR_BINARY_OP(insert_scalar_add_layer), cg, input, scalar, out, name); +} + +flexflow_error_t flexflow_computation_graph_add_op_scalar_sub( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + float scalar, + flexflow_tensor_t *out, + char *name) { + return add_op( + SCALAR_BINARY_OP(insert_scalar_sub_layer), cg, input, scalar, out, name); +} + +flexflow_error_t flexflow_computation_graph_add_op_scalar_truediv( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + float scalar, + flexflow_tensor_t *out, + char *name) { + return add_op(SCALAR_BINARY_OP(insert_scalar_truediv_layer), + cg, + input, + scalar, + out, + name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_sin(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_sin_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_cos(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_cos_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_relu(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_relu_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_identity(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_identity_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_gelu(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_gelu_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_sigmoid(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_sigmoid_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_tanh(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_tanh_layer), cg, input, out, name); +} + +flexflow_error_t + flexflow_computation_graph_add_op_elu(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return add_op(BINARY_OP(insert_tanh_layer), cg, input, out, name); +} + +flexflow_error_t flexflow_computation_graph_add_op_conv2d( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int outChannels, + int kernelH, + int kernelW, + int strideH, + int strideW, + int paddingH, + int paddingW, + flexflow_activation_t activation, + int groups, + bool use_bias, + flexflow_initializer_t kernel_initializer, + flexflow_initializer_t bias_initializer, + flexflow_regularizer_attrs_t kernel_regularizer, + char *name) { + return handle_errors(out, [&] { + return conv2d(deref_opaque(cg), + c_deref_opaque(input), + outChannels, + kernelH, + kernelW, + strideH, + strideW, + paddingH, + paddingW, + to_internal(activation), + groups, + use_bias, + c_deref_opaque(kernel_initializer), + c_deref_opaque(bias_initializer), + c_deref_opaque(kernel_regularizer), + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_dropout(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + float rate, + unsigned long long seed, + char *name) { + return handle_errors(out, [&] { + return dropout(deref_opaque(cg), + c_deref_opaque(input), + rate, + seed, + maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_embedding( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int num_entries, + int out_dim, + flexflow_aggregate_op_t aggr_op, + flexflow_datatype_t output_type, + flexflow_initializer_t initializer, + char *name) { + return handle_errors(out, [&] { + return embedding(deref_opaque(cg), + c_deref_opaque(input), + num_entries, + out_dim, + to_internal(aggr_op), + to_internal(output_type), + c_deref_opaque(initializer), + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_gather(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t index, + int dim, + flexflow_tensor_list_t *out, + char *name) { + return handle_errors(out, [&] { + return gather(deref_opaque(cg), + c_deref_opaque(input), + c_deref_opaque(index), + ff_dim_t(dim), + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_group_by(flexflow_computation_graph_t cg, + flexflow_tensor_t data, + flexflow_tensor_t assign, + int n, + float alpha, + flexflow_tensor_list_t *out, + char *name) { + return handle_errors(out, [&] { + return group_by(deref_opaque(cg), + c_deref_opaque(data), + c_deref_opaque(assign), + n, + alpha, + maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_cache( + flexflow_computation_graph_t, + flexflow_tensor_t, + flexflow_tensor_t *out, + int num_batches, + float (*score_func)(float *, void *, void *, int), + flexflow_tensor_t assign, + flexflow_tensor_t *outs, + int n, + float alpha, + char *name) { + NOT_IMPLEMENTED(); // TODO @lockshaw +} + +flexflow_error_t flexflow_computation_graph_add_op_aggregate( + flexflow_computation_graph_t cg, + flexflow_tensor_t gate_preds, + flexflow_tensor_t gate_assign, + flexflow_tensor_t true_gate_assign, + flexflow_tensor_t full_gate_gradients, + flexflow_tensor_t *exp_preds, + int n, + float lambda_bal, + flexflow_tensor_t *out, + char *name) { + return handle_errors(out, [&] { + return aggregate(deref_opaque(cg), + c_deref_opaque(gate_preds), + c_deref_opaque(gate_assign), + c_deref_opaque(true_gate_assign), + c_deref_opaque(full_gate_gradients), + c_deref_opaque_list(exp_preds, n), + lambda_bal, + maybe_string(name)); + }); +} + +flexflow_error_t aggregate_spec(flexflow_computation_graph_t cg, + flexflow_tensor_t *inputs, + flexflow_tensor_t *out, + int n, + float lambda_bal, + char *name) { + return handle_errors(out, [&] { + return aggregate_spec(deref_opaque(cg), + c_deref_opaque_list(inputs, n), + lambda_bal, + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_pool2d(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int kernelH, + int kernelW, + int strideH, + int strideW, + int paddingH, + int paddingW, + flexflow_pool_op_t pool_op, + flexflow_activation_t activation, + char *name) { + return handle_errors(out, [&] { + return pool2d(deref_opaque(cg), + c_deref_opaque(input), + kernelH, + kernelW, + strideH, + strideW, + paddingH, + paddingW, + to_internal(pool_op), + to_internal(activation), + maybe_string(name)); + }); +} + +flexflow_error_t flexcflow_computation_graph_add_op_layer_norm( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int *axes, + int num_axes, + bool elementwise_affine, + float eps, + char *name) { + return handle_errors(out, [&] { + return layer_norm(deref_opaque(cg), + c_deref_opaque(input), + make_vector(axes, num_axes), + elementwise_affine, + eps, + maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_batch_norm( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + bool relu, + char *name) { + + return handle_errors(out, [&] { + return batch_norm( + deref_opaque(cg), c_deref_opaque(input), relu, maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_batch_matmul( + flexflow_computation_graph_t cg, + flexflow_tensor_t lhs, + flexflow_tensor_t rhs, + flexflow_tensor_t *out, + int a_seq_length_dim, + int b_seq_length_dim, + char *name) { + return handle_errors(out, [&] { + return batch_matmul(deref_opaque(cg), + c_deref_opaque(lhs), + c_deref_opaque(rhs), + a_seq_length_dim, + b_seq_length_dim, + maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_dense( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int out_dim, + flexflow_activation_t activation, + bool use_bias, + flexflow_datatype_t compute_type, + flexflow_initializer_t kernel_initializer, + flexflow_initializer_t bias_initializer, + char *name) { + return handle_errors(out, [&] { + return dense(deref_opaque(cg), + c_deref_opaque(input), + out_dim, + to_internal(activation), + use_bias, + to_internal(compute_type), + deref_opaque(kernel_initializer), + deref_opaque(bias_initializer), + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_cast(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + flexflow_datatype_t out_type, + char *name) { + return handle_errors(out, [&] { + return cast(deref_opaque(cg), + c_deref_opaque(input), + to_internal(out_type), + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_concat(flexflow_computation_graph_t cg, + flexflow_tensor_t *inputs, + flexflow_tensor_t *out, + int num_inputs, + int axis, + char *name) { + return handle_errors(out, [&] { + return concat(deref_opaque(cg), + c_deref_opaque_list(inputs, num_inputs), + axis, + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_mean(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int *dims, + int num_dims, + bool keepdims, + char *name) { + return handle_errors(out, [&] { + return mean(deref_opaque(cg), + c_deref_opaque(input), + make_vector(dims, num_dims), + keepdims, + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_moe(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int num_exp, + int num_select, + int expert_hidden_size, + float alpha, + float lambda, + char *name) { + return handle_errors(out, [&] { + return moe(deref_opaque(cg), + c_deref_opaque(input), + num_exp, + num_select, + expert_hidden_size, + alpha, + lambda, + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_split(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_list_t *out, + int *splits, + int num_splits, + int axis, + char *name) { + return handle_errors(out, [&] { + return split(deref_opaque(cg), + c_deref_opaque(input), + make_vector(splits, num_splits), + axis, + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_flat(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + char *name) { + return handle_errors(out, [&] { + return flat(deref_opaque(cg), c_deref_opaque(input), maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_softmax(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int dim, + char *name) { + return handle_errors(out, [&] { + return softmax( + deref_opaque(cg), c_deref_opaque(input), dim, maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_transpose(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int *permutation, + int num_permutation_values, + char *name) { + return handle_errors(out, [&] { + return transpose(deref_opaque(cg), + c_deref_opaque(input), + make_vector(permutation, num_permutation_values), + maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_reduce_sum( + flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int *axes, + int num_axes, + bool keepdims, + char *name) { + return handle_errors(out, [&] { + return reduce_sum(deref_opaque(cg), + c_deref_opaque(input), + make_vector(axes, num_axes), + keepdims, + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_reshape(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int *shape, + int num_shape_entries, + char *name) { + return handle_errors(out, [&] { + return reshape(deref_opaque(cg), + c_deref_opaque(input), + make_vector(shape, num_shape_entries), + maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_reverse(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_t *out, + int axis, + char *name) { + return handle_errors(out, [&] { + return reverse( + deref_opaque(cg), c_deref_opaque(input), axis, maybe_string(name)); + }); +} + +flexflow_error_t + flexflow_computation_graph_add_op_topk(flexflow_computation_graph_t cg, + flexflow_tensor_t input, + flexflow_tensor_list_t *out, + int k, + bool sorted, + char *name) { + return handle_errors(out, [&] { + return top_k( + deref_opaque(cg), c_deref_opaque(input), k, sorted, maybe_string(name)); + }); +} + +flexflow_error_t flexflow_computation_graph_add_op_multihead_attention( + flexflow_computation_graph_t cg, + flexflow_tensor_t query, + flexflow_tensor_t key, + flexflow_tensor_t value, + flexflow_tensor_t *out, + int embed_dim, + int num_heads, + int kdim, + int vdim, + float dropout, + bool bias, + bool add_bias_kv, + bool add_zero_attn, + flexflow_initializer_t initializer, + char *name) { + return handle_errors(out, [&] { + return multihead_attention(deref_opaque(cg), + c_deref_opaque(query), + c_deref_opaque(key), + c_deref_opaque(value), + embed_dim, + num_heads, + kdim, + vdim, + dropout, + bias, + add_bias_kv, + add_zero_attn, + c_deref_opaque(initializer), + maybe_string(name)); + }); +} + +/* flexflow_error_t flexflow_computation_graph_add_op_cache() */ diff --git a/lib/pcg/ffi/src/pcg.h b/lib/pcg/ffi/src/pcg.h new file mode 100644 index 0000000000..fa9e00a1d0 --- /dev/null +++ b/lib/pcg/ffi/src/pcg.h @@ -0,0 +1,71 @@ +#ifndef _FLEXFLOW_LIB_PCG_FFI_SRC_PCG_H +#define _FLEXFLOW_LIB_PCG_FFI_SRC_PCG_H + +#include "flexflow/pcg.h" +#include "internal/opaque.h" +#include "pcg/computation_graph.h" +#include "pcg/initializer.h" +#include "pcg/layer.h" +#include "pcg/layer_guid_t.h" +#include "pcg/machine_specification.h" +#include "pcg/model_compilation.h" +#include "pcg/operator.h" +#include "pcg/operator_guid_t.h" +#include "pcg/parallel_computation_graph.h" +#include "pcg/parallel_tensor.h" +#include "pcg/parallel_tensor_guid_t.h" +#include "pcg/tensor_guid_t.h" +#include "utils/type_traits.h" + +using namespace FlexFlow; + +namespace FlexFlow { + +struct internal_flexflow_layer_t { + layer_guid_t guid; + Layer layer; +}; +FF_VISITABLE_STRUCT(internal_flexflow_layer_t, guid, layer); + +struct internal_flexflow_tensor_t { + tensor_guid_t guid; + Tensor tensor; +}; +FF_VISITABLE_STRUCT(internal_flexflow_tensor_t, guid, tensor); + +struct internal_flexflow_operator_t { + operator_guid_t guid; + Operator op; +}; +FF_VISITABLE_STRUCT(internal_flexflow_operator_t, guid, op); + +struct internal_flexflow_parallel_tensor_t { + parallel_tensor_guid_t guid; + ParallelTensor parallel_tensor; +}; +FF_VISITABLE_STRUCT(internal_flexflow_parallel_tensor_t, guid, parallel_tensor); + +} // namespace FlexFlow + +REGISTER_OPAQUE(flexflow_computation_graph_t, ComputationGraph); +REGISTER_OPAQUE(flexflow_parallel_computation_graph_t, + ParallelComputationGraph); +REGISTER_OPAQUE(flexflow_operator_t, internal_flexflow_operator_t); +REGISTER_OPAQUE(flexflow_parallel_tensor_t, + internal_flexflow_parallel_tensor_t); +REGISTER_OPAQUE(flexflow_layer_t, internal_flexflow_layer_t); +REGISTER_OPAQUE(flexflow_tensor_t, internal_flexflow_tensor_t); +REGISTER_OPAQUE(flexflow_machine_view_t, MachineView); +REGISTER_OPAQUE(flexflow_initializer_t, optional); +REGISTER_OPAQUE(flexflow_machine_specification_t, MachineSpecification); +REGISTER_OPAQUE(flexflow_model_compilation_input_t, ModelCompilationInput); +REGISTER_OPAQUE(flexflow_model_compilation_result_t, ModelCompilationResult); +REGISTER_OPAQUE(flexflow_tensor_list_t, + std::vector); + +struct internal_flexflow_pcg_error_t { + flexflow_pcg_error_code_t err_code; +}; +REGISTER_OPAQUE(flexflow_pcg_error_t, internal_flexflow_pcg_error_t); + +#endif diff --git a/lib/pcg/include/pcg/computation_graph.h b/lib/pcg/include/pcg/computation_graph.h index 11dad70356..a329198784 100644 --- a/lib/pcg/include/pcg/computation_graph.h +++ b/lib/pcg/include/pcg/computation_graph.h @@ -3,10 +3,13 @@ #include "layer.h" #include "operator_guid_t.h" +#include "pcg/layer.h" +#include "pcg/layer_guid_t.h" +#include "pcg/tensor_guid_t.h" #include "tensor.h" #include "utils/graph.h" #include "utils/strong_typedef.h" -#include "visit_struct/visit_struct.hpp" +#include "utils/type_traits.h" namespace FlexFlow { @@ -14,12 +17,37 @@ struct ComputationGraph : public strong_typedef> { using strong_typedef::strong_typedef; + + /* ComputationGraph(); */ + + operator OutputLabelledMultiDiGraphView() const; + + NodePort port_for_input(size_t); + NodePort port_for_weight(size_t); + NodePort port_for_output(size_t); + + size_t input_for_port(NodePort) const; + size_t weight_for_port(NodePort) const; + size_t output_for_port(NodePort) const; + size_t input_for_port(NodePort); + size_t weight_for_port(NodePort); + size_t output_for_port(NodePort); + + MultiDiInput get_input_slot(layer_guid_t, size_t); + MultiDiInput get_weight_slot(layer_guid_t, size_t); + MultiDiOutput get_output_slot(layer_guid_t, size_t); + + Tensor at(tensor_guid_t) const; + Layer at(layer_guid_t) const; }; -} // namespace FlexFlow +optional get_layer_with_name(ComputationGraph const &, + std::string const &); +optional get_tensor_with_name(ComputationGraph const &, + std::string const &); -namespace FlexFlow { -static_assert(is_well_behaved_value_type_no_hash::value, ""); -} +CHECK_WELL_BEHAVED_VALUE_TYPE_NO_EQ(ComputationGraph); + +} // namespace FlexFlow #endif diff --git a/lib/pcg/include/pcg/computation_graph_builder.h b/lib/pcg/include/pcg/computation_graph_builder.h index 7f01439712..c7c8cee14c 100644 --- a/lib/pcg/include/pcg/computation_graph_builder.h +++ b/lib/pcg/include/pcg/computation_graph_builder.h @@ -5,291 +5,317 @@ namespace FlexFlow { -struct ComputationGraphBuilder - : public use_visitable_cmp { -public: - ComputationGraphBuilder(); - - // C++ APIs for constructing models - // Add an exp layer - Tensor exp(Tensor const &, optional const &name = nullopt); - // Add an add layer - Tensor add(Tensor const &x, - Tensor const &y, - optional const &name = nullopt); - // Add a subtract layer - Tensor subtract(Tensor const &x, - Tensor const &y, - optional const &name = nullopt); - // Add a multiply layer - Tensor multiply(Tensor const &x, - Tensor const &y, - optional const &name = nullopt); - // Add a divide layer - Tensor divide(Tensor const &x, - Tensor const &y, - optional const &name = nullopt); - // Add a max layer - Tensor max(Tensor const &x, - Tensor const &y, - optional const &name = nullopt); - // Add a min layer - Tensor min(Tensor const &x, - Tensor const &y, - optional const &name = nullopt); - // Add a rsqrt layer - Tensor rsqrt(Tensor const &x, optional const &name = nullopt); - // Add a pow layer - Tensor pow(Tensor const &x, - float exponent, - optional const &name = nullopt); - // Add a scalar multiply layer - Tensor scalar_multiply(Tensor const &x, - float scalar, - optional const &name = nullopt); - Tensor scalar_add(Tensor const &x, - float scalar, - optional const &name = nullopt); - Tensor scalar_sub(Tensor const &lhs, - float rhs, - optional const &name = nullopt); - Tensor scalar_truediv(Tensor const &numerator, - float denominator, - optional const &name = nullopt); - // Add a sin layer - Tensor sin(Tensor const &x, optional const &name = nullopt); - // Add a cos layer - Tensor cos(Tensor const &x, optional const &name = nullopt); - // Add an activation layer - Tensor relu(Tensor const &x, optional const &name = nullopt); - Tensor identity(Tensor const &x, optional const &name = nullopt); - Tensor gelu(Tensor const &x, optional const &name = nullopt); - Tensor sigmoid(Tensor const &x, optional const &name = nullopt); - Tensor tanh(Tensor const &x, optional const &name = nullopt); - Tensor elu(Tensor const &x, optional const &name = nullopt); - // Add a 2D convolutional layer - Tensor conv2d(Tensor const &input, - int outChannels, - int kernelH, - int kernelW, - int strideH, - int strideW, - int paddingH, - int paddingW, - optional const &activation = nullopt, - int groups = 1, - bool use_bias = true, - optional kernel_initializer = nullopt, - optional bias_initializer = nullopt, - optional kernel_regularizer = nullopt, - optional const &name = nullopt); - // Add a dropout layer - Tensor dropout(Tensor const &input, - float rate, - unsigned long long seed = 0, - optional const &name = nullopt); - // Add an embedding layer - Tensor embedding(Tensor const &input, - int num_entries, - int outDim, - AggregateOp aggr, - DataType dtype = DataType::FLOAT, - optional kernel_initializer = nullopt, - optional const &name = nullopt); - // Add a gather layer - std::vector gather(Tensor const &input, - Tensor const &index, - ff_dim_t dim, - optional const &name = nullopt); - // Add a group_by layer - void group_by(Tensor const &data, - Tensor const &assign, - Tensor *outputs, - int n, - float alpha, - optional const &name = nullopt); - // Add a cache layer - Tensor cache(Tensor const &input, - int num_batches, - std::function - score_f = {}, - optional const &name = nullopt); - // Add aggregate layer - Tensor aggregate(Tensor const &gate_preds, - Tensor const &gate_assign, - Tensor const &true_gate_assign, - Tensor const &full_gate_gradients, - std::vector const &exp_preds, - int n, - float lambda_bal, - optional const &maybe_name); - // Add aggregate_spec layer - Tensor aggregate_spec(std::vector const &inputs, - int n, - float lambda_bal, +tensor_guid_t insert_exp_layer(ComputationGraph &, + tensor_guid_t const &, + optional const &name = nullopt); + +tensor_guid_t insert_add_layer(ComputationGraph &, + tensor_guid_t const &x, + tensor_guid_t const &y, + optional const &name = nullopt); + +tensor_guid_t + insert_subtract_layer(ComputationGraph &, + tensor_guid_t const &x, + tensor_guid_t const &y, + optional const &name = nullopt); + +tensor_guid_t + insert_multiply_layer(ComputationGraph &, + tensor_guid_t const &x, + tensor_guid_t const &y, + optional const &name = nullopt); + +tensor_guid_t insert_divide_layer(ComputationGraph &, + tensor_guid_t const &x, + tensor_guid_t const &y, + optional const &name = nullopt); + +tensor_guid_t insert_max_layer(ComputationGraph &, + tensor_guid_t const &x, + tensor_guid_t const &y, + optional const &name = nullopt); + +tensor_guid_t insert_min_layer(ComputationGraph &, + tensor_guid_t const &x, + tensor_guid_t const &y, + optional const &name = nullopt); + +tensor_guid_t insert_rsqrt_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); + +tensor_guid_t insert_pow_layer(ComputationGraph &, + tensor_guid_t const &x, + float exponent, + optional const &name = nullopt); + +tensor_guid_t + insert_scalar_multiply_layer(ComputationGraph &, + tensor_guid_t const &x, + float scalar, + optional const &name = nullopt); + +tensor_guid_t + insert_scalar_add_layer(ComputationGraph &, + tensor_guid_t const &x, + float scalar, + optional const &name = nullopt); + +tensor_guid_t + insert_scalar_sub_layer(ComputationGraph &, + tensor_guid_t const &lhs, + float rhs, + optional const &name = nullopt); + +tensor_guid_t + insert_scalar_truediv_layer(ComputationGraph &, + tensor_guid_t const &numerator, + float denominator, + optional const &name = nullopt); + +tensor_guid_t insert_sin_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); + +tensor_guid_t insert_cos_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); + +tensor_guid_t insert_relu_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); +tensor_guid_t + insert_identity_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); +tensor_guid_t insert_gelu_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); +tensor_guid_t insert_sigmoid_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); +tensor_guid_t insert_tanh_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); +tensor_guid_t insert_elu_layer(ComputationGraph &, + tensor_guid_t const &x, + optional const &name = nullopt); + +tensor_guid_t insert_conv2d_layer( + ComputationGraph &, + tensor_guid_t const &input, + int outChannels, + int kernelH, + int kernelW, + int strideH, + int strideW, + int paddingH, + int paddingW, + optional const &activation = nullopt, + int groups = 1, + bool use_bias = true, + optional const &kernel_initializer = nullopt, + optional const &bias_initializer = nullopt, + optional const &kernel_regularizer = nullopt, + optional const &name = nullopt); + +tensor_guid_t insert_dropout_layer(ComputationGraph &, + tensor_guid_t const &input, + float rate, + unsigned long long seed = 0, + optional const &name = nullopt); + +tensor_guid_t insert_embedding_layer( + ComputationGraph &, + tensor_guid_t const &input, + int num_entries, + int outDim, + AggregateOp aggr, + DataType dtype = DataType::FLOAT, + optional const &kernel_initializer = nullopt, + optional const &name = nullopt); + +std::vector + insert_gather_layer(ComputationGraph &, + tensor_guid_t const &input, + tensor_guid_t const &index, + ff_dim_t dim, optional const &name = nullopt); - // Add a 2D pooling layer - Tensor pool2d(Tensor const &input, - int kernelH, - int kernelW, - int strideH, - int strideW, - int paddingH, - int paddingW, - PoolOp type = PoolOp::MAX, - optional const &activation = nullopt, - optional const &name = nullopt); - Tensor layer_norm(Tensor const &input, - std::vector const &axes, - bool elementwise_affine, - float eps, - optional const &name = nullopt); - Tensor batch_norm(Tensor const &input, - bool relu = true, - optional const &name = nullopt); - Tensor batch_matmul(Tensor const &A, - Tensor const &B, - int a_seq_length_dim = -1, - int b_seq_length_dim = -1, - optional const &name = nullopt); - Tensor dense(Tensor const &input, - int outDim, - optional activation = nullopt, - bool use_bias = true, - DataType data_type = DataType::FLOAT, - optional kernel_initializer = nullopt, - optional bias_initializer = nullopt, - optional const &name = nullopt); - // Add a cast layer - Tensor cast(Tensor const &input, - DataType dtype, - optional const &name = nullopt); - // Add a concat layer - Tensor concat(int n, - std::vector const &tensors, - int axis, - optional const &name = nullopt); - // Add a mean layer - Tensor mean(Tensor const &input, - std::vector const &dims, - bool keepdims, - char const *name); - // Add a moe layer (wrapping topk, group_by and aggregate operators) - Tensor moe(Tensor const &input, - int num_exp, - int num_select, - int expert_hidden_size, - float alpha, - float lambda); - // Add a split layer - void split(Tensor const &input, - Tensor *outputs, - std::vector const &split, - int axis, - optional const &name = nullopt); - // Add a flat layer - Tensor flat(Tensor const &input, optional const &name = nullopt); - // Add a softmax layer - Tensor softmax(Tensor const &input, - int dim = -1, - optional const &name = nullopt); - // Create input tensors and constants - Tensor transpose(Tensor const &input, - std::vector const &perm, - optional const &name = nullopt); - Tensor reduce_sum(Tensor const &input, - std::vector const &axes, - bool keepdims = false, - optional const &name = nullopt); - Tensor reshape(Tensor const &input, - std::vector const &shape, - optional const &name = nullopt); - Tensor reverse(Tensor const &input, - int axis, - optional const &name = nullopt); - void top_k(Tensor const &input, - Tensor *outputs, - int k, - bool sorted, - optional const &name = nullopt); - Tensor - multihead_attention(Tensor const &query, - Tensor const &key, - Tensor const &value, - int embed_dim, - int num_heads, - int kdim = 0, - int vdim = 0, - float dropout = 0.0f, - bool bias = true, - bool add_bias_kv = false, - bool add_zero_attn = false, - optional initializer = nullopt, + +std::vector + insert_group_by_layer(ComputationGraph &, + tensor_guid_t const &data, + tensor_guid_t const &assign, + int n, + float alpha, optional const &name = nullopt); - Tensor create_tensor(TensorShape const &, bool create_grad = true); - Parameter create_weight(TensorShape const &, - bool create_grad = true, - optional initializer = nullopt, - optional sync_type = nullopt); - - std::vector get_outputs(Layer const &) const; - Tensor get_output(Layer const &, int idx) const; - - Tensor at(MultiDiEdge const &) const; - Layer at(Node const &) const; - -private: - Tensor broadcast(Tensor const &, TensorShape const &); - - void add_layer(Layer const &layer, - std::vector const &inputs, - std::vector const &weights, - std::vector const &outputs); - Tensor - add_layer(Layer const &layer, - std::vector const &inputs, - std::vector>> const - &weight_shapes, - TensorShape const &output_shape); - std::vector - add_layer(Layer const &layer, - std::vector const &inputs, - std::vector>> const - &weight_shapes, - std::vector const &output_shapes); - - Tensor as_type(Tensor const &, DataType, std::string const &); - - TensorShape get_broadcast_target_shape(std::vector const &); - - Tensor element_binary(OperatorType, - Tensor const &lhs, - Tensor const &rhs, + +tensor_guid_t insert_cache_layer( + ComputationGraph &, + tensor_guid_t const &input, + int num_batches, + std::function score_f = {}, + optional const &name = nullopt); + +tensor_guid_t + insert_aggregate_layer(ComputationGraph &, + tensor_guid_t const &gate_preds, + tensor_guid_t const &gate_assign, + tensor_guid_t const &true_gate_assign, + tensor_guid_t const &full_gate_gradients, + std::vector const &exp_preds, + float lambda_bal, + optional const &name = nullopt); + +tensor_guid_t + insert_aggregate_spec_layer(ComputationGraph &, + std::vector const &inputs, + float lambda_bal, + optional const &name = nullopt); + +tensor_guid_t + insert_pool2d_layer(ComputationGraph &, + tensor_guid_t const &input, + int kernelH, + int kernelW, + int strideH, + int strideW, + int paddingH, + int paddingW, + PoolOp type = PoolOp::MAX, + optional const &activation = nullopt, optional const &name = nullopt); - Tensor element_unary(OperatorType, - Tensor const &input, - optional const &name = nullopt); - Tensor element_scalar_unary(OperatorType, - Tensor const &input, - float scalar, +tensor_guid_t + insert_layer_norm_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &axes, + bool elementwise_affine, + float eps, + optional const &name = nullopt); + +tensor_guid_t + insert_batch_norm_layer(ComputationGraph &, + tensor_guid_t const &input, + bool relu = true, + optional const &name = nullopt); +tensor_guid_t + insert_batch_matmul_layer(ComputationGraph &, + tensor_guid_t const &A, + tensor_guid_t const &B, + int a_seq_length_dim = -1, + int b_seq_length_dim = -1, optional const &name = nullopt); - Tensor - element_unary(variant const &, - Tensor const &input, - optional const &name = nullopt); +tensor_guid_t insert_dense_layer( + ComputationGraph &, + tensor_guid_t const &input, + int outDim, + optional activation = nullopt, + bool use_bias = true, + DataType data_type = DataType::FLOAT, + optional const &kernel_initializer = nullopt, + optional const &bias_initializer = nullopt, + optional const &name = nullopt); -public: - ComputationGraph computation_graph; -}; +tensor_guid_t insert_cast_layer(ComputationGraph &, + tensor_guid_t const &input, + DataType dtype, + optional const &name = nullopt); -} // namespace FlexFlow +tensor_guid_t insert_concat_layer(ComputationGraph &, + std::vector const &tensors, + int axis, + optional const &name = nullopt); -VISITABLE_STRUCT(::FlexFlow::ComputationGraphBuilder, computation_graph); +tensor_guid_t insert_mean_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &dims, + bool keepdims, + optional const &name = nullopt); -namespace FlexFlow { -static_assert( - is_well_behaved_value_type_no_hash::value, ""); -} +tensor_guid_t + insert_moe_layer_composite(ComputationGraph &, + tensor_guid_t const &input, + int num_exp, + int num_select, + int expert_hidden_size, + float alpha, + float lambda, + optional const &name = nullopt); + +std::vector + insert_split_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &split, + int axis, + optional const &name = nullopt); + +tensor_guid_t insert_flat_layer(ComputationGraph &, + tensor_guid_t const &input, + optional const &name = nullopt); + +tensor_guid_t insert_softmax_layer(ComputationGraph &, + tensor_guid_t const &input, + int dim = -1, + optional const &name = nullopt); + +tensor_guid_t + insert_transpose_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &perm, + optional const &name = nullopt); + +tensor_guid_t + insert_reduce_sum_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &axes, + bool keepdims = false, + optional const &name = nullopt); + +tensor_guid_t insert_reshape_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &shape, + optional const &name = nullopt); + +tensor_guid_t insert_reverse_layer(ComputationGraph &, + tensor_guid_t const &input, + int axis, + optional const &name = nullopt); + +std::vector + insert_top_k_layer(ComputationGraph &, + tensor_guid_t const &input, + int k, + bool sorted, + optional const &name = nullopt); + +tensor_guid_t insert_multihead_attention_layer( + ComputationGraph &, + tensor_guid_t const &query, + tensor_guid_t const &key, + tensor_guid_t const &value, + int embed_dim, + int num_heads, + int kdim = 0, + int vdim = 0, + float dropout = 0.0f, + bool bias = true, + bool add_bias_kv = false, + bool add_zero_attn = false, + optional const &initializer = nullopt, + optional const &name = nullopt); + +tensor_guid_t insert_new_activation_tensor(ComputationGraph &, + TensorShape const &); +weight_guid_t insert_new_weight_tensor( + ComputationGraph &, + TensorShape const &, + optional const &initializer = nullopt); + +tensor_guid_t insert_broadcast_layer(tensor_guid_t const &, + TensorShape const &); + +} // namespace FlexFlow #endif diff --git a/lib/pcg/include/pcg/file_format/v1/graphs.h b/lib/pcg/include/pcg/file_format/v1/graphs.h index 71a8adb344..e587e11cdb 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs.h @@ -63,6 +63,7 @@ FF_VISITABLE_STRUCT( V1ComputationGraph, node_labels, outputs, output_labels, graph); CHECK_IS_JSONABLE(V1ComputationGraph); V1ComputationGraph to_v1(ComputationGraph const &); +ComputationGraph from_v1(V1ComputationGraph const &); using V1ParallelComputationGraph = V1JsonableGraph; diff --git a/lib/pcg/include/pcg/layer.h b/lib/pcg/include/pcg/layer.h index 1017036e69..e1d540d793 100644 --- a/lib/pcg/include/pcg/layer.h +++ b/lib/pcg/include/pcg/layer.h @@ -7,24 +7,13 @@ namespace FlexFlow { -struct Layer : public use_visitable_cmp { -public: - Layer() = delete; - Layer(CompGraphOperatorAttrs const &attrs, optional const &name); - -public: - optional> name; +struct Layer { CompGraphOperatorAttrs attrs; + req>> name; }; - -} // namespace FlexFlow - -VISITABLE_STRUCT(::FlexFlow::Layer, attrs, name); -MAKE_VISIT_HASHABLE(::FlexFlow::Layer); - -namespace FlexFlow { -static_assert(is_well_behaved_value_type::value, ""); +FF_VISITABLE_STRUCT(Layer, attrs, name); static_assert(is_fmtable::value, "Layer must be fmtable"); + } // namespace FlexFlow #endif diff --git a/lib/pcg/include/pcg/layer_guid_t.h b/lib/pcg/include/pcg/layer_guid_t.h new file mode 100644 index 0000000000..b07a38ca13 --- /dev/null +++ b/lib/pcg/include/pcg/layer_guid_t.h @@ -0,0 +1,17 @@ +#ifndef _FLEXFLOW_PCG_INCLUDE_PCG_LAYER_GUID_T_H +#define _FLEXFLOW_PCG_INCLUDE_PCG_LAYER_GUID_T_H + +#include "utils/graph.h" +#include "utils/strong_typedef.h" + +namespace FlexFlow { + +struct layer_guid_t : public strong_typedef { + using strong_typedef::strong_typedef; +}; +FF_TYPEDEF_HASHABLE(layer_guid_t); +FF_TYPEDEF_PRINTABLE(layer_guid_t, "layer_guid"); + +} // namespace FlexFlow + +#endif diff --git a/lib/pcg/include/pcg/parallel_tensor.h b/lib/pcg/include/pcg/parallel_tensor.h index eadc83d9fd..be902d6500 100644 --- a/lib/pcg/include/pcg/parallel_tensor.h +++ b/lib/pcg/include/pcg/parallel_tensor.h @@ -34,41 +34,18 @@ namespace FlexFlow { * @details Parallel tensor is the fundamental component to support the * representation and exploration of parallelization strategies. */ -struct ParallelTensor : public use_visitable_cmp { - ParallelTensor() = delete; - - ParallelTensor(ParallelTensorShape const &, - CreateGrad create_gradients, - optional sync_type = nullopt, - optional initializer = nullopt); - ParallelTensor(ParallelTensorDims const &, - DataType, - CreateGrad create_gradients, - optional sync_type = nullopt, - optional initializer = nullopt); - -public: +struct ParallelTensor { ParallelTensorDims dims; DataType data_type; - optional sync_type = nullopt; - optional initializer = nullopt; - CreateGrad create_gradients; + optional sync_type; + optional initializer; + req create_gradients; }; +FF_VISITABLE_STRUCT( + ParallelTensor, dims, data_type, sync_type, initializer, create_gradients); using ParallelParameter = ParallelTensor; } // namespace FlexFlow -VISITABLE_STRUCT(::FlexFlow::ParallelTensor, - dims, - data_type, - sync_type, - initializer, - create_gradients); -MAKE_VISIT_HASHABLE(::FlexFlow::ParallelTensor); - -namespace FlexFlow { -static_assert(is_well_behaved_value_type::value, ""); -} - #endif diff --git a/lib/pcg/include/pcg/tensor.h b/lib/pcg/include/pcg/tensor.h index cb79be245a..ad1193db87 100644 --- a/lib/pcg/include/pcg/tensor.h +++ b/lib/pcg/include/pcg/tensor.h @@ -9,27 +9,21 @@ namespace FlexFlow { struct Tensor { - /* Tensor() = delete; */ - /* Tensor(TensorShape const &, */ - /* CreateGrad create_gradients, */ - /* optional initializer = nullopt, */ - /* optional sync_type = nullopt); */ - size_t get_volume() const; TensorShape get_shape() const; int num_dims() const; + DataType get_data_type() const; + operator TensorShape() const; public: - TensorDims dims; - DataType data_type; - req> initializer; - req create_gradients; + TensorShape shape; + CreateGrad create_gradients; + optional initializer; req> sync_type; }; -FF_VISITABLE_STRUCT( - Tensor, dims, data_type, initializer, create_gradients, sync_type); +FF_VISITABLE_STRUCT(Tensor, shape, create_gradients, initializer, sync_type); using Parameter = Tensor; diff --git a/lib/pcg/include/pcg/tensor_guid_t.h b/lib/pcg/include/pcg/tensor_guid_t.h index 3e4e840a5f..990d60db3e 100644 --- a/lib/pcg/include/pcg/tensor_guid_t.h +++ b/lib/pcg/include/pcg/tensor_guid_t.h @@ -2,16 +2,22 @@ #define _FLEXFLOW_PCG_INCLUDE_PCG_TENSOR_GUID_T_H #include "utils/graph.h" +#include "utils/strong_typedef.h" namespace FlexFlow { -struct tensor_guid_t : strong_typedef { +struct weight_guid_t : public strong_typedef { using strong_typedef::strong_typedef; }; +FF_TYPEDEF_HASHABLE(weight_guid_t); +FF_TYPEDEF_PRINTABLE(weight_guid_t, "weight_guid"); -} // namespace FlexFlow +struct tensor_guid_t : public strong_typedef { + using strong_typedef::strong_typedef; +}; +FF_TYPEDEF_HASHABLE(tensor_guid_t); +FF_TYPEDEF_PRINTABLE(tensor_guid_t, "tensor_guid"); -MAKE_TYPEDEF_PRINTABLE(::FlexFlow::tensor_guid_t, "tensor_guid"); -MAKE_TYPEDEF_HASHABLE(::FlexFlow::tensor_guid_t); +} // namespace FlexFlow #endif diff --git a/lib/pcg/src/computation_graph_builder.cc b/lib/pcg/src/computation_graph_builder.cc index a17cd18f7c..04c713141a 100644 --- a/lib/pcg/src/computation_graph_builder.cc +++ b/lib/pcg/src/computation_graph_builder.cc @@ -1,20 +1,172 @@ #include "pcg/computation_graph_builder.h" +#include "op-attrs/datatype.h" #include "op-attrs/get_op_type.h" #include "op-attrs/get_output_shapes.h" +#include "op-attrs/operator_attrs.h" +#include "op-attrs/tensor_shape.h" +#include "pcg/computation_graph.h" +#include "pcg/create_grad.h" +#include "pcg/layer_guid_t.h" +#include "pcg/tensor_guid_t.h" +#include "utils/containers.h" +#include "utils/exception.h" #include "utils/expected.h" #include "utils/fmt.h" +#include "utils/graph/multidiedge.h" +#include "op-attrs/get_output_shapes.h" namespace FlexFlow { -Tensor ComputationGraphBuilder::as_type(Tensor const &x, - DataType data_type, - std::string const &name) { - if (x.data_type < data_type) { - return this->cast(x, data_type, name); - } else if (x.data_type > data_type) { +static TensorShape get_shape(ComputationGraph const &, tensor_guid_t const &); +static TensorShape get_shape(ComputationGraph const &, std::vector const &); + +template +static TensorShape get_output_shape(ComputationGraph const &, Attrs const &, tensor_guid_t const &); + +template +static TensorShape get_output_shape(ComputationGraph const &, Attrs const &, tensor_guid_t const &, tensor_guid_t const &); + +template +static std::vector get_output_shapes(ComputationGraph const &, Attrs const &, tensor_guid_t const &, tensor_guid_t const &); + +static TensorShape get_broadcast_target_shape(std::vector const &); +static TensorShape get_broadcast_target_shape(ComputationGraph const &, std::vector const &); + +static tensor_guid_t broadcast(ComputationGraph &, tensor_guid_t, TensorShape const &); + +static DataType get_data_type(ComputationGraph const &, tensor_guid_t const &); + +static layer_guid_t add_layer(ComputationGraph &cg, + Layer const &layer, + std::vector const &inputs, + std::vector const &weights, + std::vector const &outputs) { + return cg.fmap( + [&](OutputLabelledMultiDiGraph &g) -> layer_guid_t { + auto guid = layer_guid_t{g.add_node(layer)}; + for (int i = 0; i < inputs.size(); i++) { + g.add_edge(inputs[i], cg.get_input_slot(guid, i)); + } + for (int i = 0; i < weights.size(); i++) { + g.add_edge(weights[i], cg.get_weight_slot(guid, i)); + } + for (int i = 0; i < outputs.size(); i++) { + g.add_output(cg.get_output_slot(guid, i), outputs[i]); + } + + return guid; + }); +} + +static layer_guid_t + add_layer(ComputationGraph &cg, + Layer const &layer, + std::vector const &inputs, + std::vector>> const + &weight_shapes, + std::vector const &output_shapes) { + std::vector weights = + transform(weight_shapes, + [&](std::pair> const &kv) { + return insert_new_weight_tensor(cg, kv.first, kv.second); + }); + std::vector outputs = + transform(output_shapes, [](TensorShape const &s) -> Tensor { + return {s, CreateGrad::YES, nullopt, nullopt}; + }); + return add_layer(cg, layer, inputs, weights, outputs); +} + +static std::vector get_output_tensors(ComputationGraph const &cg, + layer_guid_t layer) { + std::unordered_set unsorted_outputs = get_outputs(cg, layer); + std::vector outputs{unsorted_outputs.cbegin(), + unsorted_outputs.cend()}; + return transform(sorted_by(outputs, + [&](MultiDiOutput const &o) { + return cg.output_for_port(o.idx); + }), + [](MultiDiOutput const &o) { return tensor_guid_t{o}; }); +} + +static tensor_guid_t get_only_output_tensor(ComputationGraph const &cg, + layer_guid_t layer) { + std::vector outputs = get_output_tensors(cg, layer); + return get_only(outputs); +} + +static std::vector + insert_layer(ComputationGraph &cg, + Layer const &layer, + std::vector const &inputs, + std::vector const &weights, + std::vector const &outputs) { + return get_output_tensors(cg, add_layer(cg, layer, inputs, weights, outputs)); +} + +static std::vector insert_layer( + ComputationGraph &cg, + Layer const &layer, + std::vector const &inputs, + std::vector>> const &weights, + std::vector const &output_shapes) { + return get_output_tensors( + cg, add_layer(cg, layer, inputs, weights, output_shapes)); +} + +static tensor_guid_t insert_layer(ComputationGraph &cg, + Layer const &layer, + std::vector const &inputs, + std::vector const &weights, + Tensor const &output) { + return get_only_output_tensor( + cg, add_layer(cg, layer, inputs, weights, {output})); +} + +static tensor_guid_t insert_layer( + ComputationGraph &cg, + Layer const &layer, + std::vector const &inputs, + std::vector>> const &weights, + TensorShape const &output_shape) { + return get_only_output_tensor( + cg, add_layer(cg, layer, inputs, weights, {output_shape})); +} + +static tensor_guid_t + element_binary(ComputationGraph &, + OperatorType, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, + optional const &name = nullopt); + +static tensor_guid_t element_unary(ComputationGraph &, + OperatorType, + tensor_guid_t const &input, + optional const &name = nullopt); +static tensor_guid_t + element_scalar_unary(ComputationGraph &, + OperatorType, + tensor_guid_t const &input, + float scalar, + optional const &name = nullopt); +static tensor_guid_t + element_unary(ComputationGraph &, + variant const &, + tensor_guid_t const &input, + optional const &name = nullopt); + +static tensor_guid_t as_type(ComputationGraph &cg, + tensor_guid_t const &x, + DataType data_type, + std::string const &name) { + Tensor tensor = cg.at(x); + if (tensor.get_data_type() < data_type) { + return insert_cast_layer(cg, x, data_type, name); + } else if (tensor.get_data_type() > data_type) { throw mk_runtime_error("Could not convert provided tensor data type {} to " "desired data type {}", - x.data_type, + tensor.get_data_type(), data_type); } return x; @@ -33,195 +185,220 @@ static std::string get_default_name(variant const &attrs) { return get_default_name(widen(attrs)); } -Tensor ComputationGraphBuilder::element_unary( +static tensor_guid_t insert_element_unary_layer( + ComputationGraph &cg, variant const &attrs, - Tensor const &x, + tensor_guid_t const &x, optional const &maybe_name) { std::string name = maybe_name.value_or(get_default_name(attrs)); - Tensor input = this->as_type(x, DataType::FLOAT, name + "input_pre_cast"); + tensor_guid_t input = + as_type(cg, x, DataType::FLOAT, name + "input_pre_cast"); Layer layer = {widen(attrs), name}; - TensorShape output_shape = get_output_shape(attrs, input); + TensorShape output_shape = get_output_shape(attrs, get_shape(cg, input)); - return this->add_layer(layer, {input}, {}, output_shape); + return insert_layer(cg, layer, {input}, {}, output_shape); } -Tensor - ComputationGraphBuilder::element_unary(OperatorType op_type, - Tensor const &input, - optional const &name) { +static tensor_guid_t + insert_element_unary_layer(ComputationGraph &cg, + OperatorType op_type, + tensor_guid_t const &input, + optional const &name) { ElementUnaryAttrs attrs = {op_type}; - return this->element_unary(attrs, input, name); + return insert_element_unary_layer(cg, attrs, input, name); } -Tensor ComputationGraphBuilder::element_scalar_unary( - OperatorType op_type, - Tensor const &input, - float scalar, - optional const &name) { +static tensor_guid_t + insert_element_scalar_unary_layer(ComputationGraph &cg, + OperatorType op_type, + tensor_guid_t const &input, + float scalar, + optional const &name) { ElementScalarUnaryAttrs attrs = {op_type, scalar}; - return this->element_unary(attrs, input, name); + return insert_element_unary_layer(cg, attrs, input, name); } -Tensor ComputationGraphBuilder::element_binary( - OperatorType op_type, - Tensor const &lhs, - Tensor const &rhs, - optional const &maybe_name) { +static tensor_guid_t + insert_element_binary_layer(ComputationGraph &cg, + OperatorType op_type, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, + optional const &maybe_name) { std::string name = maybe_name.value_or(get_default_name(op_type)); - TensorShape compute_shape = this->get_broadcast_target_shape({lhs, rhs}); - DataType compute_type = std::max(lhs.data_type, rhs.data_type); + TensorShape compute_shape = get_broadcast_target_shape(cg, {lhs, rhs}); + DataType compute_type = std::max(get_data_type(cg, lhs), get_data_type(cg, rhs)); - Tensor const lhs_input = this->as_type(this->broadcast(lhs, compute_shape), - compute_type, - name + "_inputl_pre_cast"); - Tensor const rhs_input = this->as_type(this->broadcast(rhs, compute_shape), - compute_type, - name + "_inputr_pre_cast"); + tensor_guid_t const lhs_input = as_type(cg, + broadcast(cg, lhs, compute_shape), compute_type, name + "_inputl_pre_cast"); + tensor_guid_t const rhs_input = as_type(cg, + broadcast(cg, rhs, compute_shape), compute_type, name + "_inputr_pre_cast"); ElementBinaryAttrs attrs = {op_type, compute_type, false, false}; Layer layer = {attrs, name}; - TensorShape output_shape = get_output_shape(attrs, lhs_input, rhs_input); + TensorShape output_shape = get_output_shape(cg, attrs, lhs_input, rhs_input); - return this->add_layer(layer, {lhs_input, rhs_input}, {}, output_shape); + return insert_layer(cg, layer, {lhs_input, rhs_input}, {}, output_shape); } -Tensor ComputationGraphBuilder::exp(Tensor const &input, - optional const &name) { - return this->element_unary(Op::EXP, input, name); -} - -Tensor ComputationGraphBuilder::add(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_ADD, lhs, rhs, name); +tensor_guid_t insert_exp_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::EXP, input, name); } -Tensor ComputationGraphBuilder::subtract(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_SUB, lhs, rhs, name); +tensor_guid_t insert_add_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, + optional const &name) { + return element_binary(cg, Op::EW_ADD, lhs, rhs, name); } -Tensor ComputationGraphBuilder::multiply(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_MUL, lhs, rhs, name); -} - -Tensor ComputationGraphBuilder::divide(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_DIV, lhs, rhs, name); -} - -Tensor ComputationGraphBuilder::max(Tensor const &lhs, - Tensor const &rhs, +tensor_guid_t insert_subtract_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, optional const &name) { - return this->element_binary(Op::EW_MAX, lhs, rhs, name); + return element_binary(cg, Op::EW_SUB, lhs, rhs, name); } -Tensor ComputationGraphBuilder::min(Tensor const &lhs, - Tensor const &rhs, +tensor_guid_t insert_multiply_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, optional const &name) { - return this->element_binary(Op::EW_MIN, lhs, rhs, name); + return element_binary(cg, Op::EW_MUL, lhs, rhs, name); } -Tensor ComputationGraphBuilder::rsqrt(Tensor const &input, - optional const &name) { - return this->element_unary(Op::RSQRT, input, name); +tensor_guid_t insert_divide_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, + optional const &name) { + return element_binary(cg, Op::EW_DIV, lhs, rhs, name); } -Tensor ComputationGraphBuilder::pow(Tensor const &input, - float exponent, - optional const &name) { - return this->element_scalar_unary(Op::POW, input, exponent, name); +tensor_guid_t insert_max_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, + optional const &name) { + return element_binary(cg, Op::EW_MAX, lhs, rhs, name); } -Tensor ComputationGraphBuilder::scalar_multiply( - Tensor const &input, float scalar, optional const &name) { - return this->element_scalar_unary(Op::SCALAR_MULTIPLY, input, scalar, name); +tensor_guid_t insert_min_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + tensor_guid_t const &rhs, + optional const &name) { + return element_binary(cg, Op::EW_MIN, lhs, rhs, name); } -Tensor ComputationGraphBuilder::scalar_add(Tensor const &input, - float scalar, - optional const &name) { - return this->element_scalar_unary(Op::SCALAR_ADD, input, scalar, name); -} - -Tensor ComputationGraphBuilder::scalar_sub(Tensor const &lhs, - float rhs, - optional const &name) { - return this->element_scalar_unary(Op::SCALAR_SUB, lhs, rhs, name); +tensor_guid_t insert_rsqrt_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::RSQRT, input, name); } -Tensor - ComputationGraphBuilder::scalar_truediv(Tensor const &numerator, - float denominator, - optional const &name) { - return this->element_scalar_unary( - Op::SCALAR_TRUE_DIV, numerator, denominator, name); +tensor_guid_t insert_pow_layer(ComputationGraph &cg, + tensor_guid_t const &input, + float exponent, + optional const &name) { + return element_scalar_unary(cg, Op::POW, input, exponent, name); } -Tensor ComputationGraphBuilder::sin(Tensor const &input, - optional const &name) { - return this->element_unary(Op::SIN, input, name); +tensor_guid_t insert_scalar_multiply_layer(ComputationGraph &cg, + tensor_guid_t const &input, + float scalar, + optional const &name) { + return element_scalar_unary(cg, Op::SCALAR_MULTIPLY, input, scalar, name); } -Tensor ComputationGraphBuilder::cos(Tensor const &input, - optional const &name) { - return this->element_unary(Op::COS, input, name); +tensor_guid_t insert_scalar_add_layer(ComputationGraph &cg, + tensor_guid_t const &input, + float scalar, + optional const &name) { + return element_scalar_unary(cg, Op::SCALAR_ADD, input, scalar, name); } -Tensor ComputationGraphBuilder::relu(Tensor const &input, - optional const &name) { - return this->element_unary(Op::RELU, input, name); +tensor_guid_t insert_scalar_sub_layer(ComputationGraph &cg, + tensor_guid_t const &lhs, + float rhs, + optional const &name) { + return element_scalar_unary(cg, Op::SCALAR_SUB, lhs, rhs, name); } -Tensor ComputationGraphBuilder::identity(Tensor const &input, - optional const &name) { - return this->element_unary(Op::IDENTITY, input, name); +tensor_guid_t insert_scalar_truediv_layer(ComputationGraph &cg, + tensor_guid_t const &numerator, + float denominator, + optional const &name) { + return element_scalar_unary( + cg, Op::SCALAR_TRUE_DIV, numerator, denominator, name); } -Tensor ComputationGraphBuilder::gelu(Tensor const &input, - optional const &name) { - return this->element_unary(Op::GELU, input, name); +tensor_guid_t insert_sin_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::SIN, input, name); } -Tensor ComputationGraphBuilder::sigmoid(Tensor const &input, - optional const &name) { - return this->element_unary(Op::SIGMOID, input, name); +tensor_guid_t insert_cos_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::COS, input, name); } -Tensor ComputationGraphBuilder::tanh(Tensor const &input, - optional const &name) { - return this->element_unary(Op::TANH, input, name); +tensor_guid_t insert_relu_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::RELU, input, name); } -Tensor ComputationGraphBuilder::elu(Tensor const &input, +tensor_guid_t insert_identity_layer(ComputationGraph &cg, + tensor_guid_t const &input, optional const &name) { - return this->element_unary(Op::ELU, input, name); -} - -Tensor ComputationGraphBuilder::conv2d( - Tensor const &x, - int outChannels, - int kernelH, - int kernelW, - int strideH, - int strideW, - int paddingH, - int paddingW, - optional const &activation, - int groups, - bool use_bias, - optional kernel_initializer, - optional bias_initializer, - optional kernel_regularizer, - optional const &maybe_name) { + return element_unary(cg, Op::IDENTITY, input, name); +} + +tensor_guid_t insert_gelu_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::GELU, input, name); +} + +tensor_guid_t insert_sigmoid_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::SIGMOID, input, name); +} + +tensor_guid_t insert_tanh_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::TANH, input, name); +} + +tensor_guid_t insert_elu_layer(ComputationGraph &cg, + tensor_guid_t const &input, + optional const &name) { + return element_unary(cg, Op::ELU, input, name); +} + +tensor_guid_t + insert_conv2d_layer(ComputationGraph &cg, + tensor_guid_t const &x, + int outChannels, + int kernelH, + int kernelW, + int strideH, + int strideW, + int paddingH, + int paddingW, + optional const &activation, + int groups, + bool use_bias, + optional kernel_initializer, + optional bias_initializer, + optional kernel_regularizer, + optional const &maybe_name) { Conv2DAttrs attrs = {outChannels, kernelH, kernelW, @@ -234,94 +411,101 @@ Tensor ComputationGraphBuilder::conv2d( use_bias}; std::string name = maybe_name.value_or(get_default_name(attrs)); - Tensor input = this->as_type(x, DataType::FLOAT, name + "input_pre_cast"); + tensor_guid_t input = + as_type(cg, x, DataType::FLOAT, name + "input_pre_cast"); Layer layer = {attrs, name}; - TensorShape output_shape = get_output_shape(attrs, input); + TensorShape output_shape = get_output_shape(cg, attrs, input); std::vector>> weights; - weights.push_back({get_kernel_shape(attrs, input), kernel_initializer}); + weights.push_back({get_kernel_shape(attrs, get_shape(cg, input)), kernel_initializer}); if (use_bias) { - weights.push_back({get_bias_shape(attrs, input), bias_initializer}); + weights.push_back({get_bias_shape(attrs, get_shape(cg, input)), bias_initializer}); } - return this->add_layer(layer, {input}, weights, output_shape); + return insert_layer(cg, layer, {input}, weights, output_shape); } -Tensor - ComputationGraphBuilder::dropout(Tensor const &x, - float rate, - unsigned long long seed, - optional const &maybe_name) { +tensor_guid_t insert_dropout_layer(ComputationGraph &cg, + tensor_guid_t const &x, + float rate, + unsigned long long seed, + optional const &maybe_name) { DropoutAttrs attrs = {rate, seed}; std::string name = maybe_name.value_or(get_default_name(attrs)); Layer layer = {attrs, name}; - Tensor input = this->as_type(x, DataType::FLOAT, name + "input_pre_cast"); + tensor_guid_t input = + as_type(cg, x, DataType::FLOAT, name + "input_pre_cast"); - TensorShape output_shape = get_output_shape(attrs, input); + TensorShape output_shape = get_output_shape(attrs, get_shape(cg, input)); - return this->add_layer(layer, {input}, {}, output_shape); + return insert_layer(cg, layer, {input}, {}, output_shape); } -Tensor ComputationGraphBuilder::embedding( - Tensor const &x, - int num_entries, - int outDim, - AggregateOp aggr, - DataType dtype, - optional kernel_initializer, - optional const &maybe_name) { +tensor_guid_t + insert_embedding_layer(ComputationGraph &cg, + tensor_guid_t const &x, + int num_entries, + int outDim, + AggregateOp aggr, + DataType dtype, + optional kernel_initializer, + optional const &maybe_name) { EmbeddingAttrs attrs = {num_entries, outDim, aggr, dtype}; std::string name = maybe_name.value_or(get_default_name(attrs)); Layer layer = {attrs, name}; - Tensor input = this->as_type(x, DataType::FLOAT, name + "input_pre_cast"); + tensor_guid_t input = as_type(cg, x, DataType::FLOAT, name + "input_pre_cast"); - TensorShape output_shape = get_output_shape(attrs, input); - TensorShape weights_shape = get_weights_shape(attrs, input); + TensorShape output_shape = get_output_shape(cg, attrs, input); + TensorShape weights_shape = get_weights_shape(attrs, get_shape(cg, input)); - return this->add_layer( - layer, {input}, {{weights_shape, kernel_initializer}}, output_shape); + return insert_layer( + cg, layer, {input}, {{weights_shape, kernel_initializer}}, output_shape); } -std::vector - ComputationGraphBuilder::gather(Tensor const &input, - Tensor const &index, - ff_dim_t dim, - optional const &maybe_name) { +std::vector + insert_gather_layer(ComputationGraph &cg, + tensor_guid_t const &input, + tensor_guid_t const &index, + ff_dim_t dim, + optional const &maybe_name) { GatherAttrs attrs = {dim}; std::string name = maybe_name.value_or(get_default_name(attrs)); + DataType index_dt = get_data_type(cg.at(index)); + Layer layer = {attrs, name}; - if (index.data_type != DataType::INT32 && - index.data_type != DataType::INT64) { + if (index_dt != DataType::INT32 && index_dt != DataType::INT64) { throw mk_runtime_error("Invalid data type for input tensor 2 for Gather: " "{} (should be {} or {})", - input.data_type, + index_dt, DataType::INT32, DataType::INT64); } std::vector output_shapes = - get_output_shapes(attrs, input, index); + get_output_shapes(cg, attrs, input, index); - return this->add_layer(layer, {input}, {}, output_shapes); + return insert_layer(cg, layer, {input, index}, {}, output_shapes); } -TensorShape get_shape(Tensor const &); -std::vector get_shape(std::vector const &); +tensor_guid_t + insert_aggregate_layer(ComputationGraph &cg, + tensor_guid_t const &gate_preds, + tensor_guid_t const &gate_assign, + tensor_guid_t const &true_gate_assign, + tensor_guid_t const &full_gate_gradients, + std::vector const &exp_preds, + int n, + float lambda_bal, + optional const &maybe_name) { + auto get_shape = [&](tensor_guid_t const &t) { + return cg.at(t).shape; + }; -Tensor ComputationGraphBuilder::aggregate( - Tensor const &gate_preds, - Tensor const &gate_assign, - Tensor const &true_gate_assign, - Tensor const &full_gate_gradients, - std::vector const &exp_preds, - int n, - float lambda_bal, - optional const &maybe_name) { AggregateAttrs attrs = {n, lambda_bal}; std::string name = maybe_name.value_or(get_default_name(attrs)); @@ -331,24 +515,26 @@ Tensor ComputationGraphBuilder::aggregate( get_shape(gate_assign), get_shape(true_gate_assign), get_shape(full_gate_gradients), - get_shape(exp_preds)); + transform(exp_preds, get_shape)); - std::vector inputs = { + std::vector inputs = { gate_preds, gate_assign, true_gate_assign, full_gate_gradients}; extend(inputs, exp_preds); - return this->add_layer(layer, inputs, {}, output_shape); + return insert_layer(cg, layer, inputs, {}, output_shape); } -Tensor ComputationGraphBuilder::batch_norm( - Tensor const &input, bool relu, optional const &maybe_name) { +tensor_guid_t insert_batch_norm_layer(ComputationGraph &cg, + tensor_guid_t const &input, + bool relu, + optional const &maybe_name) { BatchNormAttrs attrs = BatchNormAttrs{relu}; std::string name = maybe_name.value_or(get_default_name(attrs)); Layer layer = {attrs, name}; - TensorShape output_shape = get_output_shape(attrs, get_shape(input)); + TensorShape output_shape = get_output_shape(attrs, get_shape(cg, input)); - return this->add_layer(layer, {input}, {}, output_shape); + return insert_layer(cg, layer, {input}, {}, output_shape); } } // namespace FlexFlow diff --git a/lib/runtime/ffi/CMakeLists.txt b/lib/runtime/ffi/CMakeLists.txt index 5140c1ed7b..aee6a9fa15 100644 --- a/lib/runtime/ffi/CMakeLists.txt +++ b/lib/runtime/ffi/CMakeLists.txt @@ -1,4 +1,5 @@ set(target runtime-ffi) +set(internal_target runtime-ffi-internal) project(${target}) file(GLOB_RECURSE SRC @@ -21,12 +22,29 @@ target_link_libraries( ${target} PUBLIC compiler-ffi - utils-ffi + internal-utils-ffi substitutions-ffi pcg-ffi op-attrs-ffi PRIVATE runtime + utils ) ff_set_cxx_properties(${target}) + +add_library( + ${internal_target} + INTERFACE +) +target_link_libraries( + ${internal_target} + INTERFACE + ${target} +) +target_include_directories( + ${internal_target} + INTERFACE + internal/ +) + diff --git a/lib/runtime/ffi/include/flexflow/runtime.h b/lib/runtime/ffi/include/flexflow/runtime.h index e3f25ee6e9..e9b4f6cccf 100644 --- a/lib/runtime/ffi/include/flexflow/runtime.h +++ b/lib/runtime/ffi/include/flexflow/runtime.h @@ -8,19 +8,35 @@ #include #include -FLEXFLOW_FFI_BEGIN() +FLEXFLOW_FFI_BEGIN(); FF_NEW_OPAQUE_TYPE(flexflow_config_t); FF_NEW_OPAQUE_TYPE(flexflow_model_config_t); FF_NEW_OPAQUE_TYPE(flexflow_model_training_instance_t); FF_NEW_OPAQUE_TYPE(flexflow_void_future_t); +// Error handling + typedef enum { FLEXFLOW_RUNTIME_STATUS_OK, FLEXFLOW_RUNTIME_ERROR_UNKNOWN, FLEXFLOW_RUNTIME_ERROR_DYNAMIC_ALLOCATION_FAILED, FLEXFLOW_RUNTIME_ERROR_UNEXPECTED_EMPTY_HANDLE, -} flexflow_runtime_error_t; +} flexflow_runtime_error_code_t; + +FF_NEW_OPAQUE_TYPE(flexflow_runtime_error_t); + +flexflow_error_t flexflow_runtime_error_wrap(flexflow_runtime_error_t); +flexflow_error_t flexflow_runtime_error_unwrap(flexflow_error_t, + flexflow_runtime_error_t *); +flexflow_error_t flexflow_runtime_error_is_ok(flexflow_runtime_error_t, bool *); +char *flexflow_runtime_error_get_string(flexflow_runtime_error_t); +flexflow_runtime_error_code_t + flexflow_runtime_error_get_error_code(flexflow_runtime_error_t, + flexflow_runtime_error_code_t *out); +flexflow_error_t flexflow_runtime_error_destroy(flexflow_runtime_error_t); + +// typedef enum { FLEXFLOW_METRIC_ACCURACY, @@ -44,76 +60,71 @@ typedef enum { FLEXFLOW_COMPUTATION_MODE_INFERENCE, } flexflow_computation_mode_t; -char *flexflow_runtime_get_error_string(flexflow_runtime_error_t); - -flexflow_runtime_error_t flexflow_void_future_wait(flexflow_void_future_t); -flexflow_runtime_error_t flexflow_void_future_destroy(flexflow_void_future_t); +flexflow_error_t flexflow_void_future_wait(flexflow_void_future_t); +flexflow_error_t flexflow_void_future_destroy(flexflow_void_future_t); -flexflow_runtime_error_t flexflow_config_parse_argv(int *argc, - char **argv, - bool remove_used, - flexflow_config_t *out); -flexflow_runtime_error_t flexflow_set_config(flexflow_config_t); -flexflow_runtime_error_t flexflow_get_config(flexflow_config_t *); +flexflow_error_t flexflow_config_parse_argv(int *argc, + char **argv, + bool remove_used, + flexflow_config_t *out); +flexflow_error_t flexflow_set_config(flexflow_config_t); +flexflow_error_t flexflow_get_config(flexflow_config_t *); -flexflow_runtime_error_t flexflow_model_config_parse_argv( - int *argc, char **argv, bool remove_used, flexflow_model_config_t *out); +flexflow_error_t flexflow_model_config_parse_argv(int *argc, + char **argv, + bool remove_used, + flexflow_model_config_t *out); -flexflow_runtime_error_t +flexflow_error_t flexflow_computation_graph_set_model_config(flexflow_computation_graph_t, flexflow_model_config_t); -flexflow_runtime_error_t +flexflow_error_t flexflow_computation_graph_get_model_config(flexflow_computation_graph_t, flexflow_model_config_t *out); -flexflow_runtime_error_t flexflow_computation_graph_compile( +flexflow_error_t flexflow_computation_graph_compile( flexflow_computation_graph_t, flexflow_optimizer_t, flexflow_model_compilation_result_t *out); -flexflow_runtime_error_t flexflow_model_compilation_result_get_pcg( +flexflow_error_t flexflow_model_compilation_result_get_pcg( flexflow_model_compilation_result_t, flexflow_parallel_computation_graph_t *out); -flexflow_runtime_error_t +flexflow_error_t flexflow_model_compilation_result_get_parallel_tensor_for_tensor( flexflow_model_compilation_result_t, flexflow_tensor_t, flexflow_parallel_tensor_t *); -flexflow_runtime_error_t - flexflow_start_training(flexflow_model_compilation_result_t, - flexflow_model_training_instance_t *); -flexflow_runtime_error_t +flexflow_error_t flexflow_start_training(flexflow_model_compilation_result_t, + flexflow_model_training_instance_t *); +flexflow_error_t flexflow_model_training_instance_forward(flexflow_model_training_instance_t, flexflow_void_future_t *out); -flexflow_runtime_error_t flexflow_model_training_instance_backward( +flexflow_error_t flexflow_model_training_instance_backward( flexflow_model_training_instance_t); -flexflow_runtime_error_t - flexflow_stop_training(flexflow_model_training_instance_t); - -flexflow_runtime_error_t - flexflow_get_tensor_float(flexflow_model_training_instance_t, - flexflow_tensor_t, - float *data, - bool get_gradients); -flexflow_runtime_error_t - flexflow_get_tensor_double(flexflow_model_training_instance_t, - flexflow_tensor_t, - float *data, - bool get_gradients); -flexflow_runtime_error_t - flexflow_get_tensor_int32(flexflow_model_training_instance_t, - flexflow_tensor_t, - int32_t *data, - bool get_gradients); -flexflow_runtime_error_t - flexflow_get_tensor_int64(flexflow_model_training_instance_t, - flexflow_tensor_t, - int64_t *data, - bool get_gradients); - -flexflow_runtime_error_t flexflow_set_tensor_int( - flexflow_model_training_instance_t, flexflow_tensor_t, int32_t *data); - -FLEXFLOW_FFI_END() +flexflow_error_t flexflow_stop_training(flexflow_model_training_instance_t); + +flexflow_error_t flexflow_get_tensor_float(flexflow_model_training_instance_t, + flexflow_tensor_t, + float *data, + bool get_gradients); +flexflow_error_t flexflow_get_tensor_double(flexflow_model_training_instance_t, + flexflow_tensor_t, + float *data, + bool get_gradients); +flexflow_error_t flexflow_get_tensor_int32(flexflow_model_training_instance_t, + flexflow_tensor_t, + int32_t *data, + bool get_gradients); +flexflow_error_t flexflow_get_tensor_int64(flexflow_model_training_instance_t, + flexflow_tensor_t, + int64_t *data, + bool get_gradients); + +flexflow_error_t flexflow_set_tensor_int(flexflow_model_training_instance_t, + flexflow_tensor_t, + int32_t *data); + +FLEXFLOW_FFI_END(); #endif diff --git a/lib/runtime/ffi/internal/internal/runtime.h b/lib/runtime/ffi/internal/internal/runtime.h new file mode 100644 index 0000000000..2e9a4862f1 --- /dev/null +++ b/lib/runtime/ffi/internal/internal/runtime.h @@ -0,0 +1,6 @@ +#ifndef _FLEXFLOW_RUNTIME_FFI_INTERNAL_INTERNAL_RUNTIME_H +#define _FLEXFLOW_RUNTIME_FFI_INTERNAL_INTERNAL_RUNTIME_H + +#include "flexflow/runtime.h" + +#endif diff --git a/lib/runtime/ffi/src/runtime.cc b/lib/runtime/ffi/src/runtime.cc index afc508419e..5de26d39d0 100644 --- a/lib/runtime/ffi/src/runtime.cc +++ b/lib/runtime/ffi/src/runtime.cc @@ -1,10 +1,14 @@ #include "flexflow/runtime.h" +#include "internal/opaque.h" #include "runtime/model_training_instance.h" #include "utils/expected.h" -#include "utils/ffi/opaque.h" using namespace FlexFlow; +flexflow_error_t flexflow_runtime_error_create(flexflow_runtime_error_t) { + NOT_IMPLEMENTED(); +} + using Runtime = LibraryUtils -#include -#include +/* #include "kernels/array_shape.h" */ +/* #include "legion.h" */ +/* #include "legion_tensor_shape.h" */ +/* #include "op-attrs/datatype.h" */ +/* #include "op-attrs/param_sync.h" */ +/* #include "op-attrs/tensor_shape.h" */ +/* #include "utils/optional.h" */ +/* #include "utils/stack_vector.h" */ +/* #include "utils/visitable.h" */ +/* #include */ +/* #include */ +/* #include */ -namespace FlexFlow { +/* namespace FlexFlow { */ -struct FFModel; +/* struct Tensor { */ +/* Tensor() = delete; */ +/* Tensor(TensorShape const &, */ +/* CreateGrad create_gradients, */ +/* optional const &initializer = nullopt, */ +/* optional sync_type = nullopt); */ -struct Tensor : public use_visitable_cmp { - Tensor() = delete; - Tensor(TensorShape const &, - CreateGrad create_gradients, - optional initializer = nullopt, - optional sync_type = nullopt); +/* size_t get_volume() const; */ +/* Legion::Domain get_domain() const; */ +/* TensorShape get_shape() const; */ +/* int num_dims() const; */ - size_t get_volume() const; - Legion::Domain get_domain() const; - TensorShape get_shape() const; - int num_dims() const; +/* operator TensorShape const &() const; */ +/* public: */ +/* TensorShape shape; */ +/* CreateGrad create_gradients; */ +/* optional initializer; */ +/* optional sync_type; */ +/* }; */ +/* FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(Tensor, shape, create_gradients, + * initializer, sync_type); */ - operator TensorShape const &() const; +/* template */ +/* bool set_tensor(Tensor const &, */ +/* FFModel const *model, */ +/* std::vector const &dims, */ +/* T const *data); */ +/* template */ +/* bool get_tensor(Tensor const &, */ +/* FFModel const *model, */ +/* T *data, */ +/* bool get_gradients); */ - friend void swap(Tensor &, Tensor &); - -public: - TensorDims dims; - DataType data_type; - optional initializer; - bool create_gradients; - optional sync_type; -}; - -template -bool set_tensor(Tensor const &, - FFModel const *model, - std::vector const &dims, - T const *data); -template -bool get_tensor(Tensor const &, - FFModel const *model, - T *data, - bool get_gradients); - -static_assert(std::is_copy_constructible::value, - "Tensor must be copy constructible"); - -using Parameter = Tensor; +/* using Parameter = Tensor; */ } // namespace FlexFlow -VISITABLE_STRUCT(::FlexFlow::Tensor, - dims, - data_type, - initializer, - create_gradients, - sync_type); -MAKE_VISIT_HASHABLE(::FlexFlow::Tensor); - #endif diff --git a/lib/utils/ffi/CMakeLists.txt b/lib/utils/ffi/CMakeLists.txt index 6f5bdec4ff..87fd32fdba 100644 --- a/lib/utils/ffi/CMakeLists.txt +++ b/lib/utils/ffi/CMakeLists.txt @@ -1,4 +1,5 @@ set(target utils-ffi) +set(internal_target "internal-${target}") project(${target}) file(GLOB_RECURSE FFI_SRC @@ -16,6 +17,7 @@ target_include_directories( include/ PRIVATE src/ + internal/ ) target_link_libraries( ${target} @@ -24,3 +26,18 @@ target_link_libraries( ) ff_set_cxx_properties(${target}) + +add_library( + ${internal_target} + INTERFACE +) +target_link_libraries( + ${internal_target} + INTERFACE + ${target} +) +target_include_directories( + ${internal_target} + INTERFACE + internal/ +) diff --git a/lib/utils/ffi/include/flexflow/utils.h b/lib/utils/ffi/include/flexflow/utils.h index a87ef4c630..084cd134ba 100644 --- a/lib/utils/ffi/include/flexflow/utils.h +++ b/lib/utils/ffi/include/flexflow/utils.h @@ -23,11 +23,34 @@ typedef enum { FLEXFLOW_ERROR_SOURCE_PCG, FLEXFLOW_ERROR_SOURCE_COMPILER, FLEXFLOW_ERROR_SOURCE_OPATTRS, + FLEXFLOW_ERROR_SOURCE_UTILS, } flexflow_error_source_t; +typedef enum { + FLEXFLOW_UTILS_STATUS_OK, + FLEXFLOW_UTILS_DEALLOCATION_FAILED, + FLEXFLOW_UTILS_ALLOCATION_FAILED, + FLEXFLOW_UTILS_CAST_FAILED, + FLEXFLOW_UTILS_INVALID_ERROR_SOURCE, + FLEXFLOW_UTILS_UNEXPECTED_NULLPTR_IN_OPAQUE_HANDLE, +} flexflow_utils_error_code_t; + +typedef struct { + flexflow_utils_error_code_t err_code; +} flexflow_utils_error_t; + +#define FLEXFLOW_FFI_ERROR_BUF_SIZE 24 + typedef struct { flexflow_error_source_t error_source; - int error_code; + char buf[FLEXFLOW_FFI_ERROR_BUF_SIZE]; } flexflow_error_t; +flexflow_error_t flexflow_utils_error_is_ok(flexflow_utils_error_t, bool *); +flexflow_error_t flexflow_utils_error_create(flexflow_utils_error_code_t); +flexflow_error_t flexflow_utils_error_unwrap(flexflow_error_t, + flexflow_utils_error_t *); + +flexflow_error_t flexflow_status_is_ok(flexflow_error_t, bool *); + #endif diff --git a/lib/utils/ffi/internal/internal/enums.h b/lib/utils/ffi/internal/internal/enums.h new file mode 100644 index 0000000000..02e6b47bbd --- /dev/null +++ b/lib/utils/ffi/internal/internal/enums.h @@ -0,0 +1,54 @@ +#ifndef _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ENUMS_H +#define _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ENUMS_H + +template +struct internal_to_external; + +template +struct external_to_internal; + +template +struct enum_mapping; + +template +using internal_to_external_t = typename internal_to_external::type; + +template +using external_to_internal_t = typename external_to_internal::type; + +#define REGISTER_FFI_ENUM(EXTERNAL, INTERNAL, ERROR_CODE, ...) \ + template <> \ + struct external_to_internal { \ + using type = INTERNAL; \ + }; \ + template <> \ + struct internal_to_external { \ + using type = EXTERNAL; \ + }; \ + template <> \ + struct enum_mapping { \ + static const bidict mapping; \ + static constexpr decltype(ERROR_CODE) err_code = ERROR_CODE; \ + }; \ + const bidict enum_mapping::mapping = \ + __VA_ARGS__; + +template +external_to_internal_t to_internal_impl(ExternalEnum e) { + return enum_mapping::mapping.maybe_at_l(e) + .or_else([] { + throw make_opattrs_error(enum_mapping::err_code); + }) + .value(); +} + +template +internal_to_external_t to_external_impl(InternalEnum i) { + using Mapping = enum_mapping>; + + return Mapping::mapping.maybe_at_r(i) + .or_else([] { throw make_opattrs_error(Mapping::err_code); }) + .value(); +} + +#endif diff --git a/lib/utils/ffi/internal/internal/error.h b/lib/utils/ffi/internal/internal/error.h new file mode 100644 index 0000000000..dbf689f24b --- /dev/null +++ b/lib/utils/ffi/internal/internal/error.h @@ -0,0 +1,57 @@ +#ifndef _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ERROR_H +#define _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ERROR_H + +#include "flexflow/utils.h" +#include "utils/type_traits_core.h" + +#define RAISE_FLEXFLOW(status) \ + do { \ + bool is_ok; \ + flexflow_status_is_ok(status, &is_ok); \ + if (!is_ok) { \ + return status; \ + } \ + } while (0) + +struct flexflow_ffi_exception_t : public std::runtime_error { + flexflow_ffi_exception_t(flexflow_error_t); + + flexflow_error_t err; +}; + +flexflow_ffi_exception_t make_utils_exception(flexflow_utils_error_code_t); + +using flexflow_utils_exception_t = flexflow_ffi_exception_t; + +flexflow_error_t to_error(flexflow_utils_exception_t const &); +flexflow_error_t status_ok(); + +template +flexflow_error_t flexflow_error_wrap(flexflow_error_source_t error_source, + T const &t) { + static_assert(sizeof(T) < (FLEXFLOW_FFI_ERROR_BUF_SIZE * sizeof(char)), ""); + + flexflow_error_t result; + result.error_source = error_source; + T *buf_ptr = static_cast(result.buf); + *buf_ptr = t; + + return result; +} + +template +flexflow_error_t flexflow_error_unwrap(flexflow_error_t const &err, + flexflow_error_source_t error_source, + T *out) { + static_assert(sizeof(T) < (FLEXFLOW_FFI_ERROR_BUF_SIZE * sizeof(char)), ""); + + if (err.error_source != FLEXFLOW_ERROR_SOURCE_PCG) { + return flexflow_utils_error_create(FLEXFLOW_UTILS_CAST_FAILED); + } + + out->impl = + reinterpret_cast()))>(err.buf); + return flexflow_utils_error_create(FLEXFLOW_UTILS_STATUS_OK); +} + +#endif diff --git a/lib/utils/ffi/internal/internal/error_handling.h b/lib/utils/ffi/internal/internal/error_handling.h new file mode 100644 index 0000000000..441d12e86c --- /dev/null +++ b/lib/utils/ffi/internal/internal/error_handling.h @@ -0,0 +1,16 @@ +#ifndef _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ERROR_HANDLING_H +#define _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ERROR_HANDLING_H + +#include "error.h" +#include "opaque.h" + +flexflow_error_t handle_errors(std::function const &f); + +template +flexflow_error_t + handle_errors(Opaque *out, + std::function()> const &f) { + return handle_errors([&] { *out = new_opaque(f()); }); +} + +#endif diff --git a/lib/utils/ffi/internal/internal/opaque.h b/lib/utils/ffi/internal/internal/opaque.h new file mode 100644 index 0000000000..35c929aca0 --- /dev/null +++ b/lib/utils/ffi/internal/internal/opaque.h @@ -0,0 +1,97 @@ +#ifndef _FLEXFLOW_UTILS_INCLUDE_UTILS_FFI_OPAQUE_H +#define _FLEXFLOW_UTILS_INCLUDE_UTILS_FFI_OPAQUE_H + +#include "error.h" +#include "flexflow/utils.h" +#include "utils/containers.h" +#include "utils/expected.h" +#include +#include +#include + +template +struct opaque_to_underlying; + +template +struct underlying_to_opaque; + +template +using opaque_to_underlying_t = typename opaque_to_underlying::type; + +template +using underlying_to_opaque_t = typename underlying_to_opaque::type; + +template +opaque_to_underlying_t *unwrap_opaque(Opaque const &opaque) { + if (opaque.impl == nullptr) { + throw make_utils_exception( + FLEXFLOW_UTILS_UNEXPECTED_NULLPTR_IN_OPAQUE_HANDLE); + } + + return static_cast *>(opaque.impl); +} + +template +opaque_to_underlying_t const *c_unwrap_opaque(Opaque const &opaque) { + return unwrap_opaque(opaque); +} + +template +opaque_to_underlying_t const &c_deref_opaque(Opaque const &opaque) { + return *unwrap_opaque(opaque); +} + +template +opaque_to_underlying_t &deref_opaque(Opaque const &opaque) { + return *unwrap_opaque(opaque); +} + +template +std::vector> + c_deref_opaque_list(Opaque const *start, size_t num_elements) { + std::vector exp_preds_vector = + transform(start, start + num_elements, [](Opaque const &t) { + return c_deref_opaque(t); + }); +} + +#define REGISTER_OPAQUE(OPAQUE, UNDERLYING) \ + template <> \ + struct opaque_to_underlying { \ + using type = UNDERLYING; \ + }; \ + template <> \ + struct underlying_to_opaque { \ + using type = OPAQUE; \ + }; + +template +Opaque new_opaque(Args &&...args) { + using Underlying = opaque_to_underlying_t; + + Underlying *ptr = new (std::nothrow) Underlying(std::forward(args)...); + if (ptr == nullptr) { + throw make_utils_exception(FLEXFLOW_UTILS_ALLOCATION_FAILED); + } + return Opaque{ptr}; +} + +template +underlying_to_opaque_t new_opaque(Underlying const &underlying) { + return new_opaque>(underlying); +} + +template +void delete_opaque(Opaque const &opaque) { + using Underlying = opaque_to_underlying_t; + + Underlying *underlying = unwrap_opaque(opaque); + if (underlying == nullptr) { + throw make_utils_exception( + FLEXFLOW_UTILS_UNEXPECTED_NULLPTR_IN_OPAQUE_HANDLE); + } + + delete underlying; +} + +#endif diff --git a/lib/utils/ffi/src/error.cc b/lib/utils/ffi/src/error.cc new file mode 100644 index 0000000000..c55f602141 --- /dev/null +++ b/lib/utils/ffi/src/error.cc @@ -0,0 +1,5 @@ +#include "internal/error.h" + +flexflow_error_t status_ok() { + return flexflow_utils_error_create(FLEXFLOW_UTILS_STATUS_OK); +} diff --git a/lib/utils/ffi/src/error_handling.cc b/lib/utils/ffi/src/error_handling.cc new file mode 100644 index 0000000000..7a5f5b333f --- /dev/null +++ b/lib/utils/ffi/src/error_handling.cc @@ -0,0 +1,11 @@ +#include "internal/error_handling.h" + +flexflow_error_t handle_errors(std::function const &f) { + try { + f(); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} diff --git a/lib/utils/include/utils/bidict.h b/lib/utils/include/utils/bidict.h index 870afd0448..611a806c95 100644 --- a/lib/utils/include/utils/bidict.h +++ b/lib/utils/include/utils/bidict.h @@ -1,14 +1,25 @@ #ifndef _FLEXFLOW_UTILS_BIDICT_H #define _FLEXFLOW_UTILS_BIDICT_H +#include "optional.h" #include namespace FlexFlow { template struct bidict { + using const_iterator = typename std::unordered_map::const_iterator; + using value_type = std::pair; + using reference = value_type &; + using const_reference = value_type const &; + using key_type = L; + using mapped_type = R; + bidict() : fwd_map{}, bwd_map{} {} + bidict(std::initializer_list init) + : bidict(init.begin(), init.end()) {} + template bidict(InputIt first, InputIt last) { for (auto it = first; it != last; it++) { @@ -54,17 +65,27 @@ struct bidict { return bwd_map.at(r); } + optional maybe_at_l(L const &l) const { + if (fwd_map.count(l) != 0) { + return fwd_map.at(l); + } else { + return nullopt; + } + } + + optional maybe_at_r(R const &r) const { + if (bwd_map.count(r) != 0) { + return bwd_map.at(r); + } else { + return nullopt; + } + } + std::size_t size() const { assert(fwd_map.size() == bwd_map.size()); return fwd_map.size(); } - using const_iterator = typename std::unordered_map::const_iterator; - using value_type = std::pair; - using reference = value_type &; - using const_reference = value_type const &; - using key_type = L; - using mapped_type = R; /* struct const_iterator { */ /* using iterator_category = std::forward_iterator_tag; */ /* using difference_type = std::size_t; */ diff --git a/lib/utils/include/utils/containers.h b/lib/utils/include/utils/containers.h index a0939f7038..3a1b995275 100644 --- a/lib/utils/include/utils/containers.h +++ b/lib/utils/include/utils/containers.h @@ -372,7 +372,9 @@ std::unordered_set template typename C::value_type get_only(C const &c) { - assert(c.size() == 1); + if (c.size() != 1) { + throw mk_runtime_error("get_only called on container of size {}", c.size()); + } return *c.cbegin(); } @@ -449,6 +451,15 @@ std::vector transform(std::vector const &v, F const &f) { return result; } +template ()(*std::declval()))> +std::vector transform(It start, It end, F const &f) { + std::vector result; + std::transform(start, end, std::back_inserter(result), f); + return result; +} + template auto transform(req const &c, F const &f) -> decltype(transform(std::declval(), std::declval())) { diff --git a/lib/utils/include/utils/ffi/opaque.h b/lib/utils/include/utils/ffi/opaque.h deleted file mode 100644 index bf4f62cca8..0000000000 --- a/lib/utils/include/utils/ffi/opaque.h +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef _FLEXFLOW_UTILS_INCLUDE_UTILS_FFI_OPAQUE_H -#define _FLEXFLOW_UTILS_INCLUDE_UTILS_FFI_OPAQUE_H - -#include "utils/expected.h" -#include -#include - -namespace FlexFlow { - -template -struct LibraryUtils { - template - using err = expected; - - template - static err allocate_opaque(T const &t) { - T *ptr = new (std::nothrow) T(t); - if (ptr == nullptr) { - return StatusOK; - } - return ptr; - } - - template - static err allocate_opaque(T &&t) { - T *ptr = new (std::nothrow) T(std::move(t)); - if (ptr == nullptr) { - return AllocationFailed; - } - return ptr; - } - - template < - typename Opaque, - typename Unwrapped = decltype(*unwrap_opaque(std::declval()))> - static err new_opaque(Unwrapped const &f) { - return allocate_opaque(f).map([](Unwrapped *ptr) { return ptr; }); - } - - template - static ErrorCodeType output_stored(T const &t, T *out) { - return output_stored(new_opaque(t), out); - } - - template - static ErrorCodeType output_stored(err const &e, T *out) { - if (e.has_value()) { - *out = e.value(); - return StatusOK; - } else { - return e.error(); - } - } - - template - ErrorCodeType deallocate_opaque(T const &opaque) { - auto unwrapped = unwrap_opaque(opaque); - if (unwrapped == nullptr) { - return UnexpectedNull; - } - - delete unwrapped; - return StatusOK; - } -}; - -} // namespace FlexFlow - -#endif diff --git a/lib/utils/include/utils/graph/algorithms.h b/lib/utils/include/utils/graph/algorithms.h index 9a36b2c601..e4d7675aa9 100644 --- a/lib/utils/include/utils/graph/algorithms.h +++ b/lib/utils/include/utils/graph/algorithms.h @@ -93,7 +93,11 @@ std::unordered_set get_node_edges(UndirectedGraphView const &, Node const &); std::unordered_set get_outputs(MultiDiGraphView const &); +std::unordered_set get_outputs(MultiDiGraphView const &, + Node const &); std::unordered_set get_inputs(MultiDiGraphView const &); +std::unordered_set get_inputs(MultiDiGraphView const &, + Node const &); std::unordered_set get_incoming_edges(MultiDiGraphView const &, Node const &); diff --git a/lib/utils/include/utils/strong_typedef.h b/lib/utils/include/utils/strong_typedef.h index bc9e6872f8..b8c6babb69 100644 --- a/lib/utils/include/utils/strong_typedef.h +++ b/lib/utils/include/utils/strong_typedef.h @@ -2,6 +2,8 @@ #define _FLEXFLOW_UTILS_INCLUDE_STRONG_TYPEDEF_H #include "utils/fmt.h" +#include "utils/test_types.h" +#include "utils/type_traits.h" #include #include #include @@ -24,7 +26,7 @@ class strong_typedef { return value_; } - explicit operator T const &() const noexcept { + operator T() const noexcept { return value_; } @@ -66,7 +68,7 @@ class strong_typedef { } template - strong_typedef fmap(F const &f) { + strong_typedef fmap(F const &f) const { static_assert( std::is_same()(std::declval())), T>::value, @@ -75,6 +77,11 @@ class strong_typedef { return strong_typedef(f(this->value_)); } + template + decltype(std::declval()(std::declval())) fmap(F const &f) { + f(this->value_); + } + private: T value_; }; @@ -102,98 +109,99 @@ template struct numerical_typedef : strong_typedef { using strong_typedef::strong_typedef; - friend StrongTypedef &operator+=(StrongTypedef &lhs, T const &rhs) { - static_cast(lhs) += static_cast(rhs); - return lhs; + friend numerical_typedef operator+=(numerical_typedef lhs, T const &rhs) { + return numerical_typedef{lhs.value() += rhs}; } - friend StrongTypedef &operator++(StrongTypedef &lhs) { - static_cast(lhs) += static_cast(1); - return lhs; + friend numerical_typedef &operator++(numerical_typedef &lhs) { + return numerical_typedef{lhs.value()++}; } - friend StrongTypedef operator++(StrongTypedef &lhs, int) { - StrongTypedef tmp = lhs; - ++lhs; - return tmp; + friend numerical_typedef operator++(numerical_typedef &lhs, int) { + return numerical_typedef{++lhs.value()}; } - friend StrongTypedef operator+(StrongTypedef const &lhs, T const &rhs) { - return StrongTypedef(lhs.value() + rhs); + friend numerical_typedef operator+(numerical_typedef const &lhs, + T const &rhs) { + return numerical_typedef{lhs.value() + rhs}; } - friend StrongTypedef operator+(T const &lhs, StrongTypedef const &rhs) { - return (rhs + lhs); - } + /* friend numerical_typedef operator+(T const &lhs, numerical_typedef const + * &rhs) { */ + /* return (rhs + lhs); */ + /* } */ - friend StrongTypedef operator-=(StrongTypedef &lhs, T const &rhs) { + friend numerical_typedef operator-=(numerical_typedef &lhs, T const &rhs) { static_cast(lhs) -= static_cast(rhs); } - friend StrongTypedef &operator--(StrongTypedef &lhs) { + friend numerical_typedef &operator--(numerical_typedef &lhs) { static_cast(lhs) -= static_cast(1); return lhs; } - friend StrongTypedef operator--(StrongTypedef &lhs, int) { - StrongTypedef tmp = lhs; + friend numerical_typedef operator--(numerical_typedef &lhs, int) { + numerical_typedef tmp = lhs; --lhs; return tmp; } - friend StrongTypedef operator-(StrongTypedef const &lhs, T const &rhs) { - return StrongTypedef(lhs.value() + rhs); - } - - friend bool operator<(StrongTypedef const &lhs, T const &rhs) { - return lhs.value() < rhs; - } - - friend bool operator==(StrongTypedef const &lhs, T const &rhs) { - return lhs.value() == rhs; - } - - friend bool operator>(StrongTypedef const &lhs, T const &rhs) { - return lhs.value() > rhs; + friend numerical_typedef operator-(numerical_typedef const &lhs, + T const &rhs) { + return numerical_typedef(lhs.value() + rhs); } - friend bool operator>=(StrongTypedef const &lhs, T const &rhs) { - return lhs.value() >= rhs; - } + /* friend bool operator<(numerical_typedef const &lhs, numerical_typedef const + * &rhs) { */ + /* return lhs.value() < rhs; */ + /* } */ - friend bool operator!=(StrongTypedef const &lhs, T const &rhs) { - return lhs.value() != rhs; - } + /* friend bool operator==(numerical_typedef const &lhs, numerical_typedef + * const &rhs) { */ + /* return lhs.value() == rhs; */ + /* } */ - friend bool operator<=(StrongTypedef const &lhs, T const &rhs) { - return lhs.value() <= rhs; - } + /* friend bool operator>(numerical_typedef const &lhs, numerical_typedef const + * &rhs) { */ + /* return lhs.value() > rhs; */ + /* } */ - friend bool operator<(T const &lhs, StrongTypedef const &rhs) { - return lhs < rhs.value(); - } + /* friend bool operator>=(numerical_typedef const &lhs, numerical_typedef + * const &rhs) { */ + /* return lhs.value() >= rhs; */ + /* } */ - friend bool operator==(T const &lhs, StrongTypedef const &rhs) { - return lhs == rhs.value(); - } - - friend bool operator>(T const &lhs, StrongTypedef const &rhs) { - return lhs > rhs.value(); - } + /* friend bool operator!=(numerical_typedef const &lhs, numerical_typedef + * const &rhs) { */ + /* return lhs.value() != rhs; */ + /* } */ - friend bool operator<=(T const &lhs, StrongTypedef const &rhs) { - return lhs <= rhs.value(); - } - - friend bool operator!=(T const &lhs, StrongTypedef const &rhs) { - return lhs != rhs.value(); - } - - friend bool operator>=(T const &lhs, StrongTypedef const &rhs) { - return lhs >= rhs.value(); - } + /* friend bool operator<=(T const &lhs, T const &rhs) { */ + /* return lhs.value() <= rhs; */ + /* } */ }; +template +struct is_addable + : conjunction< + std::is_same() + std::declval())>, + implies, + std::is_same() += + std::declval())>>> {}; + +static_assert(is_neq_comparable>::value, + ""); +static_assert(is_neq_comparable>::value, + ""); +static_assert(is_lt_comparable>::value, + ""); +static_assert(is_lt_comparable>::value, + ""); +static_assert(is_addable>::value, + ""); + } // namespace FlexFlow #define MAKE_TYPEDEF_HASHABLE(TYPEDEF_NAME) \ diff --git a/lib/utils/include/utils/test_types.h b/lib/utils/include/utils/test_types.h index 514d030b6f..905517d5e8 100644 --- a/lib/utils/include/utils/test_types.h +++ b/lib/utils/include/utils/test_types.h @@ -7,13 +7,21 @@ namespace FlexFlow { namespace test_types { -enum capability { HASHABLE, EQ, CMP, DEFAULT_CONSTRUCTIBLE, COPYABLE }; +enum capability { + HASHABLE, + EQ, + CMP, + DEFAULT_CONSTRUCTIBLE, + COPYABLE, + PLUS, + PLUSEQ +}; template struct capability_implies : std::false_type {}; template <> -struct capability_implies : std::true_type {}; +struct capability_implies : std::true_type {}; template struct capability_implies : std::true_type {}; @@ -34,45 +42,60 @@ struct test_type_t { template using supports = conjunction...>; - template ::value, - bool>::type = true> + template ::value, bool>::type = true> test_type_t(); - template ::value, - bool>::type = true> + template ::value, bool>::type = true> test_type_t() = delete; - template < - typename std::enable_if::value, bool>::type = true> + template ::value, bool>::type = true> test_type_t(test_type_t const &); - template < - typename std::enable_if::value, bool>::type = true> + template ::value, bool>::type = true> test_type_t(test_type_t const &) = delete; - typename std::enable_if::value, bool>::type + template + typename std::enable_if::value, bool>::type operator==(test_type_t const &) const; - typename std::enable_if::value, bool>::type + template + typename std::enable_if::value, bool>::type operator!=(test_type_t const &) const; - typename std::enable_if::value, bool>::type + template + typename std::enable_if::value, bool>::type operator<(test_type_t const &) const; - typename std::enable_if::value, bool>::type + template + typename std::enable_if::value, bool>::type operator>(test_type_t const &) const; - typename std::enable_if::value, bool>::type + template + typename std::enable_if::value, bool>::type operator<=(test_type_t const &) const; - typename std::enable_if::value, bool>::type + template + typename std::enable_if::value, bool>::type operator>=(test_type_t const &) const; + + template + typename std::enable_if::value, test_type_t>::type + operator+(test_type_t const &); + + template + typename std::enable_if::value, test_type_t>::type + operator+=(test_type_t const &); }; using no_eq = test_type_t<>; using eq = test_type_t; using cmp = test_type_t; using hash_cmp = test_type_t; +using plusable = test_type_t; } // namespace test_types } // namespace FlexFlow diff --git a/lib/utils/include/utils/type_traits.h b/lib/utils/include/utils/type_traits.h index 7b345533fe..92f4ff1ff0 100644 --- a/lib/utils/include/utils/type_traits.h +++ b/lib/utils/include/utils/type_traits.h @@ -50,6 +50,13 @@ struct is_rc_copy_virtual_compliant std::is_move_assignable>>, std::has_virtual_destructor> {}; +template +struct is_clonable : std::false_type {}; + +template +struct is_clonable().clone())>> + : std::true_type {}; + template struct is_streamable : std::false_type {}; diff --git a/lib/utils/include/utils/type_traits_core.h b/lib/utils/include/utils/type_traits_core.h index acda070ad8..4849456731 100644 --- a/lib/utils/include/utils/type_traits_core.h +++ b/lib/utils/include/utils/type_traits_core.h @@ -1,6 +1,8 @@ #ifndef _FLEXFLOW_UTILS_INCLUDE_UTILS_TYPE_TRAITS_CORE_H #define _FLEXFLOW_UTILS_INCLUDE_UTILS_TYPE_TRAITS_CORE_H +#include +#include #include namespace FlexFlow { @@ -91,6 +93,31 @@ struct pack_contains_type> : pack_contains_type {}; template struct pack_contains_type : std::false_type {}; +template +struct pack_get { + using type = typename pack_get<(i - 1), Args...>::type; +}; + +template +struct pack_get<0, Head, Tail...> { + using type = Head; +}; + +template +struct pack_get { + static_assert(i > 0, "Out of bounds access for pack_get"); +}; + +template +struct pack_size; + +template +struct pack_size + : std::integral_constant::value + 1)> {}; + +template <> +struct pack_size<> : std::integral_constant {}; + static_assert(pack_contains_type::value, ""); static_assert(!pack_contains_type::value, ""); diff --git a/lib/utils/include/utils/visitable.h b/lib/utils/include/utils/visitable.h index aa6ecb4cbf..e6aaf878cd 100644 --- a/lib/utils/include/utils/visitable.h +++ b/lib/utils/include/utils/visitable.h @@ -337,10 +337,13 @@ struct Arbitrary< MAKE_VISIT_HASHABLE(::FlexFlow::TYPENAME); \ namespace FlexFlow { \ CHECK_WELL_BEHAVED_VISIT_TYPE(TYPENAME); \ - static_assert(is_only_visit_list_initializable::value, \ - #TYPENAME \ - " should not be list-initialializable from any sub-tuples " \ - "(you probably need to insert req<...>s)"); \ + static_assert( \ + is_only_visit_list_initializable::value, \ + #TYPENAME \ + " should not be list-initialializable from any sub-tuples " \ + "(you probably need to insert req<...>s, or make sure your " \ + "field ordering in FF_VISITABLE_STRUCT matches the order in the" \ + "struct, or use FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION"); \ static_assert(!std::is_default_constructible::value, \ #TYPENAME " should not be default-constructible (you " \ "probably need to insert req<...>s)") diff --git a/scripts/format.sh b/scripts/format.sh index 9610dc2d26..4aaf978bc1 100755 --- a/scripts/format.sh +++ b/scripts/format.sh @@ -69,7 +69,13 @@ if [[ ! -e $CLANG_FORMAT_PATH ]]; then fi CLANG_FORMAT_CONFIG="$GIT_ROOT/.clang-format-for-format-sh" -mapfile -t FILES < <(git ls-files ':!:triton/**' '*.h' '*.cc' '*.cpp' '*.cu' '*.c') +mapfile -t ALL_MODIFIED_FILES < <(git ls-files ':!:triton/**' '*.h' '*.cc' '*.cpp' '*.cu' '*.c') +mapfile -t DELETED_FILES < <(git ls-files -d) + +# set difference -- see https://unix.stackexchange.com/questions/443575/how-to-subtract-two-list-fast +# used to avoid trying to format deleted files +FILES=($(comm -3 <(printf "%s\n" "${ALL_MODIFIED_FILES[@]}" | sort) <(printf "%s\n" "${DELETED_FILES[@]}" | sort) | sort -n)) + if [[ -f $CLANG_FORMAT_CONFIG ]]; then "$CLANG_FORMAT_PATH" --style=file:"$CLANG_FORMAT_CONFIG" -i "${FILES[@]}" else