From 3e65ff8147fd469ffe64a34b3330665c13a1077d Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Wed, 2 Aug 2023 04:11:05 -0700 Subject: [PATCH 1/4] Continue implementing pcg ffi internals --- lib/CMakeLists.txt | 4 +- lib/compiler/ffi/include/flexflow/compiler.h | 25 +- lib/compiler/ffi/internal/internal/compiler.h | 6 + lib/ffi/CMakeLists.txt | 2 + lib/ffi/include/flexflow/flexflow.h | 21 +- lib/ffi/src/ffi.cc | 73 +++ lib/op-attrs/ffi/CMakeLists.txt | 18 + lib/op-attrs/ffi/include/flexflow/op-attrs.h | 31 +- lib/op-attrs/ffi/internal/internal/op-attrs.h | 37 ++ lib/op-attrs/ffi/src/op-attrs.cc | 190 ++++++ lib/pcg/ffi/CMakeLists.txt | 22 +- lib/pcg/ffi/include/flexflow/pcg.h | 57 +- lib/pcg/ffi/internal/internal/pcg.h | 5 + lib/pcg/ffi/src/pcg.cc | 592 +++++++++++++++++- lib/pcg/ffi/src/pcg.h | 33 + .../include/pcg/computation_graph_builder.h | 525 ++++++++-------- lib/pcg/include/pcg/file_format/v1/graphs.h | 1 + lib/pcg/include/pcg/tensor.h | 11 +- lib/runtime/ffi/CMakeLists.txt | 20 +- lib/runtime/ffi/include/flexflow/runtime.h | 64 +- lib/runtime/ffi/internal/internal/runtime.h | 6 + lib/runtime/ffi/src/runtime.cc | 4 +- lib/runtime/ffi/src/runtime.h | 14 + .../{src => include/runtime}/mapping_id_t.h | 0 lib/utils/ffi/CMakeLists.txt | 17 + lib/utils/ffi/include/flexflow/utils.h | 24 +- lib/utils/ffi/internal/internal/error.h | 53 ++ lib/utils/ffi/internal/internal/opaque.h | 111 ++++ lib/utils/ffi/src/error.cc | 5 + lib/utils/include/utils/bidict.h | 34 +- lib/utils/include/utils/ffi/opaque.h | 72 --- lib/utils/include/utils/type_traits.h | 8 + lib/utils/include/utils/type_traits_core.h | 26 + 33 files changed, 1696 insertions(+), 415 deletions(-) create mode 100644 lib/compiler/ffi/internal/internal/compiler.h create mode 100644 lib/op-attrs/ffi/internal/internal/op-attrs.h create mode 100644 lib/pcg/ffi/internal/internal/pcg.h create mode 100644 lib/pcg/ffi/src/pcg.h create mode 100644 lib/runtime/ffi/internal/internal/runtime.h create mode 100644 lib/runtime/ffi/src/runtime.h rename lib/runtime/{src => include/runtime}/mapping_id_t.h (100%) create mode 100644 lib/utils/ffi/internal/internal/error.h create mode 100644 lib/utils/ffi/internal/internal/opaque.h create mode 100644 lib/utils/ffi/src/error.cc delete mode 100644 lib/utils/include/utils/ffi/opaque.h 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..dec1f44188 100644 --- a/lib/compiler/ffi/include/flexflow/compiler.h +++ b/lib/compiler/ffi/include/flexflow/compiler.h @@ -7,10 +7,6 @@ FLEXFLOW_FFI_BEGIN() -typedef enum { - FLEXFLOW_COMPILER_STATUS_OK, - FLEXFLOW_COMPILER_ERROR_UNKNOWN -} flexflow_compiler_error_t; typedef enum { FLEXFLOW_SEARCH_ALGORITHM_DATA_PARALLEL @@ -18,7 +14,26 @@ typedef enum { 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..7f47a6c527 100644 --- a/lib/ffi/include/flexflow/flexflow.h +++ b/lib/ffi/include/flexflow/flexflow.h @@ -5,13 +5,21 @@ #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 +27,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..c21d26d93f 100644 --- a/lib/ffi/src/ffi.cc +++ b/lib/ffi/src/ffi.cc @@ -1 +1,74 @@ #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..9b31050fd9 100644 --- a/lib/op-attrs/ffi/include/flexflow/op-attrs.h +++ b/lib/op-attrs/ffi/include/flexflow/op-attrs.h @@ -5,6 +5,29 @@ 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 +52,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 +61,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..395f91e9aa --- /dev/null +++ b/lib/op-attrs/ffi/internal/internal/op-attrs.h @@ -0,0 +1,37 @@ +#ifndef _FLEXFLOW_OPATTRS_FFI_INTERNAL_INTERNAL_OPATTRS_H +#define _FLEXFLOW_OPATTRS_FFI_INTERNAL_INTERNAL_OPATTRS_H + +#include "flexflow/op-attrs.h" +#include "op-attrs/ops/linear.h" +#include "op-attrs/param_sync.h" +#include "op-attrs/datatype.h" +#include "op-attrs/ops/pool_2d.h" +#include "op-attrs/activation.h" +#include "op-attrs/datatype.h" +#include "op-attrs/ops/embedding.h" +#include "internal/opaque.h" +#include "op-attrs/op.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..19dd5d748a 100644 --- a/lib/op-attrs/ffi/src/op-attrs.cc +++ b/lib/op-attrs/ffi/src/op-attrs.cc @@ -1 +1,191 @@ #include "flexflow/op-attrs.h" +#include "internal/op-attrs.h" +#include "internal/error.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); + +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(); +} + +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/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..e1e06fd629 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,41 @@ 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); + +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_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, @@ -68,8 +82,7 @@ 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_destroy(flexflow_tensor_t); flexflow_error_t flexflow_computation_graph_add_op_exp(flexflow_computation_graph_t, @@ -120,14 +133,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 +216,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 +231,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, 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..6a89f7649d 100644 --- a/lib/pcg/ffi/src/pcg.cc +++ b/lib/pcg/ffi/src/pcg.cc @@ -1,2 +1,592 @@ +#include "pcg.h" #include "flexflow/pcg.h" -#include "pcg/model_compilation.h" +#include "flexflow/utils.h" +#include "internal/opaque.h" +#include "internal/error.h" +#include "pcg/computation_graph.h" +#include "pcg/computation_graph_builder.h" +#include "pcg/file_format/v1/graphs.h" +#include "utils/exception.h" +#include "internal/op-attrs.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_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(); +} + +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) { + try { + *out = new_opaque(); + return status_ok(); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } +} + +flexflow_error_t flexflow_computation_graph_destroy(flexflow_computation_graph_t opaque) { + try { + return status_ok(); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } +} + +template +std::string get_v1_string(T const &t) { + return json{to_v1(t)}.dump(); +} + +template +T from_v1_string(char const *s) { + auto v1 = json::parse(s).template get>()))>(); + return new_opaque(from_v1(v1)); +} + +flexflow_error_t flexflow_computation_graph_serialize_to_buffer(flexflow_computation_graph_t opaque, char **out) { + try { + 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'); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_computation_graph_deserialize_from_buffer(char *buf, flexflow_computation_graph_t *out) { + try { + *out = from_v1_string(buf); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_computation_graph_serialize_to_file(flexflow_computation_graph_t cg, FILE *f) { + try { + if (f == nullptr) { + return make_pcg_error(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()) { + return make_pcg_error(FLEXFLOW_PCG_STATUS_FILE_WRITE_FAILED); + } + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_computation_graph_deserialize_from_file(FILE *f, flexflow_computation_graph_t *out) { + size_t constexpr BUF_SIZE = 256; + + try { + 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 { + return make_pcg_error(FLEXFLOW_PCG_STATUS_FILE_READ_FAILED); + } + } + } + + *out = from_v1_string(oss.str().c_str()); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_create(flexflow_computation_graph_t cg_handle, + int num_dims, + int *dims, + flexflow_datatype_t datatype, + bool create_grad, + flexflow_tensor_t *out) { + try { + ComputationGraph *cg = unwrap_opaque(cg_handle); + TensorDims ordered_dims{dims, dims + num_dims}; + TensorShape shape = { ordered_dims, to_internal(datatype) }; + Tensor result = create_tensor(*cg, shape, create_grad); + *out = new_opaque(result); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_get_create_grad(flexflow_tensor_t opaque, bool *out) { + try { + Tensor const *t = unwrap_opaque(opaque); + *out = t->create_gradients; + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_get_initializer(flexflow_tensor_t opaque, flexflow_initializer_t *out) { + try { + Tensor const *t = unwrap_opaque(opaque); + *out = new_opaque(t->initializer); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_get_sync_type(flexflow_tensor_t opaque, flexflow_param_sync_t *out) { + try { + Tensor const *t = unwrap_opaque(opaque); + *out = to_external(t->sync_type); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_get_datatype(flexflow_tensor_t opaque, flexflow_datatype_t *out) { + try { + Tensor const *t = unwrap_opaque(opaque); + *out = to_external(t->data_type); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_get_num_dims(flexflow_tensor_t opaque, int *out) { + try { + Tensor const *t = unwrap_opaque(opaque); + *out = t->num_dims(); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_get_dims(flexflow_tensor_t opaque, int *out) { + try { + Tensor const *t = unwrap_opaque(opaque); + for (int i = 0; i < t->num_dims(); i++) { + out[i] = t->dims.at(ff_dim_t(i)); + } + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_tensor_destroy(flexflow_tensor_t opaque) { + try { + delete_opaque(opaque); + } catch (flexflow_utils_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +optional maybe_string(char *); +using BinaryFunc = std::function const &)>; +using UnaryFunc = std::function)>; +using ScalarUnaryFunc = std::function)>; + +#define BINARY_OP(func) \ + [](ComputationGraph &cg, Tensor const &lhs, Tensor const &rhs, optional const &name) -> Tensor { return func(cg, lhs, rhs, name); } + +#define UNARY_OP(func) \ + [](ComputationGraph &cg, Tensor const &t, optional const &name) -> Tensor { return func(cg, t, name); } + +#define SCALAR_UNARY_OP(func) \ + [](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 opaque_cg, + flexflow_tensor_t opaque_lhs, + flexflow_tensor_t opaque_rhs, + flexflow_tensor_t *out, + char *name) { + try { + ComputationGraph *cg = unwrap_opaque(opaque_cg); + Tensor const *lhs = unwrap_opaque(opaque_lhs); + Tensor const *rhs = unwrap_opaque(opaque_rhs); + Tensor output = f(*cg, *lhs, *rhs, maybe_string(name)); + *out = new_opaque(output); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t add_op(UnaryFunc const &f, + flexflow_computation_graph_t opaque_cg, + flexflow_tensor_t opaque_input, + flexflow_tensor_t *out, + char *name) { + try { + ComputationGraph *cg = unwrap_opaque(opaque_cg); + Tensor const *input = unwrap_opaque(opaque_input); + Tensor output = f(*cg, *input, maybe_string(name)); + *out = new_opaque(output); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +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 *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, 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(UNARY_OP(exp), 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(add), 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(subtract), 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(multiply), 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(divide), 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(max), 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(min), 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(UNARY_OP(rsqrt), 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_UNARY_OP(scalar_multiply), 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_UNARY_OP(scalar_add), 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_UNARY_OP(scalar_sub), 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_UNARY_OP(scalar_truediv), 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(UNARY_OP(sin), 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(UNARY_OP(cos), 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(UNARY_OP(relu), 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(UNARY_OP(identity), 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(UNARY_OP(gelu), 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(UNARY_OP(sigmoid), 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(UNARY_OP(tanh), 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(UNARY_OP(tanh), 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) { + try { + Tensor output = 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)); + *out = new_opaque(output); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +flexflow_error_t flexflow_computation_graph_add_op_dropout(flexflow_computation_graph_t opaque_cg, + flexflow_tensor_t opaque_input, + float rate, + flexflow_tensor_t *out, + unsigned long long seed, + char *name) { + UnaryFunc f = [&](ComputationGraph &cg, + Tensor const &input, + optional const &name) -> Tensor { + return dropout(cg, input, rate, seed, name); + }; + return add_op(f, opaque_cg, opaque_input, out, name); +} + +flexflow_error_t flexflow_computation_graph_add_op_embedding(flexflow_computation_graph_t opaque_cg, + flexflow_tensor_t opaque_input, + int num_entries, + int out_dim, + flexflow_aggregate_op_t aggr_op, + flexflow_tensor_t *out, + flexflow_datatype_t output_type, + flexflow_initializer_t initializer, + char *name) { + UnaryFunc f = [&](ComputationGraph &cg, + Tensor const &input, + optional const &name) -> Tensor { + return embedding(cg, input, num_entries, out_dim, to_internal(aggr_op), to_internal(output_type), c_deref_opaque(initializer), name); + }; + return add_op(f, opaque_cg, opaque_input, out, 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) { + try { + std::vector outputs = gather(deref_opaque(cg), + c_deref_opaque(input), + c_deref_opaque(index), + ff_dim_t(dim), + maybe_string(name)); + *out = new_opaque(outputs); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +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) { + try { + std::vector outputs = group_by(deref_opaque(cg), + c_deref_opaque(data), + c_deref_opaque(assign), + n, + alpha, + maybe_string(name)); + *out = new_opaque(outputs); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + +/* 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..a0d26a7afb --- /dev/null +++ b/lib/pcg/ffi/src/pcg.h @@ -0,0 +1,33 @@ +#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/machine_specification.h" +#include "pcg/parallel_computation_graph.h" +#include "pcg/operator.h" +#include "pcg/model_compilation.h" + +using namespace FlexFlow; + +REGISTER_OPAQUE(flexflow_computation_graph_t, ComputationGraph); +REGISTER_OPAQUE(flexflow_parallel_computation_graph_t, ParallelComputationGraph); +REGISTER_OPAQUE(flexflow_operator_t, Operator); +REGISTER_OPAQUE(flexflow_parallel_tensor_t, ParallelTensor); +REGISTER_OPAQUE(flexflow_layer_t, Layer); +REGISTER_OPAQUE(flexflow_tensor_t, Tensor); +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_builder.h b/lib/pcg/include/pcg/computation_graph_builder.h index 7f01439712..9d010cd232 100644 --- a/lib/pcg/include/pcg/computation_graph_builder.h +++ b/lib/pcg/include/pcg/computation_graph_builder.h @@ -4,284 +4,291 @@ #include "computation_graph.h" 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, +// +// C++ APIs for constructing models +// Add an exp layer +Tensor exp(ComputationGraph &, Tensor const &, optional const &name = nullopt); +// Add an add layer +Tensor add(ComputationGraph &, + Tensor const &x, + Tensor const &y, + optional const &name = nullopt); +// Add a subtract layer +Tensor subtract(ComputationGraph &, + 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, - 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, +// Add a multiply layer +Tensor multiply(ComputationGraph &, + Tensor const &x, + Tensor const &y, 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, +// Add a divide layer +Tensor divide(ComputationGraph &, + Tensor const &x, + Tensor const &y, + optional const &name = nullopt); +// Add a max layer +Tensor max(ComputationGraph &, + Tensor const &x, + Tensor const &y, + optional const &name = nullopt); +// Add a min layer +Tensor min(ComputationGraph &, + Tensor const &x, + Tensor const &y, + optional const &name = nullopt); +// Add a rsqrt layer +Tensor rsqrt(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +// Add a pow layer +Tensor pow(ComputationGraph &, + Tensor const &x, + float exponent, + optional const &name = nullopt); +// Add a scalar multiply layer +Tensor scalar_multiply(ComputationGraph &, + Tensor const &x, + float scalar, + optional const &name = nullopt); +Tensor scalar_add(ComputationGraph &, + Tensor const &x, + float scalar, + optional const &name = nullopt); +Tensor scalar_sub(ComputationGraph &, + Tensor const &lhs, + float rhs, + optional const &name = nullopt); +Tensor scalar_truediv(ComputationGraph &, + Tensor const &numerator, + float denominator, 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, +// Add a sin layer +Tensor sin(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +// Add a cos layer +Tensor cos(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +// Add an activation layer +Tensor relu(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +Tensor identity(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +Tensor gelu(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +Tensor sigmoid(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +Tensor tanh(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +Tensor elu(ComputationGraph &, Tensor const &x, optional const &name = nullopt); +// Add a 2D convolutional layer +Tensor conv2d(ComputationGraph &, + 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 const &kernel_initializer = nullopt, + optional const &bias_initializer = nullopt, + optional const &kernel_regularizer = nullopt, + optional const &name = nullopt); +// Add a dropout layer +Tensor dropout(ComputationGraph &, + Tensor const &input, + float rate, + unsigned long long seed = 0, optional const &name = nullopt); - // Add a cast layer - Tensor cast(Tensor const &input, - DataType dtype, +// Add an embedding layer +Tensor embedding(ComputationGraph &, + Tensor const &input, + int num_entries, + int outDim, + AggregateOp aggr, + DataType dtype = DataType::FLOAT, + optional const &kernel_initializer = nullopt, + optional const &name = nullopt); +// Add a gather layer +std::vector gather(ComputationGraph &, + Tensor const &input, + Tensor const &index, + ff_dim_t dim, + optional const &name = nullopt); +// Add a group_by layer +std::vector group_by(ComputationGraph &, + Tensor const &data, + Tensor const &assign, + int n, + float alpha, 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, +// Add a cache layer +Tensor cache(Tensor const &input, + int num_batches, + std::function + score_f = {}, 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, +// 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, + 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 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, +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); - 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, - 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); +// 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, + optional const &name = nullopt); +Tensor create_tensor(ComputationGraph &, 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; +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; +Tensor at(MultiDiEdge const &) const; +Layer at(Node const &) const; private: - Tensor broadcast(Tensor const &, TensorShape const &); +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); +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 &); +Tensor as_type(Tensor const &, DataType, std::string const &); - TensorShape get_broadcast_target_shape(std::vector const &); +TensorShape get_broadcast_target_shape(std::vector const &); - Tensor element_binary(OperatorType, - Tensor const &lhs, - Tensor const &rhs, - optional const &name = nullopt); +Tensor element_binary(OperatorType, + Tensor const &lhs, + Tensor const &rhs, + 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, - optional const &name = nullopt); - Tensor - element_unary(variant const &, - Tensor const &input, - 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, + optional const &name = nullopt); +Tensor + element_unary(variant const &, + Tensor const &input, + optional const &name = nullopt); -public: - ComputationGraph computation_graph; -}; } // namespace FlexFlow 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/tensor.h b/lib/pcg/include/pcg/tensor.h index cb79be245a..60ed3ec49e 100644 --- a/lib/pcg/include/pcg/tensor.h +++ b/lib/pcg/include/pcg/tensor.h @@ -9,23 +9,16 @@ 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; operator TensorShape() const; - public: TensorDims dims; DataType data_type; - req> initializer; - req create_gradients; + optional initializer; + bool create_gradients; req> sync_type; }; FF_VISITABLE_STRUCT( 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..eec6475844 100644 --- a/lib/runtime/ffi/include/flexflow/runtime.h +++ b/lib/runtime/ffi/include/flexflow/runtime.h @@ -1,26 +1,40 @@ #ifndef _FLEXFLOW_RUNTIME_INCLUDE_RUNTIME_FFI_RUNTIME_H #define _FLEXFLOW_RUNTIME_INCLUDE_RUNTIME_FFI_RUNTIME_H +#include "flexflow/utils.h" #include "flexflow/compiler.h" #include "flexflow/op-attrs.h" #include "flexflow/pcg.h" -#include "flexflow/utils.h" #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 +58,74 @@ 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, +flexflow_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_set_config(flexflow_config_t); +flexflow_error_t flexflow_get_config(flexflow_config_t *); -flexflow_runtime_error_t flexflow_model_config_parse_argv( +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_error_t flexflow_start_training(flexflow_model_compilation_result_t, flexflow_model_training_instance_t *); -flexflow_runtime_error_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_error_t flexflow_stop_training(flexflow_model_training_instance_t); -flexflow_runtime_error_t +flexflow_error_t flexflow_get_tensor_float(flexflow_model_training_instance_t, flexflow_tensor_t, float *data, bool get_gradients); -flexflow_runtime_error_t +flexflow_error_t flexflow_get_tensor_double(flexflow_model_training_instance_t, flexflow_tensor_t, float *data, bool get_gradients); -flexflow_runtime_error_t +flexflow_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_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_error_t flexflow_set_tensor_int( flexflow_model_training_instance_t, flexflow_tensor_t, int32_t *data); -FLEXFLOW_FFI_END() +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..46cdcd766a 100644 --- a/lib/runtime/ffi/src/runtime.cc +++ b/lib/runtime/ffi/src/runtime.cc @@ -1,10 +1,12 @@ #include "flexflow/runtime.h" #include "runtime/model_training_instance.h" +#include "internal/opaque.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 +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/opaque.h b/lib/utils/ffi/internal/internal/opaque.h new file mode 100644 index 0000000000..7ee5bbff9c --- /dev/null +++ b/lib/utils/ffi/internal/internal/opaque.h @@ -0,0 +1,111 @@ +#ifndef _FLEXFLOW_UTILS_INCLUDE_UTILS_FFI_OPAQUE_H +#define _FLEXFLOW_UTILS_INCLUDE_UTILS_FFI_OPAQUE_H + +#include "utils/expected.h" +#include +#include +#include "flexflow/utils.h" +#include "error.h" + +template +struct opaque_to_underlying; + +template +struct underlying_to_opaque; + +template +struct internal_to_external; + +template +struct external_to_internal; + +template +struct enum_mapping; + +template +using opaque_to_underlying_t = typename opaque_to_underlying::type; + +template +using underlying_to_opaque_t = typename underlying_to_opaque::type; + +template +using internal_to_external_t = typename internal_to_external::type; + +template +using external_to_internal_t = typename external_to_internal::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); +} + +#define REGISTER_OPAQUE(OPAQUE, UNDERLYING) \ + template <> \ + struct opaque_to_underlying { \ + using type = UNDERLYING; \ + }; \ + template <> \ + struct underlying_to_opaque { \ + using type = OPAQUE; \ + }; + +#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 +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 +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/include/utils/bidict.h b/lib/utils/include/utils/bidict.h index 870afd0448..6432c97cad 100644 --- a/lib/utils/include/utils/bidict.h +++ b/lib/utils/include/utils/bidict.h @@ -2,13 +2,25 @@ #define _FLEXFLOW_UTILS_BIDICT_H #include +#include "optional.h" 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 +66,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/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/type_traits.h b/lib/utils/include/utils/type_traits.h index 7b345533fe..45389a03e3 100644 --- a/lib/utils/include/utils/type_traits.h +++ b/lib/utils/include/utils/type_traits.h @@ -50,6 +50,14 @@ 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..fb5fd62a0c 100644 --- a/lib/utils/include/utils/type_traits_core.h +++ b/lib/utils/include/utils/type_traits_core.h @@ -2,6 +2,8 @@ #define _FLEXFLOW_UTILS_INCLUDE_UTILS_TYPE_TRAITS_CORE_H #include +#include +#include namespace FlexFlow { template @@ -91,6 +93,30 @@ 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, ""); From afc315e8dfd3b87f3cf3abbca11671f86856439c Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Thu, 3 Aug 2023 21:08:26 -0700 Subject: [PATCH 2/4] Centralize error handling wrapper --- lib/op-attrs/ffi/src/op-attrs.cc | 19 +- lib/pcg/ffi/include/flexflow/pcg.h | 22 +- lib/pcg/ffi/src/pcg.cc | 601 ++++++++++++------ lib/pcg/include/pcg/computation_graph.h | 2 + .../include/pcg/computation_graph_builder.h | 78 ++- lib/utils/ffi/internal/internal/enums.h | 54 ++ .../ffi/internal/internal/error_handling.h | 30 + lib/utils/ffi/internal/internal/opaque.h | 45 +- lib/utils/include/utils/containers.h | 9 + 9 files changed, 589 insertions(+), 271 deletions(-) create mode 100644 lib/utils/ffi/internal/internal/enums.h create mode 100644 lib/utils/ffi/internal/internal/error_handling.h diff --git a/lib/op-attrs/ffi/src/op-attrs.cc b/lib/op-attrs/ffi/src/op-attrs.cc index 19dd5d748a..18854a8e9c 100644 --- a/lib/op-attrs/ffi/src/op-attrs.cc +++ b/lib/op-attrs/ffi/src/op-attrs.cc @@ -4,6 +4,7 @@ #include "op-attrs/op.h" #include "op-attrs/ops/embedding.h" #include "utils/bidict.h" +#include "internal/enums.h" flexflow_utils_exception_t make_opattrs_exception(flexflow_opattrs_error_code_t); @@ -154,24 +155,6 @@ REGISTER_FFI_ENUM( flexflow_error_t make_opattrs_error(flexflow_opattrs_error_code_t); -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(); -} - 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); } diff --git a/lib/pcg/ffi/include/flexflow/pcg.h b/lib/pcg/ffi/include/flexflow/pcg.h index e1e06fd629..9b738512f6 100644 --- a/lib/pcg/ffi/include/flexflow/pcg.h +++ b/lib/pcg/ffi/include/flexflow/pcg.h @@ -27,12 +27,15 @@ 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; @@ -81,7 +84,7 @@ 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_get_dims(flexflow_tensor_t, flexflow_int_list_t *out); flexflow_error_t flexflow_tensor_destroy(flexflow_tensor_t); flexflow_error_t @@ -268,9 +271,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, @@ -297,6 +300,7 @@ flexflow_error_t flexflow_tensor_t, flexflow_tensor_t *out, int *axes, + int num_axes, bool elementwise_affine, float eps, char *name = NULL); @@ -343,6 +347,7 @@ flexflow_error_t flexflow_tensor_t, flexflow_tensor_t *out, int *dims, + int num_dims, bool keepdims, char *name = NULL); flexflow_error_t @@ -358,8 +363,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 @@ -376,13 +382,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 @@ -390,6 +399,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, @@ -400,15 +410,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, @@ -416,6 +427,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/src/pcg.cc b/lib/pcg/ffi/src/pcg.cc index 6a89f7649d..431f500d9a 100644 --- a/lib/pcg/ffi/src/pcg.cc +++ b/lib/pcg/ffi/src/pcg.cc @@ -16,6 +16,7 @@ flexflow_error_t flexflow_pcg_error_wrap(flexflow_pcg_error_t 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) { @@ -27,6 +28,15 @@ flexflow_error_t flexflow_pcg_error_is_ok(flexflow_pcg_error_t, bool *out) { 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)); @@ -60,20 +70,11 @@ flexflow_error_t flexflow_pcg_error_destroy(flexflow_pcg_error_t) { } flexflow_error_t flexflow_computation_graph_create(flexflow_computation_graph_t *out) { - try { - *out = new_opaque(); - return status_ok(); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } + return handle_errors(out, [&] { return ComputationGraph{}; }); } flexflow_error_t flexflow_computation_graph_destroy(flexflow_computation_graph_t opaque) { - try { - return status_ok(); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } + return handle_errors([&] { }); } template @@ -83,12 +84,11 @@ std::string get_v1_string(T const &t) { template T from_v1_string(char const *s) { - auto v1 = json::parse(s).template get>()))>(); - return new_opaque(from_v1(v1)); + return json::parse(s).template get>()))>(); } flexflow_error_t flexflow_computation_graph_serialize_to_buffer(flexflow_computation_graph_t opaque, char **out) { - try { + return handle_errors([&] { ComputationGraph const *cg = unwrap_opaque(opaque); std::string json_str = get_v1_string(*cg); @@ -96,46 +96,32 @@ flexflow_error_t flexflow_computation_graph_serialize_to_buffer(flexflow_computa *out = new char[json_str.size()+1]; strncpy(*out, json_str.c_str(), json_str.size()+1); assert((*out)[json_str.size()] == '\x00'); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + }); } flexflow_error_t flexflow_computation_graph_deserialize_from_buffer(char *buf, flexflow_computation_graph_t *out) { - try { - *out = from_v1_string(buf); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + 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) { - try { + return handle_errors([&] { if (f == nullptr) { - return make_pcg_error(FLEXFLOW_PCG_STATUS_INVALID_FILE_PTR); + 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()) { - return make_pcg_error(FLEXFLOW_PCG_STATUS_FILE_WRITE_FAILED); + throw make_pcg_exception(FLEXFLOW_PCG_STATUS_FILE_WRITE_FAILED); } - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + }); } flexflow_error_t flexflow_computation_graph_deserialize_from_file(FILE *f, flexflow_computation_graph_t *out) { - size_t constexpr BUF_SIZE = 256; + return handle_errors(out, [&] { + size_t constexpr BUF_SIZE = 256; - try { char buf[BUF_SIZE]; std::ostringstream oss; @@ -146,114 +132,54 @@ flexflow_error_t flexflow_computation_graph_deserialize_from_file(FILE *f, flexf if (feof(f)) { break; } else { - return make_pcg_error(FLEXFLOW_PCG_STATUS_FILE_READ_FAILED); + throw make_pcg_exception(FLEXFLOW_PCG_STATUS_FILE_READ_FAILED); } } } - *out = from_v1_string(oss.str().c_str()); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + return from_v1_string(oss.str().c_str()); + }); } -flexflow_error_t flexflow_tensor_create(flexflow_computation_graph_t cg_handle, +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) { - try { - ComputationGraph *cg = unwrap_opaque(cg_handle); - TensorDims ordered_dims{dims, dims + num_dims}; + return handle_errors(out, [&] { + TensorDims ordered_dims{make_vector(dims, num_dims)}; TensorShape shape = { ordered_dims, to_internal(datatype) }; - Tensor result = create_tensor(*cg, shape, create_grad); - *out = new_opaque(result); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + return create_tensor(deref_opaque(cg), shape, create_grad); + }); } flexflow_error_t flexflow_tensor_get_create_grad(flexflow_tensor_t opaque, bool *out) { - try { - Tensor const *t = unwrap_opaque(opaque); - *out = t->create_gradients; - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + 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) { - try { - Tensor const *t = unwrap_opaque(opaque); - *out = new_opaque(t->initializer); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + 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) { - try { - Tensor const *t = unwrap_opaque(opaque); - *out = to_external(t->sync_type); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + 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) { - try { - Tensor const *t = unwrap_opaque(opaque); - *out = to_external(t->data_type); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + 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) { - try { - Tensor const *t = unwrap_opaque(opaque); - *out = t->num_dims(); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + return handle_errors(out, [&] { return c_deref_opaque(opaque).num_dims(); }); } flexflow_error_t flexflow_tensor_get_dims(flexflow_tensor_t opaque, int *out) { - try { - Tensor const *t = unwrap_opaque(opaque); - for (int i = 0; i < t->num_dims(); i++) { - out[i] = t->dims.at(ff_dim_t(i)); - } - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + return handle_errors(out, [&] { return c_deref_opaque(opaque).dims; }); } flexflow_error_t flexflow_tensor_destroy(flexflow_tensor_t opaque) { - try { - delete_opaque(opaque); - } catch (flexflow_utils_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + return handle_errors([&] { delete_opaque(opaque); }); } optional maybe_string(char *); @@ -272,39 +198,25 @@ using ScalarUnaryFunc = std::function(output); - } catch (flexflow_ffi_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + 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 opaque_cg, - flexflow_tensor_t opaque_input, + flexflow_computation_graph_t cg, + flexflow_tensor_t input, flexflow_tensor_t *out, char *name) { - try { - ComputationGraph *cg = unwrap_opaque(opaque_cg); - Tensor const *input = unwrap_opaque(opaque_input); - Tensor output = f(*cg, *input, maybe_string(name)); - *out = new_opaque(output); - } catch (flexflow_ffi_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + return handle_errors(out, [&] { return f(deref_opaque(cg), + c_deref_opaque(input), + maybe_string(name)); }); } flexflow_error_t add_op(ScalarUnaryFunc const &f, @@ -312,13 +224,13 @@ flexflow_error_t add_op(ScalarUnaryFunc const &f, flexflow_tensor_t opaque_input, float scalar, flexflow_tensor_t *out, - char *name) { + 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, name); + return add_op(ff, opaque_cg, opaque_input, out, raw_name); } @@ -491,8 +403,7 @@ flexflow_error_t flexflow_computation_graph_add_op_conv2d( flexflow_initializer_t bias_initializer, flexflow_regularizer_attrs_t kernel_regularizer, char *name) { - try { - Tensor output = conv2d(deref_opaque(cg), + return handle_errors(out, [&] { return conv2d(deref_opaque(cg), c_deref_opaque(input), outChannels, kernelH, @@ -507,44 +418,41 @@ flexflow_error_t flexflow_computation_graph_add_op_conv2d( c_deref_opaque(kernel_initializer), c_deref_opaque(bias_initializer), c_deref_opaque(kernel_regularizer), - maybe_string(name)); - *out = new_opaque(output); - } catch (flexflow_ffi_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + maybe_string(name)); }); } -flexflow_error_t flexflow_computation_graph_add_op_dropout(flexflow_computation_graph_t opaque_cg, - flexflow_tensor_t opaque_input, - float rate, +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) { - UnaryFunc f = [&](ComputationGraph &cg, - Tensor const &input, - optional const &name) -> Tensor { - return dropout(cg, input, rate, seed, name); - }; - return add_op(f, opaque_cg, opaque_input, out, 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 opaque_cg, - flexflow_tensor_t opaque_input, +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_tensor_t *out, flexflow_datatype_t output_type, flexflow_initializer_t initializer, char *name) { - UnaryFunc f = [&](ComputationGraph &cg, - Tensor const &input, - optional const &name) -> Tensor { - return embedding(cg, input, num_entries, out_dim, to_internal(aggr_op), to_internal(output_type), c_deref_opaque(initializer), name); - }; - return add_op(f, opaque_cg, opaque_input, out, 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, @@ -553,18 +461,12 @@ flexflow_error_t flexflow_computation_graph_add_op_gather(flexflow_computation_g int dim, flexflow_tensor_list_t *out, char *name) { - try { - std::vector outputs = gather(deref_opaque(cg), + return handle_errors(out, + [&] { return gather(deref_opaque(cg), c_deref_opaque(input), c_deref_opaque(index), ff_dim_t(dim), - maybe_string(name)); - *out = new_opaque(outputs); - } catch (flexflow_ffi_exception_t const &e) { - return to_error(e); - } - - return status_ok(); + maybe_string(name)); }); } flexflow_error_t flexflow_computation_graph_add_op_group_by(flexflow_computation_graph_t cg, @@ -574,19 +476,348 @@ flexflow_error_t flexflow_computation_graph_add_op_group_by(flexflow_computation float alpha, flexflow_tensor_list_t *out, char *name) { - try { - std::vector outputs = group_by(deref_opaque(cg), + return handle_errors(out, + [&] { return group_by(deref_opaque(cg), c_deref_opaque(data), c_deref_opaque(assign), n, alpha, - maybe_string(name)); - *out = new_opaque(outputs); - } catch (flexflow_ffi_exception_t const &e) { - return to_error(e); - } + maybe_string(name)); }); +} - return status_ok(); +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/include/pcg/computation_graph.h b/lib/pcg/include/pcg/computation_graph.h index 11dad70356..5258242bf1 100644 --- a/lib/pcg/include/pcg/computation_graph.h +++ b/lib/pcg/include/pcg/computation_graph.h @@ -13,6 +13,8 @@ namespace FlexFlow { struct ComputationGraph : public strong_typedef> { + ComputationGraph(); + using strong_typedef::strong_typedef; }; diff --git a/lib/pcg/include/pcg/computation_graph_builder.h b/lib/pcg/include/pcg/computation_graph_builder.h index 9d010cd232..ebe65e12e0 100644 --- a/lib/pcg/include/pcg/computation_graph_builder.h +++ b/lib/pcg/include/pcg/computation_graph_builder.h @@ -119,27 +119,29 @@ std::vector group_by(ComputationGraph &, float alpha, optional const &name = nullopt); // Add a cache layer -Tensor cache(Tensor const &input, +Tensor cache(ComputationGraph &, + Tensor const &input, int num_batches, std::function score_f = {}, optional const &name = nullopt); // Add aggregate layer -Tensor aggregate(Tensor const &gate_preds, +Tensor aggregate(ComputationGraph &, + 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); + optional const &name = nullopt); // Add aggregate_spec layer -Tensor aggregate_spec(std::vector const &inputs, - int n, +Tensor aggregate_spec(ComputationGraph &, + std::vector const &inputs, float lambda_bal, optional const &name = nullopt); // Add a 2D pooling layer -Tensor pool2d(Tensor const &input, +Tensor pool2d(ComputationGraph &, + Tensor const &input, int kernelH, int kernelW, int strideH, @@ -149,81 +151,95 @@ Tensor pool2d(Tensor const &input, PoolOp type = PoolOp::MAX, optional const &activation = nullopt, optional const &name = nullopt); -Tensor layer_norm(Tensor const &input, +Tensor layer_norm(ComputationGraph &, + Tensor const &input, std::vector const &axes, bool elementwise_affine, float eps, optional const &name = nullopt); -Tensor batch_norm(Tensor const &input, +Tensor batch_norm(ComputationGraph &, + Tensor const &input, bool relu = true, optional const &name = nullopt); -Tensor batch_matmul(Tensor const &A, +Tensor batch_matmul(ComputationGraph &, + 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, +Tensor dense(ComputationGraph &, + 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 &kernel_initializer = nullopt, + optional const &bias_initializer = nullopt, optional const &name = nullopt); // Add a cast layer -Tensor cast(Tensor const &input, +Tensor cast(ComputationGraph &, + Tensor const &input, DataType dtype, optional const &name = nullopt); // Add a concat layer -Tensor concat(int n, +Tensor concat(ComputationGraph &, std::vector const &tensors, int axis, optional const &name = nullopt); // Add a mean layer -Tensor mean(Tensor const &input, +Tensor mean(ComputationGraph &, + Tensor const &input, std::vector const &dims, bool keepdims, - char const *name); + optional const &name = nullopt); // Add a moe layer (wrapping topk, group_by and aggregate operators) -Tensor moe(Tensor const &input, +Tensor moe(ComputationGraph &, + Tensor const &input, int num_exp, int num_select, int expert_hidden_size, float alpha, - float lambda); + float lambda, + optional const &name = nullopt); // Add a split layer -void split(Tensor const &input, - Tensor *outputs, +std::vector split(ComputationGraph &, + Tensor const &input, std::vector const &split, int axis, optional const &name = nullopt); // Add a flat layer -Tensor flat(Tensor const &input, optional const &name = nullopt); +Tensor flat(ComputationGraph &, Tensor const &input, optional const &name = nullopt); // Add a softmax layer -Tensor softmax(Tensor const &input, +Tensor softmax(ComputationGraph &, + Tensor const &input, int dim = -1, optional const &name = nullopt); // Create input tensors and constants -Tensor transpose(Tensor const &input, +Tensor transpose(ComputationGraph &, + Tensor const &input, std::vector const &perm, optional const &name = nullopt); -Tensor reduce_sum(Tensor const &input, +Tensor reduce_sum(ComputationGraph &, + Tensor const &input, std::vector const &axes, bool keepdims = false, optional const &name = nullopt); -Tensor reshape(Tensor const &input, +Tensor reshape(ComputationGraph &, + Tensor const &input, std::vector const &shape, optional const &name = nullopt); -Tensor reverse(Tensor const &input, +Tensor reverse(ComputationGraph &, + Tensor const &input, int axis, optional const &name = nullopt); -void top_k(Tensor const &input, - Tensor *outputs, +std::vector top_k(ComputationGraph &, + Tensor const &input, int k, bool sorted, optional const &name = nullopt); Tensor - multihead_attention(Tensor const &query, + multihead_attention(ComputationGraph &, + Tensor const &query, Tensor const &key, Tensor const &value, int embed_dim, @@ -234,7 +250,7 @@ Tensor bool bias = true, bool add_bias_kv = false, bool add_zero_attn = false, - optional initializer = nullopt, + optional const &initializer = nullopt, optional const &name = nullopt); Tensor create_tensor(ComputationGraph &, TensorShape const &, bool create_grad = true); Parameter create_weight(TensorShape const &, diff --git a/lib/utils/ffi/internal/internal/enums.h b/lib/utils/ffi/internal/internal/enums.h new file mode 100644 index 0000000000..2604c05a11 --- /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_handling.h b/lib/utils/ffi/internal/internal/error_handling.h new file mode 100644 index 0000000000..a25f75b939 --- /dev/null +++ b/lib/utils/ffi/internal/internal/error_handling.h @@ -0,0 +1,30 @@ +#ifndef _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ERROR_HANDLING_H +#define _FLEXFLOW_UTILS_FFI_INTERNAL_INTERNAL_ERROR_HANDLING_H + +#include "error.h" +#include "opaque.h" + +template +flexflow_error_t handle_errors(Opaque *out, std::function()> const &f) { + try { + *out = new_opaque(f()); + } catch (flexflow_ffi_exception_t const &e) { + return to_error(e); + } + + return status_ok(); +} + + +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(); +} + + +#endif diff --git a/lib/utils/ffi/internal/internal/opaque.h b/lib/utils/ffi/internal/internal/opaque.h index 7ee5bbff9c..78cf7e34a2 100644 --- a/lib/utils/ffi/internal/internal/opaque.h +++ b/lib/utils/ffi/internal/internal/opaque.h @@ -6,6 +6,8 @@ #include #include "flexflow/utils.h" #include "error.h" +#include +#include "utils/containers.h" template struct opaque_to_underlying; @@ -13,27 +15,12 @@ struct opaque_to_underlying; template struct underlying_to_opaque; -template -struct internal_to_external; - -template -struct external_to_internal; - -template -struct enum_mapping; - template using opaque_to_underlying_t = typename opaque_to_underlying::type; template using underlying_to_opaque_t = typename underlying_to_opaque::type; -template -using internal_to_external_t = typename internal_to_external::type; - -template -using external_to_internal_t = typename external_to_internal::type; - template opaque_to_underlying_t *unwrap_opaque(Opaque const &opaque) { if (opaque.impl == nullptr) { @@ -58,6 +45,12 @@ 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 { \ @@ -68,23 +61,6 @@ opaque_to_underlying_t &deref_opaque(Opaque const &opaque) { using type = OPAQUE; \ }; -#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 Opaque new_opaque(Args &&... args) { using Underlying = opaque_to_underlying_t; @@ -96,6 +72,11 @@ Opaque new_opaque(Args &&... args) { 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; diff --git a/lib/utils/include/utils/containers.h b/lib/utils/include/utils/containers.h index a0939f7038..568e6451a2 100644 --- a/lib/utils/include/utils/containers.h +++ b/lib/utils/include/utils/containers.h @@ -449,6 +449,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())) { From a4a677643c1d87826320fc8eb1e01b7750861b9e Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Thu, 17 Aug 2023 04:04:49 -0700 Subject: [PATCH 3/4] Format plus other stuff --- .clang-format-for-format-sh | 182 +++ examples/cpp/DLRM/strategies/dlrm_strategy.cc | 2 +- lib/compiler/ffi/include/flexflow/compiler.h | 14 +- lib/compiler/test/test_dp.cc | 2 +- lib/ffi/include/flexflow/flexflow.h | 11 +- lib/ffi/src/ffi.cc | 15 +- lib/kernels/include/kernels/reshape_kernels.h | 2 +- lib/kernels/src/hip/aggregate_kernels.cpp | 2 +- lib/kernels/src/hip/optimizer_kernel.cpp | 2 +- lib/op-attrs/ffi/include/flexflow/op-attrs.h | 7 +- lib/op-attrs/ffi/internal/internal/op-attrs.h | 13 +- lib/op-attrs/ffi/src/op-attrs.cc | 318 +++-- lib/op-attrs/include/op-attrs/ff_dim.h | 2 + lib/op-attrs/include/op-attrs/tensor_shape.h | 8 +- lib/pcg/ffi/include/flexflow/pcg.h | 20 +- lib/pcg/ffi/src/pcg.cc | 1251 +++++++++-------- lib/pcg/ffi/src/pcg.h | 58 +- lib/pcg/include/pcg/computation_graph.h | 40 +- .../include/pcg/computation_graph_builder.h | 609 ++++---- lib/pcg/include/pcg/layer.h | 19 +- lib/pcg/include/pcg/layer_guid_t.h | 17 + lib/pcg/include/pcg/parallel_tensor.h | 35 +- lib/pcg/include/pcg/tensor.h | 11 +- lib/pcg/include/pcg/tensor_guid_t.h | 14 +- lib/pcg/src/computation_graph_builder.cc | 528 ++++--- lib/runtime/ffi/include/flexflow/runtime.h | 77 +- lib/runtime/ffi/src/runtime.cc | 6 +- lib/runtime/ffi/src/runtime.h | 4 +- lib/runtime/src/create_grad.h | 10 - lib/runtime/src/tensor.h | 100 +- lib/runtime/test/src/test_op_task_spec.cc | 2 +- lib/runtime/test/src/test_serialization.cc | 2 +- lib/utils/ffi/include/flexflow/utils.h | 5 +- lib/utils/ffi/internal/internal/enums.h | 50 +- lib/utils/ffi/internal/internal/error.h | 34 +- .../ffi/internal/internal/error_handling.h | 26 +- lib/utils/ffi/internal/internal/opaque.h | 41 +- lib/utils/ffi/src/error_handling.cc | 13 + lib/utils/include/utils/bidict.h | 5 +- lib/utils/include/utils/containers.h | 4 +- lib/utils/include/utils/graph/algorithms.h | 4 + .../labelled_downward_open_interfaces.h | 3 +- .../utils/graph/labelled/open_algorithms.h | 35 +- .../include/utils/graph/labelled/open_views.h | 107 +- .../utils/graph/open_graph_interfaces.h | 34 +- lib/utils/include/utils/graph/query_set.h | 5 +- lib/utils/include/utils/strong_typedef.h | 138 +- lib/utils/include/utils/test_types.h | 114 +- lib/utils/include/utils/type_traits.h | 1 - lib/utils/include/utils/type_traits_core.h | 13 +- lib/utils/include/utils/variant.h | 12 +- lib/utils/include/utils/visitable.h | 11 +- lib/utils/src/graph/algorithms.cc | 20 +- python/flexflow_c.cc | 4 +- scripts/format.sh | 20 +- 55 files changed, 2377 insertions(+), 1705 deletions(-) create mode 100644 .clang-format-for-format-sh create mode 100644 lib/pcg/include/pcg/layer_guid_t.h delete mode 100644 lib/runtime/src/create_grad.h create mode 100644 lib/utils/ffi/src/error_handling.cc diff --git a/.clang-format-for-format-sh b/.clang-format-for-format-sh new file mode 100644 index 0000000000..17e9f8935d --- /dev/null +++ b/.clang-format-for-format-sh @@ -0,0 +1,182 @@ +--- +Language: Cpp +# BasedOnStyle: LLVM +AccessModifierOffset: -2 +AlignAfterOpenBracket: Align +AlignArrayOfStructures: None +AlignConsecutiveMacros: None +AlignConsecutiveAssignments: None +AlignConsecutiveBitFields: None +AlignConsecutiveDeclarations: None +AlignEscapedNewlines: Right +AlignOperands: Align +AlignTrailingComments: true +AllowAllArgumentsOnNextLine: true +AllowAllConstructorInitializersOnNextLine: true +AllowAllParametersOfDeclarationOnNextLine: true +AllowShortEnumsOnASingleLine: true +AllowShortBlocksOnASingleLine: Never +AllowShortCaseLabelsOnASingleLine: false +AllowShortFunctionsOnASingleLine: Empty +AllowShortLambdasOnASingleLine: All +AllowShortIfStatementsOnASingleLine: Never +AllowShortLoopsOnASingleLine: false +AlwaysBreakAfterDefinitionReturnType: None +AlwaysBreakAfterReturnType: None +AlwaysBreakBeforeMultilineStrings: false +AlwaysBreakTemplateDeclarations: Yes +AttributeMacros: + - __capability +BinPackArguments: false +BinPackParameters: false +BraceWrapping: + AfterCaseLabel: false + AfterClass: false + AfterControlStatement: Never + AfterEnum: false + AfterFunction: false + AfterNamespace: false + AfterObjCDeclaration: false + AfterStruct: false + AfterUnion: false + AfterExternBlock: false + BeforeCatch: false + BeforeElse: false + BeforeLambdaBody: false + BeforeWhile: false + IndentBraces: false + SplitEmptyFunction: true + SplitEmptyRecord: true + SplitEmptyNamespace: true +BreakBeforeBinaryOperators: None +BreakBeforeConceptDeclarations: true +BreakBeforeBraces: Attach +BreakBeforeInheritanceComma: false +BreakInheritanceList: BeforeColon +BreakBeforeTernaryOperators: true +BreakConstructorInitializersBeforeComma: false +BreakConstructorInitializers: BeforeColon +BreakAfterJavaFieldAnnotations: false +BreakStringLiterals: true +ColumnLimit: 80 +CommentPragmas: '^ IWYU pragma:' +CompactNamespaces: false +ConstructorInitializerAllOnOneLineOrOnePerLine: false +ConstructorInitializerIndentWidth: 4 +ContinuationIndentWidth: 4 +Cpp11BracedListStyle: true +DeriveLineEnding: true +DerivePointerAlignment: false +DisableFormat: false +EmptyLineAfterAccessModifier: Never +EmptyLineBeforeAccessModifier: LogicalBlock +ExperimentalAutoDetectBinPacking: false +FixNamespaceComments: true +ForEachMacros: + - foreach + - Q_FOREACH + - BOOST_FOREACH +IfMacros: + - KJ_IF_MAYBE +IncludeBlocks: Preserve +IncludeCategories: + - Regex: '^"(llvm|llvm-c|clang|clang-c)/' + Priority: 2 + SortPriority: 0 + CaseSensitive: false + - Regex: '^(<|"(gtest|gmock|isl|json)/)' + Priority: 3 + SortPriority: 0 + CaseSensitive: false + - Regex: '.*' + Priority: 1 + SortPriority: 0 + CaseSensitive: false +IncludeIsMainRegex: '(Test)?$' +IncludeIsMainSourceRegex: '' +IndentAccessModifiers: false +IndentCaseLabels: true +IndentCaseBlocks: false +IndentGotoLabels: true +IndentPPDirectives: None +IndentExternBlock: AfterExternBlock +IndentRequires: false +IndentWidth: 2 +IndentWrappedFunctionNames: true +InsertBraces: true +InsertTrailingCommas: None +InsertNewlineAtEOF: true +JavaScriptQuotes: Leave +JavaScriptWrapImports: true +KeepEmptyLinesAtTheStartOfBlocks: true +LambdaBodyIndentation: Signature +LineEnding: LF +MacroBlockBegin: '' +MacroBlockEnd: '' +MaxEmptyLinesToKeep: 1 +NamespaceIndentation: None +ObjCBinPackProtocolList: Auto +ObjCBlockIndentWidth: 2 +ObjCBreakBeforeNestedBlockParam: true +ObjCSpaceAfterProperty: false +ObjCSpaceBeforeProtocolList: true +PenaltyBreakAssignment: 2 +PenaltyBreakBeforeFirstCallParameter: 19 +PenaltyBreakComment: 300 +PenaltyBreakFirstLessLess: 120 +PenaltyBreakString: 1000 +PenaltyBreakTemplateDeclaration: 10 +PenaltyExcessCharacter: 1000000 +PenaltyReturnTypeOnItsOwnLine: 60 +PenaltyIndentedWhitespace: 0 +PointerAlignment: Right +PPIndentWidth: -1 +QualifierAlignment: Right +ReferenceAlignment: Pointer +ReflowComments: true +ShortNamespaceLines: 1 +SortIncludes: CaseSensitive +SortJavaStaticImport: Before +SortUsingDeclarations: true +SpaceAfterCStyleCast: false +SpaceAfterLogicalNot: false +SpaceAfterTemplateKeyword: true +SpaceBeforeAssignmentOperators: true +SpaceBeforeCaseColon: false +SpaceBeforeCpp11BracedList: false +SpaceBeforeCtorInitializerColon: true +SpaceBeforeInheritanceColon: true +SpaceBeforeParens: ControlStatements +SpaceAroundPointerQualifiers: Default +SpaceBeforeRangeBasedForLoopColon: true +SpaceInEmptyBlock: false +SpaceInEmptyParentheses: false +SpacesBeforeTrailingComments: 1 +SpacesInAngles: Never +SpacesInConditionalStatement: false +SpacesInContainerLiterals: true +SpacesInCStyleCastParentheses: false +SpacesInLineCommentPrefix: + Minimum: 1 + Maximum: -1 +SpacesInParentheses: false +SpacesInSquareBrackets: false +SpaceBeforeSquareBrackets: false +BitFieldColonSpacing: Both +Standard: Latest +StatementAttributeLikeMacros: + - Q_EMIT +StatementMacros: + - Q_UNUSED + - QT_REQUIRE_VERSION +TabWidth: 8 +UseCRLF: false +UseTab: Never +WhitespaceSensitiveMacros: + - STRINGIZE + - PP_STRINGIZE + - BOOST_PP_STRINGIZE + - NS_SWIFT_NAME + - CF_SWIFT_NAME +... + diff --git a/examples/cpp/DLRM/strategies/dlrm_strategy.cc b/examples/cpp/DLRM/strategies/dlrm_strategy.cc index 2fcc4173c9..a7fab8c3a8 100644 --- a/examples/cpp/DLRM/strategies/dlrm_strategy.cc +++ b/examples/cpp/DLRM/strategies/dlrm_strategy.cc @@ -356,4 +356,4 @@ int main(int argc, char **argv) { std::to_string(gpu) + "gpus.pb"; std::fstream outputFile(output.c_str(), std::ios::out | std::ios::trunc); strategy.SerializeToOstream(&outputFile); */ -} \ No newline at end of file +} diff --git a/lib/compiler/ffi/include/flexflow/compiler.h b/lib/compiler/ffi/include/flexflow/compiler.h index dec1f44188..14bd66150f 100644 --- a/lib/compiler/ffi/include/flexflow/compiler.h +++ b/lib/compiler/ffi/include/flexflow/compiler.h @@ -7,7 +7,6 @@ FLEXFLOW_FFI_BEGIN() - typedef enum { FLEXFLOW_SEARCH_ALGORITHM_DATA_PARALLEL } flexflow_search_algorithm_t; @@ -22,15 +21,18 @@ typedef enum { 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_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_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); // diff --git a/lib/compiler/test/test_dp.cc b/lib/compiler/test/test_dp.cc index 18b098a202..01e4189839 100644 --- a/lib/compiler/test/test_dp.cc +++ b/lib/compiler/test/test_dp.cc @@ -51,4 +51,4 @@ TEST_CASE("optimal_cost") { optimal_cost(pcg, allowed_machine_views, TestCostEstimator{}, resource); // TODO: check result -} \ No newline at end of file +} diff --git a/lib/ffi/include/flexflow/flexflow.h b/lib/ffi/include/flexflow/flexflow.h index 7f47a6c527..5b5eb1c473 100644 --- a/lib/ffi/include/flexflow/flexflow.h +++ b/lib/ffi/include/flexflow/flexflow.h @@ -13,11 +13,12 @@ FLEXFLOW_FFI_BEGIN(); #define CHECK_FLEXFLOW(status) \ do { \ - 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)); \ + 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 error at %s:%d : %s\n", \ __FILE__, \ diff --git a/lib/ffi/src/ffi.cc b/lib/ffi/src/ffi.cc index c21d26d93f..f6cf67cf2c 100644 --- a/lib/ffi/src/ffi.cc +++ b/lib/ffi/src/ffi.cc @@ -6,32 +6,27 @@ flexflow_error_t flexflow_status_is_ok(flexflow_error_t e, bool *result) { switch (e.error_source) { - case FLEXFLOW_ERROR_SOURCE_RUNTIME: - { + 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: - { + 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: - { + 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: - { + 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: - { + 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)); diff --git a/lib/kernels/include/kernels/reshape_kernels.h b/lib/kernels/include/kernels/reshape_kernels.h index fa752b6c2b..972f8ee9b2 100644 --- a/lib/kernels/include/kernels/reshape_kernels.h +++ b/lib/kernels/include/kernels/reshape_kernels.h @@ -30,4 +30,4 @@ void backward_kernel(ffStream_t stream, } // namespace Kernels } // namespace FlexFlow -#endif // _FLEXFLOW_OPS_KERNELS_RESHAPE_KERNELS_H \ No newline at end of file +#endif // _FLEXFLOW_OPS_KERNELS_RESHAPE_KERNELS_H diff --git a/lib/kernels/src/hip/aggregate_kernels.cpp b/lib/kernels/src/hip/aggregate_kernels.cpp index ff50a8c7ad..40faaaeeff 100644 --- a/lib/kernels/src/hip/aggregate_kernels.cpp +++ b/lib/kernels/src/hip/aggregate_kernels.cpp @@ -296,4 +296,4 @@ AggregatePerDeviceState::~AggregatePerDeviceState(void) { } // namespace Aggregate } // namespace Kernels -} // namespace FlexFlow \ No newline at end of file +} // namespace FlexFlow diff --git a/lib/kernels/src/hip/optimizer_kernel.cpp b/lib/kernels/src/hip/optimizer_kernel.cpp index 7f57d6a2fb..c22ecd7f5a 100644 --- a/lib/kernels/src/hip/optimizer_kernel.cpp +++ b/lib/kernels/src/hip/optimizer_kernel.cpp @@ -245,4 +245,4 @@ __host__ void AdamOptimizer::nccl_update_task_gpu(AdamOptimizer const *op, } #endif -}; // namespace FlexFlow \ No newline at end of file +}; // namespace FlexFlow diff --git a/lib/op-attrs/ffi/include/flexflow/op-attrs.h b/lib/op-attrs/ffi/include/flexflow/op-attrs.h index 9b31050fd9..08b8e26f83 100644 --- a/lib/op-attrs/ffi/include/flexflow/op-attrs.h +++ b/lib/op-attrs/ffi/include/flexflow/op-attrs.h @@ -18,10 +18,11 @@ typedef enum { 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_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_get_string(flexflow_opattrs_error_t, + char **); flexflow_error_t flexflow_opattrs_error_destroy(flexflow_opattrs_error_t); // diff --git a/lib/op-attrs/ffi/internal/internal/op-attrs.h b/lib/op-attrs/ffi/internal/internal/op-attrs.h index 395f91e9aa..df0d6ce61c 100644 --- a/lib/op-attrs/ffi/internal/internal/op-attrs.h +++ b/lib/op-attrs/ffi/internal/internal/op-attrs.h @@ -2,15 +2,14 @@ #define _FLEXFLOW_OPATTRS_FFI_INTERNAL_INTERNAL_OPATTRS_H #include "flexflow/op-attrs.h" -#include "op-attrs/ops/linear.h" -#include "op-attrs/param_sync.h" -#include "op-attrs/datatype.h" -#include "op-attrs/ops/pool_2d.h" +#include "internal/opaque.h" #include "op-attrs/activation.h" #include "op-attrs/datatype.h" -#include "op-attrs/ops/embedding.h" -#include "internal/opaque.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; @@ -34,4 +33,4 @@ 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 +#endif diff --git a/lib/op-attrs/ffi/src/op-attrs.cc b/lib/op-attrs/ffi/src/op-attrs.cc index 18854a8e9c..828574dc25 100644 --- a/lib/op-attrs/ffi/src/op-attrs.cc +++ b/lib/op-attrs/ffi/src/op-attrs.cc @@ -1,174 +1,184 @@ #include "flexflow/op-attrs.h" -#include "internal/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" -#include "internal/enums.h" -flexflow_utils_exception_t make_opattrs_exception(flexflow_opattrs_error_code_t); +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_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_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_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_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_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 }, - }); +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); } +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); } +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); } +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); } +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); } +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); } +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..ee1e54194f 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 { @@ -14,5 +15,6 @@ struct ff_dim_t : public numerical_typedef { MAKE_TYPEDEF_HASHABLE(::FlexFlow::ff_dim_t); MAKE_TYPEDEF_PRINTABLE(::FlexFlow::ff_dim_t, "ff_dim"); +static_assert(FlexFlow::is_neq_comparable<::FlexFlow::ff_dim_t>::value); #endif 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/include/flexflow/pcg.h b/lib/pcg/ffi/include/flexflow/pcg.h index 9b738512f6..03dee5ec95 100644 --- a/lib/pcg/ffi/include/flexflow/pcg.h +++ b/lib/pcg/ffi/include/flexflow/pcg.h @@ -43,15 +43,19 @@ extern flexflow_initializer_t NO_INITIALIZER; extern flexflow_regularizer_attrs_t NO_REGULARIZER; flexflow_error_t flexflow_pcg_error_wrap(flexflow_pcg_error_t); -flexflow_error_t flexflow_pcg_error_unwrap(flexflow_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_get_error_code(flexflow_pcg_error_t, + flexflow_pcg_error_code_t *); flexflow_error_t flexflow_pcg_error_destroy(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_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 @@ -59,8 +63,9 @@ flexflow_error_t flexflow_error_t flexflow_computation_graph_destroy(flexflow_computation_graph_t); -flexflow_error_t flexflow_computation_graph_serialize_to_buffer( - flexflow_computation_graph_t, char **out); +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); @@ -84,7 +89,8 @@ 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, flexflow_int_list_t *out); +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 diff --git a/lib/pcg/ffi/src/pcg.cc b/lib/pcg/ffi/src/pcg.cc index 431f500d9a..c958ed6c18 100644 --- a/lib/pcg/ffi/src/pcg.cc +++ b/lib/pcg/ffi/src/pcg.cc @@ -1,13 +1,14 @@ #include "pcg.h" #include "flexflow/pcg.h" #include "flexflow/utils.h" -#include "internal/opaque.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" -#include "internal/op-attrs.h" using namespace FlexFlow; @@ -34,18 +35,18 @@ std::vector make_vector(T const *ptr, int num_values) { throw make_pcg_exception(FLEXFLOW_PCG_STATUS_NEGATIVE_ARRAY_LENGTH_FOUND); } - return {ptr, ptr+num_values}; + return {ptr, ptr + num_values}; } -flexflow_error_t flexflow_pcg_error_get_string(flexflow_pcg_error_t err, char **m_out) { +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: - { + case FLEXFLOW_PCG_ERROR_UNKNOWN: { *out = "Unknown error"; break; } @@ -56,7 +57,9 @@ flexflow_error_t flexflow_pcg_error_get_string(flexflow_pcg_error_t err, char ** return status_ok(); } -flexflow_error_t flexflow_pcg_error_get_error_code(flexflow_error_t err, flexflow_pcg_error_code_t *out) { +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); @@ -69,12 +72,14 @@ 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) { +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([&] { }); +flexflow_error_t + flexflow_computation_graph_destroy(flexflow_computation_graph_t opaque) { + return handle_errors([&] {}); } template @@ -84,26 +89,32 @@ std::string get_v1_string(T const &t) { template T from_v1_string(char const *s) { - return json::parse(s).template get>()))>(); + return json::parse(s) + .template get>()))>(); } -flexflow_error_t flexflow_computation_graph_serialize_to_buffer(flexflow_computation_graph_t opaque, char **out) { +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); + *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_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) { +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); @@ -111,19 +122,21 @@ flexflow_error_t flexflow_computation_graph_serialize_to_file(flexflow_computati 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); + 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) { +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; + std::ostringstream oss; while (true) { size_t num_bytes_read = fread(buf, sizeof(char), BUF_SIZE, f); @@ -134,7 +147,7 @@ flexflow_error_t flexflow_computation_graph_deserialize_from_file(FILE *f, flexf } else { throw make_pcg_exception(FLEXFLOW_PCG_STATUS_FILE_READ_FAILED); } - } + } } return from_v1_string(oss.str().c_str()); @@ -147,30 +160,36 @@ flexflow_error_t flexflow_tensor_create(flexflow_computation_graph_t cg, flexflow_datatype_t datatype, bool create_grad, flexflow_tensor_t *out) { - return handle_errors(out, [&] { + return handle_errors(out, [&] { TensorDims ordered_dims{make_vector(dims, num_dims)}; - TensorShape shape = { ordered_dims, to_internal(datatype) }; + 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_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) { +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) { +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) { +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) { +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(); }); } @@ -183,19 +202,34 @@ flexflow_error_t flexflow_tensor_destroy(flexflow_tensor_t opaque) { } optional maybe_string(char *); -using BinaryFunc = std::function const &)>; -using UnaryFunc = std::function)>; -using ScalarUnaryFunc = std::function)>; - -#define BINARY_OP(func) \ - [](ComputationGraph &cg, Tensor const &lhs, Tensor const &rhs, optional const &name) -> Tensor { return func(cg, lhs, rhs, name); } - -#define UNARY_OP(func) \ - [](ComputationGraph &cg, Tensor const &t, optional const &name) -> Tensor { return func(cg, t, name); } +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 SCALAR_UNARY_OP(func) \ - [](ComputationGraph &cg, Tensor const &t, float scalar, optional const &name) -> Tensor { return func(cg, t, scalar, 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, @@ -203,186 +237,224 @@ flexflow_error_t add_op(BinaryFunc const &f, 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)); }); + 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_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) { + 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(UNARY_OP(exp), 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(add), 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(subtract), 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(multiply), 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(divide), 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(max), 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(min), 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(UNARY_OP(rsqrt), cg, input, out, 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_scalar_multiply(flexflow_computation_graph_t cg, - flexflow_tensor_t input, - float scalar, - flexflow_tensor_t *out, - char *name) { - return add_op(SCALAR_UNARY_OP(scalar_multiply), cg, input, scalar, 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_UNARY_OP(scalar_add), 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_UNARY_OP(scalar_sub), 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_UNARY_OP(scalar_truediv), 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(UNARY_OP(sin), cg, input, 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(UNARY_OP(cos), 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(UNARY_OP(relu), 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(UNARY_OP(identity), 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(UNARY_OP(gelu), 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(UNARY_OP(sigmoid), 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(UNARY_OP(tanh), 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(UNARY_OP(tanh), 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( @@ -403,86 +475,96 @@ flexflow_error_t flexflow_computation_graph_add_op_conv2d( 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), + 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)); }); + 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_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( @@ -499,325 +581,356 @@ flexflow_error_t flexflow_computation_graph_add_op_cache( 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 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_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_batch_norm(flexflow_computation_graph_t cg, - flexflow_tensor_t input, - flexflow_tensor_t *out, - bool relu, - char *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)); + }); +} - 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 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 index a0d26a7afb..fa9e00a1d0 100644 --- a/lib/pcg/ffi/src/pcg.h +++ b/lib/pcg/ffi/src/pcg.h @@ -3,28 +3,66 @@ #include "flexflow/pcg.h" #include "internal/opaque.h" -#include "pcg/computation_graph.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/parallel_computation_graph.h" -#include "pcg/operator.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, Operator); -REGISTER_OPAQUE(flexflow_parallel_tensor_t, ParallelTensor); -REGISTER_OPAQUE(flexflow_layer_t, Layer); -REGISTER_OPAQUE(flexflow_tensor_t, Tensor); +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); - +REGISTER_OPAQUE(flexflow_tensor_list_t, + std::vector); + struct internal_flexflow_pcg_error_t { flexflow_pcg_error_code_t err_code; }; diff --git a/lib/pcg/include/pcg/computation_graph.h b/lib/pcg/include/pcg/computation_graph.h index 5258242bf1..a329198784 100644 --- a/lib/pcg/include/pcg/computation_graph.h +++ b/lib/pcg/include/pcg/computation_graph.h @@ -3,25 +3,51 @@ #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 { struct ComputationGraph : public strong_typedef> { - ComputationGraph(); - 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 ebe65e12e0..c7c8cee14c 100644 --- a/lib/pcg/include/pcg/computation_graph_builder.h +++ b/lib/pcg/include/pcg/computation_graph_builder.h @@ -4,315 +4,318 @@ #include "computation_graph.h" namespace FlexFlow { -// -// C++ APIs for constructing models -// Add an exp layer -Tensor exp(ComputationGraph &, Tensor const &, optional const &name = nullopt); -// Add an add layer -Tensor add(ComputationGraph &, - Tensor const &x, - Tensor const &y, - optional const &name = nullopt); -// Add a subtract layer -Tensor subtract(ComputationGraph &, - Tensor const &x, - Tensor const &y, - optional const &name = nullopt); -// Add a multiply layer -Tensor multiply(ComputationGraph &, - Tensor const &x, - Tensor const &y, - optional const &name = nullopt); -// Add a divide layer -Tensor divide(ComputationGraph &, - Tensor const &x, - Tensor const &y, - optional const &name = nullopt); -// Add a max layer -Tensor max(ComputationGraph &, - Tensor const &x, - Tensor const &y, - optional const &name = nullopt); -// Add a min layer -Tensor min(ComputationGraph &, - Tensor const &x, - Tensor const &y, - optional const &name = nullopt); -// Add a rsqrt layer -Tensor rsqrt(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -// Add a pow layer -Tensor pow(ComputationGraph &, - Tensor const &x, - float exponent, - optional const &name = nullopt); -// Add a scalar multiply layer -Tensor scalar_multiply(ComputationGraph &, - Tensor const &x, - float scalar, - optional const &name = nullopt); -Tensor scalar_add(ComputationGraph &, - Tensor const &x, - float scalar, - optional const &name = nullopt); -Tensor scalar_sub(ComputationGraph &, - Tensor const &lhs, - float rhs, - optional const &name = nullopt); -Tensor scalar_truediv(ComputationGraph &, - Tensor const &numerator, - float denominator, - optional const &name = nullopt); -// Add a sin layer -Tensor sin(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -// Add a cos layer -Tensor cos(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -// Add an activation layer -Tensor relu(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -Tensor identity(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -Tensor gelu(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -Tensor sigmoid(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -Tensor tanh(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -Tensor elu(ComputationGraph &, Tensor const &x, optional const &name = nullopt); -// Add a 2D convolutional layer -Tensor conv2d(ComputationGraph &, - 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 const &kernel_initializer = nullopt, - optional const &bias_initializer = nullopt, - optional const &kernel_regularizer = nullopt, - optional const &name = nullopt); -// Add a dropout layer -Tensor dropout(ComputationGraph &, - Tensor const &input, - float rate, - unsigned long long seed = 0, - optional const &name = nullopt); -// Add an embedding layer -Tensor embedding(ComputationGraph &, - Tensor const &input, - int num_entries, - int outDim, - AggregateOp aggr, - DataType dtype = DataType::FLOAT, - optional const &kernel_initializer = nullopt, - optional const &name = nullopt); -// Add a gather layer -std::vector gather(ComputationGraph &, - Tensor const &input, - Tensor const &index, - ff_dim_t dim, + +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); + +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_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); -// Add a group_by layer -std::vector group_by(ComputationGraph &, - Tensor const &data, - Tensor const &assign, - int n, - float alpha, - optional const &name = nullopt); -// Add a cache layer -Tensor cache(ComputationGraph &, - Tensor const &input, - int num_batches, - std::function - score_f = {}, - optional const &name = nullopt); -// Add aggregate layer -Tensor aggregate(ComputationGraph &, - Tensor const &gate_preds, - Tensor const &gate_assign, - Tensor const &true_gate_assign, - Tensor const &full_gate_gradients, - std::vector const &exp_preds, - float lambda_bal, - optional const &name = nullopt); -// Add aggregate_spec layer -Tensor aggregate_spec(ComputationGraph &, - std::vector const &inputs, - float lambda_bal, - optional const &name = nullopt); -// Add a 2D pooling layer -Tensor pool2d(ComputationGraph &, - 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(ComputationGraph &, - Tensor const &input, - std::vector const &axes, - bool elementwise_affine, - float eps, - optional const &name = nullopt); -Tensor batch_norm(ComputationGraph &, - Tensor const &input, - bool relu = true, - optional const &name = nullopt); -Tensor batch_matmul(ComputationGraph &, - Tensor const &A, - Tensor const &B, - int a_seq_length_dim = -1, - int b_seq_length_dim = -1, - optional const &name = nullopt); -Tensor dense(ComputationGraph &, - Tensor 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); -// Add a cast layer -Tensor cast(ComputationGraph &, - Tensor const &input, - DataType dtype, - optional const &name = nullopt); -// Add a concat layer -Tensor concat(ComputationGraph &, - std::vector const &tensors, - int axis, - optional const &name = nullopt); -// Add a mean layer -Tensor mean(ComputationGraph &, - Tensor const &input, - std::vector const &dims, - bool keepdims, - optional const &name = nullopt); -// Add a moe layer (wrapping topk, group_by and aggregate operators) -Tensor moe(ComputationGraph &, - Tensor const &input, - int num_exp, - int num_select, - int expert_hidden_size, - float alpha, - float lambda, - optional const &name = nullopt); -// Add a split layer -std::vector split(ComputationGraph &, - Tensor const &input, - std::vector const &split, - int axis, - optional const &name = nullopt); -// Add a flat layer -Tensor flat(ComputationGraph &, Tensor const &input, optional const &name = nullopt); -// Add a softmax layer -Tensor softmax(ComputationGraph &, - Tensor const &input, - int dim = -1, - optional const &name = nullopt); -// Create input tensors and constants -Tensor transpose(ComputationGraph &, - Tensor const &input, - std::vector const &perm, - optional const &name = nullopt); -Tensor reduce_sum(ComputationGraph &, - Tensor const &input, - std::vector const &axes, - bool keepdims = false, - optional const &name = nullopt); -Tensor reshape(ComputationGraph &, - Tensor const &input, - std::vector const &shape, - optional const &name = nullopt); -Tensor reverse(ComputationGraph &, - Tensor const &input, - int axis, - optional const &name = nullopt); -std::vector top_k(ComputationGraph &, - Tensor const &input, - int k, - bool sorted, - optional const &name = nullopt); -Tensor - multihead_attention(ComputationGraph &, - 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 const &initializer = 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 create_tensor(ComputationGraph &, 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, - 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 - element_unary(variant const &, - Tensor const &input, - 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_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); -} // namespace FlexFlow +tensor_guid_t insert_cast_layer(ComputationGraph &, + tensor_guid_t const &input, + DataType dtype, + optional const &name = nullopt); -VISITABLE_STRUCT(::FlexFlow::ComputationGraphBuilder, computation_graph); +tensor_guid_t insert_concat_layer(ComputationGraph &, + std::vector const &tensors, + int axis, + optional const &name = nullopt); -namespace FlexFlow { -static_assert( - is_well_behaved_value_type_no_hash::value, ""); -} +tensor_guid_t insert_mean_layer(ComputationGraph &, + tensor_guid_t const &input, + std::vector const &dims, + bool keepdims, + optional const &name = nullopt); + +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/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..b6a60a2e3a --- /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/strong_typedef.h" +#include "utils/graph.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"); + +} + +#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 60ed3ec49e..ad1193db87 100644 --- a/lib/pcg/include/pcg/tensor.h +++ b/lib/pcg/include/pcg/tensor.h @@ -13,16 +13,17 @@ struct Tensor { TensorShape get_shape() const; int num_dims() const; + DataType get_data_type() const; + operator TensorShape() const; + public: - TensorDims dims; - DataType data_type; + TensorShape shape; + CreateGrad create_gradients; optional initializer; - bool create_gradients; 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..e2c070a8c0 100644 --- a/lib/pcg/src/computation_graph_builder.cc +++ b/lib/pcg/src/computation_graph_builder.cc @@ -1,20 +1,151 @@ #include "pcg/computation_graph_builder.h" #include "op-attrs/get_op_type.h" #include "op-attrs/get_output_shapes.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" 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 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 TensorShape get_broadcast_target_shape(std::vector const &); + +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 +164,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); - 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 &, + 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(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 &, + 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(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({lhs, rhs}); + DataType compute_type = std::max(get_data_type(lhs), get_data_type(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( + broadcast(lhs, compute_shape), compute_type, name + "_inputl_pre_cast"); + tensor_guid_t const rhs_input = as_type( + broadcast(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); - return this->add_layer(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 ComputationGraphBuilder::subtract(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_SUB, lhs, rhs, name); + return insert_layer(cg, layer, {lhs_input, rhs_input}, {}, output_shape); } -Tensor ComputationGraphBuilder::multiply(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_MUL, 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::divide(Tensor const &lhs, - Tensor const &rhs, - optional const &name) { - return this->element_binary(Op::EW_DIV, 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::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); -} - -Tensor ComputationGraphBuilder::rsqrt(Tensor const &input, - optional const &name) { - return this->element_unary(Op::RSQRT, input, name); + return element_binary(cg, Op::EW_MUL, 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_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::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_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_add(Tensor const &input, - float scalar, - optional const &name) { - return this->element_scalar_unary(Op::SCALAR_ADD, 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_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,7 +390,8 @@ 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); @@ -247,81 +404,88 @@ Tensor ComputationGraphBuilder::conv2d( weights.push_back({get_bias_shape(attrs, 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); - 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(x, DataType::FLOAT, name + "input_pre_cast"); TensorShape output_shape = get_output_shape(attrs, input); TensorShape weights_shape = get_weights_shape(attrs, 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)); + Tensor index_tensor = cg.at(index); + DataType index_dt = get_data_type(index_tensor); + 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(attrs, input, index_tensor); - 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 get_data_type(cg.at(t)); + }; -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,16 +495,18 @@ 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)); @@ -348,7 +514,7 @@ Tensor ComputationGraphBuilder::batch_norm( TensorShape output_shape = get_output_shape(attrs, get_shape(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/include/flexflow/runtime.h b/lib/runtime/ffi/include/flexflow/runtime.h index eec6475844..e9b4f6cccf 100644 --- a/lib/runtime/ffi/include/flexflow/runtime.h +++ b/lib/runtime/ffi/include/flexflow/runtime.h @@ -1,10 +1,10 @@ #ifndef _FLEXFLOW_RUNTIME_INCLUDE_RUNTIME_FFI_RUNTIME_H #define _FLEXFLOW_RUNTIME_INCLUDE_RUNTIME_FFI_RUNTIME_H -#include "flexflow/utils.h" #include "flexflow/compiler.h" #include "flexflow/op-attrs.h" #include "flexflow/pcg.h" +#include "flexflow/utils.h" #include #include @@ -23,18 +23,20 @@ typedef enum { FLEXFLOW_RUNTIME_ERROR_DYNAMIC_ALLOCATION_FAILED, FLEXFLOW_RUNTIME_ERROR_UNEXPECTED_EMPTY_HANDLE, } 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_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_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, @@ -62,14 +64,16 @@ flexflow_error_t flexflow_void_future_wait(flexflow_void_future_t); flexflow_error_t flexflow_void_future_destroy(flexflow_void_future_t); flexflow_error_t flexflow_config_parse_argv(int *argc, - char **argv, - bool remove_used, - flexflow_config_t *out); + 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_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_error_t flexflow_computation_graph_set_model_config(flexflow_computation_graph_t, @@ -91,40 +95,35 @@ flexflow_error_t flexflow_tensor_t, flexflow_parallel_tensor_t *); -flexflow_error_t - flexflow_start_training(flexflow_model_compilation_result_t, - flexflow_model_training_instance_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_error_t flexflow_model_training_instance_backward( flexflow_model_training_instance_t); -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_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(); diff --git a/lib/runtime/ffi/src/runtime.cc b/lib/runtime/ffi/src/runtime.cc index 46cdcd766a..5de26d39d0 100644 --- a/lib/runtime/ffi/src/runtime.cc +++ b/lib/runtime/ffi/src/runtime.cc @@ -1,11 +1,13 @@ #include "flexflow/runtime.h" -#include "runtime/model_training_instance.h" #include "internal/opaque.h" +#include "runtime/model_training_instance.h" #include "utils/expected.h" using namespace FlexFlow; -flexflow_error_t flexflow_runtime_error_create(flexflow_runtime_error_t) { NOT_IMPLEMENTED(); } +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/runtime/test/src/test_op_task_spec.cc b/lib/runtime/test/src/test_op_task_spec.cc index 821ef0dba6..bb0bee567c 100644 --- a/lib/runtime/test/src/test_op_task_spec.cc +++ b/lib/runtime/test/src/test_op_task_spec.cc @@ -44,4 +44,4 @@ TEST_CASE("OpTaskBinding") { correct_bwd.bind_grad(2, input_tensor(2).grad()); CHECK(correct_bwd == bwd); -} \ No newline at end of file +} diff --git a/lib/runtime/test/src/test_serialization.cc b/lib/runtime/test/src/test_serialization.cc index caf5cd7c93..ef18764efb 100644 --- a/lib/runtime/test/src/test_serialization.cc +++ b/lib/runtime/test/src/test_serialization.cc @@ -44,4 +44,4 @@ TEST_CASE("Serialization") { RC_ASSERT(post_op == pre_op); })) } -} \ No newline at end of file +} diff --git a/lib/utils/ffi/include/flexflow/utils.h b/lib/utils/ffi/include/flexflow/utils.h index eed9bace98..084cd134ba 100644 --- a/lib/utils/ffi/include/flexflow/utils.h +++ b/lib/utils/ffi/include/flexflow/utils.h @@ -27,7 +27,7 @@ typedef enum { } flexflow_error_source_t; typedef enum { - FLEXFLOW_UTILS_STATUS_OK, + FLEXFLOW_UTILS_STATUS_OK, FLEXFLOW_UTILS_DEALLOCATION_FAILED, FLEXFLOW_UTILS_ALLOCATION_FAILED, FLEXFLOW_UTILS_CAST_FAILED, @@ -48,7 +48,8 @@ typedef struct { 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_utils_error_unwrap(flexflow_error_t, + flexflow_utils_error_t *); flexflow_error_t flexflow_status_is_ok(flexflow_error_t, bool *); diff --git a/lib/utils/ffi/internal/internal/enums.h b/lib/utils/ffi/internal/internal/enums.h index 2604c05a11..02e6b47bbd 100644 --- a/lib/utils/ffi/internal/internal/enums.h +++ b/lib/utils/ffi/internal/internal/enums.h @@ -7,7 +7,7 @@ struct internal_to_external; template struct external_to_internal; -template +template struct enum_mapping; template @@ -16,39 +16,39 @@ 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__; +#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(); + 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(); + 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 index af8ea7fdd2..dbf689f24b 100644 --- a/lib/utils/ffi/internal/internal/error.h +++ b/lib/utils/ffi/internal/internal/error.h @@ -4,16 +4,16 @@ #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 { +#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; @@ -27,10 +27,11 @@ 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) { +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; + flexflow_error_t result; result.error_source = error_source; T *buf_ptr = static_cast(result.buf); *buf_ptr = t; @@ -39,14 +40,17 @@ flexflow_error_t flexflow_error_wrap(flexflow_error_source_t error_source, T con } 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)), ""); +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); + out->impl = + reinterpret_cast()))>(err.buf); return flexflow_utils_error_create(FLEXFLOW_UTILS_STATUS_OK); } diff --git a/lib/utils/ffi/internal/internal/error_handling.h b/lib/utils/ffi/internal/internal/error_handling.h index a25f75b939..441d12e86c 100644 --- a/lib/utils/ffi/internal/internal/error_handling.h +++ b/lib/utils/ffi/internal/internal/error_handling.h @@ -4,27 +4,13 @@ #include "error.h" #include "opaque.h" -template -flexflow_error_t handle_errors(Opaque *out, std::function()> const &f) { - try { - *out = new_opaque(f()); - } catch (flexflow_ffi_exception_t const &e) { - return to_error(e); - } - - return status_ok(); -} - +flexflow_error_t handle_errors(std::function const &f); -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(); +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 index 78cf7e34a2..35c929aca0 100644 --- a/lib/utils/ffi/internal/internal/opaque.h +++ b/lib/utils/ffi/internal/internal/opaque.h @@ -1,13 +1,13 @@ #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 "flexflow/utils.h" -#include "error.h" #include -#include "utils/containers.h" template struct opaque_to_underlying; @@ -24,7 +24,8 @@ 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); + throw make_utils_exception( + FLEXFLOW_UTILS_UNEXPECTED_NULLPTR_IN_OPAQUE_HANDLE); } return static_cast *>(opaque.impl); @@ -46,23 +47,26 @@ opaque_to_underlying_t &deref_opaque(Opaque const &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); }); +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; \ - }; +#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) { +Opaque new_opaque(Args &&...args) { using Underlying = opaque_to_underlying_t; Underlying *ptr = new (std::nothrow) Underlying(std::forward(args)...); @@ -83,7 +87,8 @@ void delete_opaque(Opaque const &opaque) { Underlying *underlying = unwrap_opaque(opaque); if (underlying == nullptr) { - throw make_utils_exception(FLEXFLOW_UTILS_UNEXPECTED_NULLPTR_IN_OPAQUE_HANDLE); + throw make_utils_exception( + FLEXFLOW_UTILS_UNEXPECTED_NULLPTR_IN_OPAQUE_HANDLE); } delete underlying; diff --git a/lib/utils/ffi/src/error_handling.cc b/lib/utils/ffi/src/error_handling.cc new file mode 100644 index 0000000000..6b30021e4b --- /dev/null +++ b/lib/utils/ffi/src/error_handling.cc @@ -0,0 +1,13 @@ +#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 6432c97cad..611a806c95 100644 --- a/lib/utils/include/utils/bidict.h +++ b/lib/utils/include/utils/bidict.h @@ -1,8 +1,8 @@ #ifndef _FLEXFLOW_UTILS_BIDICT_H #define _FLEXFLOW_UTILS_BIDICT_H -#include #include "optional.h" +#include namespace FlexFlow { @@ -18,8 +18,7 @@ struct bidict { bidict() : fwd_map{}, bwd_map{} {} bidict(std::initializer_list init) - : bidict(init.begin(), init.end()) - { } + : bidict(init.begin(), init.end()) {} template bidict(InputIt first, InputIt last) { diff --git a/lib/utils/include/utils/containers.h b/lib/utils/include/utils/containers.h index 568e6451a2..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(); } 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/graph/labelled/labelled_downward_open_interfaces.h b/lib/utils/include/utils/graph/labelled/labelled_downward_open_interfaces.h index e77787294d..98b579df13 100644 --- a/lib/utils/include/utils/graph/labelled/labelled_downward_open_interfaces.h +++ b/lib/utils/include/utils/graph/labelled/labelled_downward_open_interfaces.h @@ -14,7 +14,8 @@ struct ILabelledDownwardOpenMultiDiGraphView public IDownwardOpenMultiDiGraphView { virtual ~ILabelledDownwardOpenMultiDiGraphView() = default; - std::unordered_set query_edges(MultiDiEdgeQuery const &q) const final { + std::unordered_set + query_edges(MultiDiEdgeQuery const &q) const final { return this->query_edges(static_cast(q)); } diff --git a/lib/utils/include/utils/graph/labelled/open_algorithms.h b/lib/utils/include/utils/graph/labelled/open_algorithms.h index 89f8674373..4b54aeeece 100644 --- a/lib/utils/include/utils/graph/labelled/open_algorithms.h +++ b/lib/utils/include/utils/graph/labelled/open_algorithms.h @@ -150,16 +150,39 @@ ResultType get_subgraph(LabelledOpenMultiDiGraph(as_view(g), nodes); } -template +template LabelledUpwardOpenMultiDiGraphView -as_upward_open(LabelledOpenMultiDiGraphView const &g) { - return LabelledUpwardOpenMultiDiGraphView::template create>(g); + as_upward_open(LabelledOpenMultiDiGraphView const &g) { + return LabelledUpwardOpenMultiDiGraphView:: + template create>(g); } -template +template LabelledDownwardOpenMultiDiGraphView -as_downward_open(LabelledOpenMultiDiGraphView const &g) { - return LabelledDownwardOpenMultiDiGraphView::template create>(g); + as_downward_open(LabelledOpenMultiDiGraphView const &g) { + return LabelledDownwardOpenMultiDiGraphView:: + template create>( + g); } } // namespace FlexFlow diff --git a/lib/utils/include/utils/graph/labelled/open_views.h b/lib/utils/include/utils/graph/labelled/open_views.h index 4243e6ba43..4c58af9386 100644 --- a/lib/utils/include/utils/graph/labelled/open_views.h +++ b/lib/utils/include/utils/graph/labelled/open_views.h @@ -9,12 +9,12 @@ #include "labelled_upward_open_interfaces.h" #include "standard_labelled.h" #include "utils/exception.h" +#include "utils/graph/algorithms.h" #include "utils/graph/multidiedge.h" #include "utils/graph/open_graph_interfaces.h" #include "utils/graph/open_graphs.h" #include "utils/type_traits.h" #include "utils/visitable.h" -#include "utils/graph/algorithms.h" namespace FlexFlow { @@ -51,48 +51,97 @@ struct LabelledDownwardMultiDiSubgraphView { std::unordered_set const &); }; -template -struct ViewLabelledOpenMultiDiGraphAsUpwardOpen : public ILabelledUpwardOpenMultiDiGraphView { +template +struct ViewLabelledOpenMultiDiGraphAsUpwardOpen + : public ILabelledUpwardOpenMultiDiGraphView { public: - using InputType = LabelledOpenMultiDiGraphView; + using InputType = LabelledOpenMultiDiGraphView; - explicit ViewLabelledOpenMultiDiGraphAsUpwardOpen(InputType const &g) : g(g) { } + explicit ViewLabelledOpenMultiDiGraphAsUpwardOpen(InputType const &g) + : g(g) {} - std::unordered_set query_nodes(NodeQuery const &q) const override { return this->g.query_nodes(q); } + std::unordered_set query_nodes(NodeQuery const &q) const override { + return this->g.query_nodes(q); + } - std::unordered_set query_edges(UpwardOpenMultiDiEdgeQuery const &q) const override { - return value_all(narrow(this->g.query_edges(q))); + std::unordered_set + query_edges(UpwardOpenMultiDiEdgeQuery const &q) const override { + return value_all( + narrow(this->g.query_edges(q))); + } + + NodeLabel const &at(Node const &n) const override { + return this->g.at(n); + } + InputLabel const &at(InputMultiDiEdge const &e) const override { + return this->g.at(e); + } + EdgeLabel const &at(MultiDiEdge const &e) const override { + return this->g.at(e); } - NodeLabel const &at(Node const &n) const override { return this->g.at(n); } - InputLabel const &at(InputMultiDiEdge const &e) const override { return this->g.at(e); } - EdgeLabel const &at(MultiDiEdge const &e) const override { return this->g.at(e); } private: InputType g; }; -CHECK_NOT_ABSTRACT(ViewLabelledOpenMultiDiGraphAsUpwardOpen); +CHECK_NOT_ABSTRACT( + ViewLabelledOpenMultiDiGraphAsUpwardOpen); -template -struct ViewLabelledOpenMultiDiGraphAsDownwardOpen : public ILabelledDownwardOpenMultiDiGraphView { +template +struct ViewLabelledOpenMultiDiGraphAsDownwardOpen + : public ILabelledDownwardOpenMultiDiGraphView { public: - using InputType = LabelledOpenMultiDiGraphView; + using InputType = LabelledOpenMultiDiGraphView; - explicit ViewLabelledOpenMultiDiGraphAsDownwardOpen(InputType const &g) : g(g) { } + explicit ViewLabelledOpenMultiDiGraphAsDownwardOpen(InputType const &g) + : g(g) {} - std::unordered_set query_nodes(NodeQuery const &q) const override { return this->g.query_nodes(q); } + std::unordered_set query_nodes(NodeQuery const &q) const override { + return this->g.query_nodes(q); + } - std::unordered_set query_edges(DownwardOpenMultiDiEdgeQuery const &q) const override { - return value_all(narrow(this->g.query_edges(q))); + std::unordered_set + query_edges(DownwardOpenMultiDiEdgeQuery const &q) const override { + return value_all( + narrow(this->g.query_edges(q))); } - NodeLabel const &at(Node const &n) const override { return this->g.at(n); } - OutputLabel const &at(OutputMultiDiEdge const &e) const override { return this->g.at(e); } - EdgeLabel const &at(MultiDiEdge const &e) const override { return this->g.at(e); } + NodeLabel const &at(Node const &n) const override { + return this->g.at(n); + } + OutputLabel const &at(OutputMultiDiEdge const &e) const override { + return this->g.at(e); + } + EdgeLabel const &at(MultiDiEdge const &e) const override { + return this->g.at(e); + } private: InputType g; }; -CHECK_NOT_ABSTRACT(ViewLabelledOpenMultiDiGraphAsUpwardOpen); +CHECK_NOT_ABSTRACT( + ViewLabelledOpenMultiDiGraphAsUpwardOpen); template ; - LabelledOpenMultiDiSubgraphView(LabelledOpenMultiDiGraphView const &g, std::unordered_set const &nodes) - : g(g), nodes(nodes) - { } + LabelledOpenMultiDiSubgraphView( + LabelledOpenMultiDiGraphView const &g, + std::unordered_set const &nodes) + : g(g), nodes(nodes) {} std::unordered_set query_edges(UpwardOpenMultiDiEdgeQuery const &q) const override { @@ -186,7 +239,7 @@ struct LabelledOpenMultiDiSubgraphView query_nodes(NodeQuery const &q) const override { return static_cast(this->g).query_nodes(q); - + NOT_IMPLEMENTED(); } diff --git a/lib/utils/include/utils/graph/open_graph_interfaces.h b/lib/utils/include/utils/graph/open_graph_interfaces.h index 6d05c21f37..3d7e66119d 100644 --- a/lib/utils/include/utils/graph/open_graph_interfaces.h +++ b/lib/utils/include/utils/graph/open_graph_interfaces.h @@ -88,12 +88,15 @@ FF_VISITABLE_STRUCT(OpenMultiDiEdgeQuery, struct DownwardOpenMultiDiEdgeQuery { DownwardOpenMultiDiEdgeQuery() = delete; DownwardOpenMultiDiEdgeQuery(OutputMultiDiEdgeQuery const &output_edge_query, - MultiDiEdgeQuery const &standard_edge_query) - : output_edge_query(output_edge_query), standard_edge_query(standard_edge_query) { } - DownwardOpenMultiDiEdgeQuery(OutputMultiDiEdgeQuery const &output_edge_query) - : DownwardOpenMultiDiEdgeQuery(output_edge_query, MultiDiEdgeQuery::none()) { } - DownwardOpenMultiDiEdgeQuery(MultiDiEdgeQuery const &standard_edge_query) - : DownwardOpenMultiDiEdgeQuery(OutputMultiDiEdgeQuery::all(), standard_edge_query) { }; + MultiDiEdgeQuery const &standard_edge_query) + : output_edge_query(output_edge_query), + standard_edge_query(standard_edge_query) {} + DownwardOpenMultiDiEdgeQuery(OutputMultiDiEdgeQuery const &output_edge_query) + : DownwardOpenMultiDiEdgeQuery(output_edge_query, + MultiDiEdgeQuery::none()) {} + DownwardOpenMultiDiEdgeQuery(MultiDiEdgeQuery const &standard_edge_query) + : DownwardOpenMultiDiEdgeQuery(OutputMultiDiEdgeQuery::all(), + standard_edge_query){}; operator OpenMultiDiEdgeQuery() const { NOT_IMPLEMENTED(); @@ -103,8 +106,8 @@ struct DownwardOpenMultiDiEdgeQuery { MultiDiEdgeQuery standard_edge_query; }; FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(DownwardOpenMultiDiEdgeQuery, - output_edge_query, - standard_edge_query); + output_edge_query, + standard_edge_query); struct UpwardOpenMultiDiEdgeQuery { UpwardOpenMultiDiEdgeQuery() = delete; @@ -133,10 +136,11 @@ struct IDownwardOpenMultiDiGraphView : public IOpenMultiDiGraphView { virtual std::unordered_set query_edges(DownwardOpenMultiDiEdgeQuery const &) const = 0; - std::unordered_set query_edges(OpenMultiDiEdgeQuery const &q) const final { + std::unordered_set + query_edges(OpenMultiDiEdgeQuery const &q) const final { return widen( - this->query_edges(DownwardOpenMultiDiEdgeQuery{ q.output_edge_query, q.standard_edge_query }) - ); + this->query_edges(DownwardOpenMultiDiEdgeQuery{q.output_edge_query, + q.standard_edge_query})); } }; CHECK_RC_COPY_VIRTUAL_COMPLIANT(IDownwardOpenMultiDiGraphView); @@ -145,10 +149,10 @@ struct IUpwardOpenMultiDiGraphView : public IOpenMultiDiGraphView { virtual std::unordered_set query_edges(UpwardOpenMultiDiEdgeQuery const &) const = 0; - std::unordered_set query_edges(OpenMultiDiEdgeQuery const &q) const final { - return widen( - this->query_edges(UpwardOpenMultiDiEdgeQuery{ q.input_edge_query, q.standard_edge_query }) - ); + std::unordered_set + query_edges(OpenMultiDiEdgeQuery const &q) const final { + return widen(this->query_edges( + UpwardOpenMultiDiEdgeQuery{q.input_edge_query, q.standard_edge_query})); } }; CHECK_RC_COPY_VIRTUAL_COMPLIANT(IUpwardOpenMultiDiGraphView); diff --git a/lib/utils/include/utils/graph/query_set.h b/lib/utils/include/utils/graph/query_set.h index 656d64d531..e7f4f60812 100644 --- a/lib/utils/include/utils/graph/query_set.h +++ b/lib/utils/include/utils/graph/query_set.h @@ -20,9 +20,8 @@ struct query_set { query_set(optional> const &) { NOT_IMPLEMENTED(); } - query_set(std::initializer_list const &l) - : query_set(std::unordered_set{l}) - { } + query_set(std::initializer_list const &l) + : query_set(std::unordered_set{l}) {} friend bool operator==(query_set const &lhs, query_set const &rhs) { return lhs.value == rhs.value; 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 2b7876496b..6ec21f7dbe 100644 --- a/lib/utils/include/utils/test_types.h +++ b/lib/utils/include/utils/test_types.h @@ -12,84 +12,122 @@ enum capability { EQ, CMP, DEFAULT_CONSTRUCTIBLE, - COPYABLE + COPYABLE, + PLUS, + PLUSEQ }; -template -struct capability_implies : std::false_type { }; +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 { }; +struct capability_implies : std::true_type {}; +template +struct has_capability; -template struct has_capability; - -template -struct has_capability : disjunction< - capability_implies, - has_capability> { }; +template +struct has_capability + : disjunction, + has_capability> {}; template -struct has_capability : std::false_type { }; +struct has_capability : std::false_type {}; + +template +struct test_type_t; + +/* template */ +/* struct has_capability> : std::true_type + * { }; */ + +static_assert(has_capability::value, ""); -template +template struct test_type_t { - template + template using supports = conjunction...>; - template::value, bool>::type = true> + /* template */ + /* using supports = conjunction>; */ + + template ::value, bool>::type = true> test_type_t(); - template::value, bool>::type = true> + template ::value, bool>::type = true> test_type_t() = delete; - template::value, bool>::type = true> + template ::value, bool>::type = true> test_type_t(test_type_t const &); - template::value, bool>::type = true> + template ::value, bool>::type = true> test_type_t(test_type_t const &) = delete; - typename std::enable_if::value, bool>::type - operator==(test_type_t const &) const; + template + typename std::enable_if::value, bool>::type + operator==(test_type_t const &) const; + + template + typename std::enable_if::value, bool>::type + operator!=(test_type_t const &) const; + + template + typename std::enable_if::value, bool>::type + operator<(test_type_t const &) const; - typename std::enable_if::value, bool>::type - operator!=(test_type_t const &) const; + template + typename std::enable_if::value, bool>::type + operator>(test_type_t const &) const; - typename std::enable_if::value, bool>::type - operator<(test_type_t const &) const; + template + typename std::enable_if::value, bool>::type + operator<=(test_type_t const &) const; - typename std::enable_if::value, bool>::type - operator>(test_type_t const &) const; + template + typename std::enable_if::value, bool>::type + operator>=(test_type_t const &) const; - 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 &); - 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 &); }; 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 namespace std { -template <::FlexFlow::test_types::capability ...CAPABILITIES> //, typename = typename std::enable_if<::FlexFlow::test_types::has_capability<::FlexFlow::test_types::HASHABLE>::value, bool>::type> +template < + ::FlexFlow::test_types:: + capability... CAPABILITIES> //, typename = typename + // std::enable_if<::FlexFlow::test_types::has_capability<::FlexFlow::test_types::HASHABLE>::value, + // bool>::type> struct hash<::FlexFlow::test_types::test_type_t> { typename std::enable_if< - ::FlexFlow::test_types::has_capability<::FlexFlow::test_types::HASHABLE, CAPABILITIES...>::value, - size_t - >::type - operator()(::FlexFlow::test_types::test_type_t const &) const; + ::FlexFlow::test_types::has_capability<::FlexFlow::test_types::HASHABLE, + CAPABILITIES...>::value, + size_t>::type + operator()( + ::FlexFlow::test_types::test_type_t const &) const; }; -} +} // namespace std #endif diff --git a/lib/utils/include/utils/type_traits.h b/lib/utils/include/utils/type_traits.h index 45389a03e3..92f4ff1ff0 100644 --- a/lib/utils/include/utils/type_traits.h +++ b/lib/utils/include/utils/type_traits.h @@ -57,7 +57,6 @@ 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 fb5fd62a0c..4849456731 100644 --- a/lib/utils/include/utils/type_traits_core.h +++ b/lib/utils/include/utils/type_traits_core.h @@ -1,9 +1,9 @@ #ifndef _FLEXFLOW_UTILS_INCLUDE_UTILS_TYPE_TRAITS_CORE_H #define _FLEXFLOW_UTILS_INCLUDE_UTILS_TYPE_TRAITS_CORE_H -#include #include #include +#include namespace FlexFlow { template @@ -95,11 +95,11 @@ struct pack_contains_type : std::false_type {}; template struct pack_get { - using type = typename pack_get<(i-1), Args...>::type; + using type = typename pack_get<(i - 1), Args...>::type; }; template -struct pack_get<0, Head, Tail...> { +struct pack_get<0, Head, Tail...> { using type = Head; }; @@ -111,11 +111,12 @@ struct pack_get { template struct pack_size; -template -struct pack_size : std::integral_constant::value + 1)> { }; +template +struct pack_size + : std::integral_constant::value + 1)> {}; template <> -struct pack_size<> : std::integral_constant { }; +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/variant.h b/lib/utils/include/utils/variant.h index 61758c79d9..a03a781c75 100644 --- a/lib/utils/include/utils/variant.h +++ b/lib/utils/include/utils/variant.h @@ -174,9 +174,8 @@ template < typename Container, typename VariantIn = typename Container::value_type, typename = std::enable_if::value>> -auto widen(Container const &c) - -> decltype(transform(c, std::declval>())) -{ +auto widen(Container const &c) -> decltype(transform( + c, std::declval>())) { return transform(c, [](VariantIn const &i) { return widen(i); }); } @@ -193,13 +192,12 @@ template < typename Container, typename VariantIn = typename Container::value_type, typename = std::enable_if::value>> -auto narrow(Container const &c) - -> decltype(transform(c, std::declval(VariantIn const &)>>())) -{ +auto narrow(Container const &c) -> decltype(transform( + c, + std::declval(VariantIn const &)>>())) { return transform(c, [](VariantIn const &i) { return narrow(i); }); } - template ::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/lib/utils/src/graph/algorithms.cc b/lib/utils/src/graph/algorithms.cc index 5d5b38e81c..28bfaa77b7 100644 --- a/lib/utils/src/graph/algorithms.cc +++ b/lib/utils/src/graph/algorithms.cc @@ -115,9 +115,7 @@ bool contains_edge(DiGraphView const &g, DirectedEdge const &e) { } bool contains_edge(UndirectedGraphView const &g, UndirectedEdge const &e) { - UndirectedEdgeQuery q = { - {e.bigger, e.smaller} - }; + UndirectedEdgeQuery q = {{e.bigger, e.smaller}}; return contains(g.query_edges(q), e); } @@ -475,10 +473,18 @@ std::unordered_set get_cut(OpenMultiDiGraphView const &g, return keys(get_edge_splits(g, s)); } -Node get_src_node(MultiDiEdge const &) { NOT_IMPLEMENTED(); } -Node get_dst_node(MultiDiEdge const &) { NOT_IMPLEMENTED(); } -Node get_src_node(InputMultiDiEdge const &) { NOT_IMPLEMENTED(); } -Node get_dst_node(OutputMultiDiEdge const &) { NOT_IMPLEMENTED(); } +Node get_src_node(MultiDiEdge const &) { + NOT_IMPLEMENTED(); +} +Node get_dst_node(MultiDiEdge const &) { + NOT_IMPLEMENTED(); +} +Node get_src_node(InputMultiDiEdge const &) { + NOT_IMPLEMENTED(); +} +Node get_dst_node(OutputMultiDiEdge const &) { + NOT_IMPLEMENTED(); +} UndirectedGraphView get_subgraph(UndirectedGraphView const &g, std::unordered_set const &nodes) { diff --git a/python/flexflow_c.cc b/python/flexflow_c.cc index 4c16dd82e9..1d3d77afe5 100644 --- a/python/flexflow_c.cc +++ b/python/flexflow_c.cc @@ -33,7 +33,9 @@ class FFCObjectWrapper { t_.impl = const_cast(static_cast(t)); \ return t_; \ } \ - static T unwrap(T_ t_) { return static_cast(t_.impl); } \ + static T unwrap(T_ t_) { \ + return static_cast(t_.impl); \ + } \ static const T unwrap_const(const T_ t_) { \ return static_cast(t_.impl); \ } diff --git a/scripts/format.sh b/scripts/format.sh index 2ed97b8f0a..13a6531b93 100755 --- a/scripts/format.sh +++ b/scripts/format.sh @@ -6,8 +6,8 @@ GIT_ROOT="$(git rev-parse --show-toplevel)" cd "$GIT_ROOT" TOOLS_PATH="$GIT_ROOT/.tools" -RELEASE="master-1d7ec53d" -CLANG_FORMAT_VERSION="15" +RELEASE="master-f4f85437" +CLANG_FORMAT_VERSION="16" CLANG_FORMAT_PATH="$TOOLS_PATH/clang-format-$CLANG_FORMAT_VERSION-$RELEASE" mkdir -p "$TOOLS_PATH" @@ -68,5 +68,17 @@ if [[ ! -e $CLANG_FORMAT_PATH ]]; then chmod u+x "$CLANG_FORMAT_PATH" fi -mapfile -t FILES < <(git ls-files ':!:triton/**' '*.h' '*.cc' '*.cpp' '*.cu' '*.c') -"$CLANG_FORMAT_PATH" -i "${FILES[@]}" + +CLANG_FORMAT_CONFIG="$GIT_ROOT/.clang-format-for-format-sh" +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 + echo "error" +fi From 366d9b77b438150f8f7c7e19b98add96f022633e Mon Sep 17 00:00:00 2001 From: Colin Unger Date: Fri, 18 Aug 2023 04:21:34 -0700 Subject: [PATCH 4/4] ComputationGraphBuilder functions fixes --- cmake/nccl.cmake | 2 +- lib/op-attrs/include/op-attrs/ff_dim.h | 1 - lib/pcg/src/computation_graph_builder.cc | 72 +++++++++++++++--------- 3 files changed, 47 insertions(+), 28 deletions(-) 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/op-attrs/include/op-attrs/ff_dim.h b/lib/op-attrs/include/op-attrs/ff_dim.h index ee1e54194f..298e134e08 100644 --- a/lib/op-attrs/include/op-attrs/ff_dim.h +++ b/lib/op-attrs/include/op-attrs/ff_dim.h @@ -15,6 +15,5 @@ struct ff_dim_t : public numerical_typedef { MAKE_TYPEDEF_HASHABLE(::FlexFlow::ff_dim_t); MAKE_TYPEDEF_PRINTABLE(::FlexFlow::ff_dim_t, "ff_dim"); -static_assert(FlexFlow::is_neq_comparable<::FlexFlow::ff_dim_t>::value); #endif diff --git a/lib/pcg/src/computation_graph_builder.cc b/lib/pcg/src/computation_graph_builder.cc index e2c070a8c0..04c713141a 100644 --- a/lib/pcg/src/computation_graph_builder.cc +++ b/lib/pcg/src/computation_graph_builder.cc @@ -1,6 +1,9 @@ #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" @@ -10,9 +13,29 @@ #include "utils/expected.h" #include "utils/fmt.h" #include "utils/graph/multidiedge.h" +#include "op-attrs/get_output_shapes.h" namespace FlexFlow { +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, @@ -110,8 +133,6 @@ static tensor_guid_t insert_layer( cg, add_layer(cg, layer, inputs, weights, {output_shape})); } -static TensorShape get_broadcast_target_shape(std::vector const &); - static tensor_guid_t element_binary(ComputationGraph &, OperatorType, @@ -175,28 +196,28 @@ static tensor_guid_t insert_element_unary_layer( 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 insert_layer(cg, layer, {input}, {}, output_shape); } static tensor_guid_t - insert_element_unary_layer(ComputationGraph &, + insert_element_unary_layer(ComputationGraph &cg, OperatorType op_type, tensor_guid_t const &input, optional const &name) { ElementUnaryAttrs attrs = {op_type}; - return insert_element_unary_layer(attrs, input, name); + return insert_element_unary_layer(cg, attrs, input, name); } static tensor_guid_t - insert_element_scalar_unary_layer(ComputationGraph &, + 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 insert_element_unary_layer(attrs, input, name); + return insert_element_unary_layer(cg, attrs, input, name); } static tensor_guid_t @@ -207,18 +228,18 @@ static tensor_guid_t optional const &maybe_name) { std::string name = maybe_name.value_or(get_default_name(op_type)); - TensorShape compute_shape = get_broadcast_target_shape({lhs, rhs}); - DataType compute_type = std::max(get_data_type(lhs), get_data_type(rhs)); + 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_guid_t const lhs_input = as_type( - broadcast(lhs, compute_shape), compute_type, name + "_inputl_pre_cast"); - tensor_guid_t const rhs_input = as_type( - 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 insert_layer(cg, layer, {lhs_input, rhs_input}, {}, output_shape); } @@ -394,14 +415,14 @@ tensor_guid_t 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 insert_layer(cg, layer, {input}, weights, output_shape); @@ -419,7 +440,7 @@ tensor_guid_t insert_dropout_layer(ComputationGraph &cg, 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 insert_layer(cg, layer, {input}, {}, output_shape); } @@ -437,10 +458,10 @@ tensor_guid_t std::string name = maybe_name.value_or(get_default_name(attrs)); Layer layer = {attrs, name}; - tensor_guid_t input = 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 insert_layer( cg, layer, {input}, {{weights_shape, kernel_initializer}}, output_shape); @@ -455,8 +476,7 @@ std::vector GatherAttrs attrs = {dim}; std::string name = maybe_name.value_or(get_default_name(attrs)); - Tensor index_tensor = cg.at(index); - DataType index_dt = get_data_type(index_tensor); + DataType index_dt = get_data_type(cg.at(index)); Layer layer = {attrs, name}; if (index_dt != DataType::INT32 && index_dt != DataType::INT64) { @@ -467,7 +487,7 @@ std::vector DataType::INT64); } std::vector output_shapes = - get_output_shapes(attrs, input, index_tensor); + get_output_shapes(cg, attrs, input, index); return insert_layer(cg, layer, {input, index}, {}, output_shapes); } @@ -483,7 +503,7 @@ tensor_guid_t float lambda_bal, optional const &maybe_name) { auto get_shape = [&](tensor_guid_t const &t) { - return get_data_type(cg.at(t)); + return cg.at(t).shape; }; AggregateAttrs attrs = {n, lambda_bal}; @@ -512,7 +532,7 @@ tensor_guid_t insert_batch_norm_layer(ComputationGraph &cg, 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 insert_layer(cg, layer, {input}, {}, output_shape); }