From 756b4ad898a471f6714ee0c61af47c0649bfdb2b Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Tue, 29 Sep 2026 14:58:21 -0400 Subject: [PATCH 1/3] Add NASNet-a model --- lib/models/include/models/nasnet/index.dox | 11 + lib/models/include/models/nasnet/nasnet.h | 31 + .../nasnet/nasnet_a_large_config.dtg.toml | 44 ++ lib/models/index.dox | 1 + lib/models/src/models/nasnet/nasnet.cc | 697 ++++++++++++++++++ lib/models/test/src/models/nasnet/nasnet.cc | 141 ++++ 6 files changed, 925 insertions(+) create mode 100644 lib/models/include/models/nasnet/index.dox create mode 100644 lib/models/include/models/nasnet/nasnet.h create mode 100644 lib/models/include/models/nasnet/nasnet_a_large_config.dtg.toml create mode 100644 lib/models/src/models/nasnet/nasnet.cc create mode 100644 lib/models/test/src/models/nasnet/nasnet.cc diff --git a/lib/models/include/models/nasnet/index.dox b/lib/models/include/models/nasnet/index.dox new file mode 100644 index 0000000000..6d10a8143a --- /dev/null +++ b/lib/models/include/models/nasnet/index.dox @@ -0,0 +1,11 @@ +namespace FlexFlow { +/** + +\page models-nasnet models/nasnet + +NASNet-A Large image classification model based on the implementation in + +timm. + +*/ +} diff --git a/lib/models/include/models/nasnet/nasnet.h b/lib/models/include/models/nasnet/nasnet.h new file mode 100644 index 0000000000..f8d58e2590 --- /dev/null +++ b/lib/models/include/models/nasnet/nasnet.h @@ -0,0 +1,31 @@ +/** + * @file nasnet.h + * + * @brief NASNet model computation graphs. + */ + +#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_NASNET_NASNET_H +#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_NASNET_NASNET_H + +#include "models/nasnet/nasnet_a_large_config.dtg.h" +#include "pcg/computation_graph.dtg.h" + +namespace FlexFlow { + +/** + * @brief Get the default NASNet-A Large configuration. + * + * The defaults match the timm NASNet-A Large definition: RGB 331x331 input, + * 1000 output classes, 4032 final features, and no classifier dropout. + */ +NASNetALargeConfig get_default_nasnet_a_large_config(); + +/** + * @brief Build a NASNet-A Large computation graph. + */ +ComputationGraph + get_nasnet_a_large_computation_graph(NASNetALargeConfig const &config); + +} // namespace FlexFlow + +#endif diff --git a/lib/models/include/models/nasnet/nasnet_a_large_config.dtg.toml b/lib/models/include/models/nasnet/nasnet_a_large_config.dtg.toml new file mode 100644 index 0000000000..7efa0249e6 --- /dev/null +++ b/lib/models/include/models/nasnet/nasnet_a_large_config.dtg.toml @@ -0,0 +1,44 @@ +namespace = "FlexFlow" +name = "NASNetALargeConfig" +type = "struct" + +features = [ + "eq", + "ord", + "hash", + "json", + "rapidcheck", + "fmt", +] + +includes = [ + "utils/positive_int/positive_int.h", +] + +docstring = ''' +\brief Configuration for a NASNet-A Large image classifier. +''' + +[[fields]] +name = "batch_size" +type = "::FlexFlow::positive_int" + +[[fields]] +name = "num_classes" +type = "::FlexFlow::positive_int" + +[[fields]] +name = "input_channels" +type = "::FlexFlow::positive_int" + +[[fields]] +name = "image_height" +type = "::FlexFlow::positive_int" + +[[fields]] +name = "image_width" +type = "::FlexFlow::positive_int" + +[[fields]] +name = "dropout" +type = "float" diff --git a/lib/models/index.dox b/lib/models/index.dox index 8046ef125f..987c3ae4f9 100644 --- a/lib/models/index.dox +++ b/lib/models/index.dox @@ -10,6 +10,7 @@ namespace FlexFlow { - \subpage models-candle-uno "Candle UNO" - \subpage models-dlrm "DLRM" - \subpage models-inception "Inception v3" +- \subpage models-nasnet "NASNet-A Large" - \subpage models-transformer "Transformer" \section test-models Artificial Models for Testing diff --git a/lib/models/src/models/nasnet/nasnet.cc b/lib/models/src/models/nasnet/nasnet.cc new file mode 100644 index 0000000000..25249bff9d --- /dev/null +++ b/lib/models/src/models/nasnet/nasnet.cc @@ -0,0 +1,697 @@ +/** + * @file nasnet.cc + * @brief Implement NASNet-A Large. + */ + +#include "models/nasnet/nasnet.h" +#include "op-attrs/tensor_dims.h" +#include "op-attrs/tensor_shape.h" +#include "pcg/computation_graph_builder.h" + +#include +#include +#include + +namespace FlexFlow { + +namespace { + +static nonnegative_int same_padding(positive_int kernel_size) { + return nonnegative_int{kernel_size.int_from_positive_int() / 2}; +} + +static positive_int get_num_channels(ComputationGraphBuilder const &cgb, + tensor_guid_t const &input) { + return dim_at_idx(cgb.get_shape(input).dims, relative_ff_dim_t{1}); +} + +static tensor_guid_t create_same_pool2d(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + positive_int kernel_size, + positive_int stride, + PoolOp type, + std::string const &name) { + // TODO: timm uses dynamically computed asymmetric SAME padding when stride + // is greater than one. ComputationGraphBuilder currently exposes only + // symmetric padding, which preserves the NASNet tensor shapes but differs at + // boundary values for even-sized inputs. + nonnegative_int padding = same_padding(kernel_size); + return cgb.pool2d(input, + /*kernelH=*/kernel_size, + /*kernelW=*/kernel_size, + /*strideH=*/stride, + /*strideW=*/stride, + /*paddingH=*/padding, + /*paddingW=*/padding, + /*type=*/type, + /*activation=*/std::nullopt, + /*name=*/name); +} + +static tensor_guid_t create_act_conv_bn(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + positive_int out_channels, + std::string const &name) { + tensor_guid_t x = cgb.relu(input, name + ".act"); + x = cgb.conv2d(x, + /*outChannels=*/out_channels, + /*kernelH=*/1_p, + /*kernelW=*/1_p, + /*strideH=*/1_p, + /*strideW=*/1_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*activation=*/std::nullopt, + /*groups=*/1_p, + /*use_bias=*/false, + /*kernel_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*kernel_regularizer=*/std::nullopt, + /*name=*/name + ".conv"); + return cgb.batch_norm(x, + /*affine=*/true, + /*activation=*/std::nullopt, + /*eps=*/0.001f, + /*momentum=*/0.1f, + /*name=*/name + ".bn"); +} + +static tensor_guid_t create_separable_conv2d(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + positive_int out_channels, + positive_int kernel_size, + positive_int stride, + std::string const &name) { + positive_int in_channels = get_num_channels(cgb, input); + nonnegative_int padding = same_padding(kernel_size); + + // TODO: use asymmetric SAME padding for stride-two depthwise convolutions + // once the graph builder can represent it. Symmetric padding gives the same + // output dimensions for this architecture. + tensor_guid_t x = cgb.conv2d(input, + /*outChannels=*/in_channels, + /*kernelH=*/kernel_size, + /*kernelW=*/kernel_size, + /*strideH=*/stride, + /*strideW=*/stride, + /*paddingH=*/padding, + /*paddingW=*/padding, + /*activation=*/std::nullopt, + /*groups=*/in_channels, + /*use_bias=*/false, + /*kernel_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*kernel_regularizer=*/std::nullopt, + /*name=*/name + ".depthwise"); + return cgb.conv2d(x, + /*outChannels=*/out_channels, + /*kernelH=*/1_p, + /*kernelW=*/1_p, + /*strideH=*/1_p, + /*strideW=*/1_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*activation=*/std::nullopt, + /*groups=*/1_p, + /*use_bias=*/false, + /*kernel_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*kernel_regularizer=*/std::nullopt, + /*name=*/name + ".pointwise"); +} + +static tensor_guid_t create_branch_separables(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + positive_int out_channels, + positive_int kernel_size, + positive_int stride, + bool stem_cell, + std::string const &name) { + positive_int middle_channels = + stem_cell ? out_channels : get_num_channels(cgb, input); + + tensor_guid_t x = cgb.relu(input, name + ".act_1"); + x = create_separable_conv2d( + cgb, x, middle_channels, kernel_size, stride, name + ".separable_1"); + x = cgb.batch_norm(x, + /*affine=*/true, + /*activation=*/std::nullopt, + /*eps=*/0.001f, + /*momentum=*/0.1f, + /*name=*/name + ".bn_sep_1"); + x = cgb.relu(x, name + ".act_2"); + x = create_separable_conv2d(cgb, + x, + out_channels, + kernel_size, + /*stride=*/1_p, + name + ".separable_2"); + return cgb.batch_norm(x, + /*affine=*/true, + /*activation=*/std::nullopt, + /*eps=*/0.001f, + /*momentum=*/0.1f, + /*name=*/name + ".bn_sep_2"); +} + +static tensor_guid_t + create_factorized_reduction(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + positive_int out_channels_per_path, + std::string const &name) { + tensor_guid_t activated = cgb.relu(input, name + ".act"); + + tensor_guid_t path_1 = cgb.pool2d(activated, + /*kernelH=*/1_p, + /*kernelW=*/1_p, + /*strideH=*/2_p, + /*strideW=*/2_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*type=*/PoolOp::AVG, + /*activation=*/std::nullopt, + /*name=*/name + ".path_1.avgpool"); + path_1 = cgb.conv2d(path_1, + /*outChannels=*/out_channels_per_path, + /*kernelH=*/1_p, + /*kernelW=*/1_p, + /*strideH=*/1_p, + /*strideW=*/1_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*activation=*/std::nullopt, + /*groups=*/1_p, + /*use_bias=*/false, + /*kernel_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*kernel_regularizer=*/std::nullopt, + /*name=*/name + ".path_1.conv"); + + // TODO: NASNet shifts path_2 by cropping one pixel from the top and left and + // padding the bottom and right before pooling. FlexFlow does not currently + // expose tensor padding/cropping, so this path temporarily uses the + // unshifted activation. It has the correct shape but samples the same spatial + // positions as path_1. + tensor_guid_t path_2 = cgb.pool2d(activated, + /*kernelH=*/1_p, + /*kernelW=*/1_p, + /*strideH=*/2_p, + /*strideW=*/2_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*type=*/PoolOp::AVG, + /*activation=*/std::nullopt, + /*name=*/name + ".path_2.avgpool"); + path_2 = cgb.conv2d(path_2, + /*outChannels=*/out_channels_per_path, + /*kernelH=*/1_p, + /*kernelW=*/1_p, + /*strideH=*/1_p, + /*strideW=*/1_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*activation=*/std::nullopt, + /*groups=*/1_p, + /*use_bias=*/false, + /*kernel_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*kernel_regularizer=*/std::nullopt, + /*name=*/name + ".path_2.conv"); + + tensor_guid_t output = cgb.concat({path_1, path_2}, + /*axis=*/relative_ff_dim_t{1}, + /*name=*/name + ".concat"); + return cgb.batch_norm(output, + /*affine=*/true, + /*activation=*/std::nullopt, + /*eps=*/0.001f, + /*momentum=*/0.1f, + /*name=*/name + ".final_path_bn"); +} + +static tensor_guid_t create_nasnet_stem_cell_0(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + positive_int channels, + std::string const &name) { + tensor_guid_t x1 = + create_act_conv_bn(cgb, input, channels, name + ".conv_1x1"); + + tensor_guid_t iter_0 = + cgb.add(create_branch_separables(cgb, + x1, + channels, + /*kernel_size=*/5_p, + /*stride=*/2_p, + /*stem_cell=*/false, + name + ".comb_iter_0_left"), + create_branch_separables(cgb, + input, + channels, + /*kernel_size=*/7_p, + /*stride=*/2_p, + /*stem_cell=*/true, + name + ".comb_iter_0_right"), + name + ".comb_iter_0"); + + tensor_guid_t iter_1 = + cgb.add(create_same_pool2d(cgb, + x1, + /*kernel_size=*/3_p, + /*stride=*/2_p, + PoolOp::MAX, + name + ".comb_iter_1_left"), + create_branch_separables(cgb, + input, + channels, + /*kernel_size=*/7_p, + /*stride=*/2_p, + /*stem_cell=*/true, + name + ".comb_iter_1_right"), + name + ".comb_iter_1"); + + tensor_guid_t iter_2 = + cgb.add(create_same_pool2d(cgb, + x1, + /*kernel_size=*/3_p, + /*stride=*/2_p, + PoolOp::AVG, + name + ".comb_iter_2_left"), + create_branch_separables(cgb, + input, + channels, + /*kernel_size=*/5_p, + /*stride=*/2_p, + /*stem_cell=*/true, + name + ".comb_iter_2_right"), + name + ".comb_iter_2"); + + tensor_guid_t iter_3 = + cgb.add(create_same_pool2d(cgb, + iter_0, + /*kernel_size=*/3_p, + /*stride=*/1_p, + PoolOp::AVG, + name + ".comb_iter_3_right"), + iter_1, + name + ".comb_iter_3"); + + tensor_guid_t iter_4 = + cgb.add(create_branch_separables(cgb, + iter_0, + channels, + /*kernel_size=*/3_p, + /*stride=*/1_p, + /*stem_cell=*/false, + name + ".comb_iter_4_left"), + create_same_pool2d(cgb, + x1, + /*kernel_size=*/3_p, + /*stride=*/2_p, + PoolOp::MAX, + name + ".comb_iter_4_right"), + name + ".comb_iter_4"); + + return cgb.concat({iter_1, iter_2, iter_3, iter_4}, + /*axis=*/relative_ff_dim_t{1}, + /*name=*/name + ".output"); +} + +static tensor_guid_t create_nasnet_stem_cell_1(ComputationGraphBuilder &cgb, + tensor_guid_t const &x_conv0, + tensor_guid_t const &x_stem_0, + positive_int channels, + std::string const &name) { + tensor_guid_t x_left = + create_act_conv_bn(cgb, x_stem_0, channels, name + ".conv_1x1"); + tensor_guid_t x_right = create_factorized_reduction( + cgb, + x_conv0, + positive_int{channels.int_from_positive_int() / 2}, + name + ".previous_reduction"); + + tensor_guid_t iter_0 = + cgb.add(create_branch_separables(cgb, + x_left, + channels, + /*kernel_size=*/5_p, + /*stride=*/2_p, + /*stem_cell=*/false, + name + ".comb_iter_0_left"), + create_branch_separables(cgb, + x_right, + channels, + /*kernel_size=*/7_p, + /*stride=*/2_p, + /*stem_cell=*/false, + name + ".comb_iter_0_right"), + name + ".comb_iter_0"); + + tensor_guid_t iter_1 = + cgb.add(create_same_pool2d(cgb, + x_left, + /*kernel_size=*/3_p, + /*stride=*/2_p, + PoolOp::MAX, + name + ".comb_iter_1_left"), + create_branch_separables(cgb, + x_right, + channels, + /*kernel_size=*/7_p, + /*stride=*/2_p, + /*stem_cell=*/false, + name + ".comb_iter_1_right"), + name + ".comb_iter_1"); + + tensor_guid_t iter_2 = + cgb.add(create_same_pool2d(cgb, + x_left, + /*kernel_size=*/3_p, + /*stride=*/2_p, + PoolOp::AVG, + name + ".comb_iter_2_left"), + create_branch_separables(cgb, + x_right, + channels, + /*kernel_size=*/5_p, + /*stride=*/2_p, + /*stem_cell=*/false, + name + ".comb_iter_2_right"), + name + ".comb_iter_2"); + + tensor_guid_t iter_3 = + cgb.add(create_same_pool2d(cgb, + iter_0, + /*kernel_size=*/3_p, + /*stride=*/1_p, + PoolOp::AVG, + name + ".comb_iter_3_right"), + iter_1, + name + ".comb_iter_3"); + + tensor_guid_t iter_4 = + cgb.add(create_branch_separables(cgb, + iter_0, + channels, + /*kernel_size=*/3_p, + /*stride=*/1_p, + /*stem_cell=*/false, + name + ".comb_iter_4_left"), + create_same_pool2d(cgb, + x_left, + /*kernel_size=*/3_p, + /*stride=*/2_p, + PoolOp::MAX, + name + ".comb_iter_4_right"), + name + ".comb_iter_4"); + + return cgb.concat({iter_1, iter_2, iter_3, iter_4}, + /*axis=*/relative_ff_dim_t{1}, + /*name=*/name + ".output"); +} + +static tensor_guid_t + create_nasnet_first_cell(ComputationGraphBuilder &cgb, + tensor_guid_t const &x, + tensor_guid_t const &x_prev, + positive_int out_channels_per_left_path, + positive_int out_channels, + std::string const &name) { + tensor_guid_t x_left = create_factorized_reduction( + cgb, x_prev, out_channels_per_left_path, name + ".previous_reduction"); + tensor_guid_t x_right = + create_act_conv_bn(cgb, x, out_channels, name + ".conv_1x1"); + + tensor_guid_t iter_0 = + cgb.add(create_branch_separables(cgb, + x_right, + out_channels, + 5_p, + 1_p, + false, + name + ".comb_iter_0_left"), + create_branch_separables(cgb, + x_left, + out_channels, + 3_p, + 1_p, + false, + name + ".comb_iter_0_right"), + name + ".comb_iter_0"); + tensor_guid_t iter_1 = + cgb.add(create_branch_separables(cgb, + x_left, + out_channels, + 5_p, + 1_p, + false, + name + ".comb_iter_1_left"), + create_branch_separables(cgb, + x_left, + out_channels, + 3_p, + 1_p, + false, + name + ".comb_iter_1_right"), + name + ".comb_iter_1"); + tensor_guid_t iter_2 = cgb.add( + create_same_pool2d( + cgb, x_right, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_2_left"), + x_left, + name + ".comb_iter_2"); + tensor_guid_t iter_3 = cgb.add( + create_same_pool2d( + cgb, x_left, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_3_left"), + create_same_pool2d( + cgb, x_left, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_3_right"), + name + ".comb_iter_3"); + tensor_guid_t iter_4 = + cgb.add(create_branch_separables(cgb, + x_right, + out_channels, + 3_p, + 1_p, + false, + name + ".comb_iter_4_left"), + x_right, + name + ".comb_iter_4"); + + return cgb.concat({x_left, iter_0, iter_1, iter_2, iter_3, iter_4}, + relative_ff_dim_t{1}, + name + ".output"); +} + +static tensor_guid_t create_nasnet_normal_cell(ComputationGraphBuilder &cgb, + tensor_guid_t const &x, + tensor_guid_t const &x_prev, + positive_int channels, + std::string const &name) { + tensor_guid_t x_left = + create_act_conv_bn(cgb, x_prev, channels, name + ".conv_prev_1x1"); + tensor_guid_t x_right = + create_act_conv_bn(cgb, x, channels, name + ".conv_1x1"); + + tensor_guid_t iter_0 = cgb.add( + create_branch_separables( + cgb, x_right, channels, 5_p, 1_p, false, name + ".comb_iter_0_left"), + create_branch_separables( + cgb, x_left, channels, 3_p, 1_p, false, name + ".comb_iter_0_right"), + name + ".comb_iter_0"); + tensor_guid_t iter_1 = cgb.add( + create_branch_separables( + cgb, x_left, channels, 5_p, 1_p, false, name + ".comb_iter_1_left"), + create_branch_separables( + cgb, x_left, channels, 3_p, 1_p, false, name + ".comb_iter_1_right"), + name + ".comb_iter_1"); + tensor_guid_t iter_2 = cgb.add( + create_same_pool2d( + cgb, x_right, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_2_left"), + x_left, + name + ".comb_iter_2"); + tensor_guid_t iter_3 = cgb.add( + create_same_pool2d( + cgb, x_left, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_3_left"), + create_same_pool2d( + cgb, x_left, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_3_right"), + name + ".comb_iter_3"); + tensor_guid_t iter_4 = cgb.add( + create_branch_separables( + cgb, x_right, channels, 3_p, 1_p, false, name + ".comb_iter_4_left"), + x_right, + name + ".comb_iter_4"); + + return cgb.concat({x_left, iter_0, iter_1, iter_2, iter_3, iter_4}, + relative_ff_dim_t{1}, + name + ".output"); +} + +static tensor_guid_t create_nasnet_reduction_cell(ComputationGraphBuilder &cgb, + tensor_guid_t const &x, + tensor_guid_t const &x_prev, + positive_int channels, + std::string const &name) { + tensor_guid_t x_left = + create_act_conv_bn(cgb, x_prev, channels, name + ".conv_prev_1x1"); + tensor_guid_t x_right = + create_act_conv_bn(cgb, x, channels, name + ".conv_1x1"); + + tensor_guid_t iter_0 = cgb.add( + create_branch_separables( + cgb, x_right, channels, 5_p, 2_p, false, name + ".comb_iter_0_left"), + create_branch_separables( + cgb, x_left, channels, 7_p, 2_p, false, name + ".comb_iter_0_right"), + name + ".comb_iter_0"); + tensor_guid_t iter_1 = cgb.add( + create_same_pool2d( + cgb, x_right, 3_p, 2_p, PoolOp::MAX, name + ".comb_iter_1_left"), + create_branch_separables( + cgb, x_left, channels, 7_p, 2_p, false, name + ".comb_iter_1_right"), + name + ".comb_iter_1"); + tensor_guid_t iter_2 = cgb.add( + create_same_pool2d( + cgb, x_right, 3_p, 2_p, PoolOp::AVG, name + ".comb_iter_2_left"), + create_branch_separables( + cgb, x_left, channels, 5_p, 2_p, false, name + ".comb_iter_2_right"), + name + ".comb_iter_2"); + tensor_guid_t iter_3 = cgb.add( + create_same_pool2d( + cgb, iter_0, 3_p, 1_p, PoolOp::AVG, name + ".comb_iter_3_right"), + iter_1, + name + ".comb_iter_3"); + tensor_guid_t iter_4 = cgb.add( + create_branch_separables( + cgb, iter_0, channels, 3_p, 1_p, false, name + ".comb_iter_4_left"), + create_same_pool2d( + cgb, x_right, 3_p, 2_p, PoolOp::MAX, name + ".comb_iter_4_right"), + name + ".comb_iter_4"); + + return cgb.concat( + {iter_1, iter_2, iter_3, iter_4}, relative_ff_dim_t{1}, name + ".output"); +} + +} // namespace + +NASNetALargeConfig get_default_nasnet_a_large_config() { + return NASNetALargeConfig{ + /*batch_size=*/32_p, + /*num_classes=*/1000_p, + /*input_channels=*/3_p, + /*image_height=*/331_p, + /*image_width=*/331_p, + /*dropout=*/0.0f, + }; +} + +ComputationGraph + get_nasnet_a_large_computation_graph(NASNetALargeConfig const &config) { + if (!(config.dropout >= 0.0f && config.dropout < 1.0f)) { + throw mk_runtime_error(fmt::format( + "NASNet-A Large dropout must be in the range [0, 1), but found {}", + config.dropout)); + } + + ComputationGraphBuilder cgb; + + TensorShape input_shape = TensorShape{ + TensorDims{FFOrdered{config.batch_size, + config.input_channels, + config.image_height, + config.image_width}}, + DataType::FLOAT, + }; + tensor_guid_t input = cgb.create_input(input_shape, CreateGrad::YES, "input"); + + tensor_guid_t x_conv0 = cgb.conv2d(input, + /*outChannels=*/96_p, + /*kernelH=*/3_p, + /*kernelW=*/3_p, + /*strideH=*/2_p, + /*strideW=*/2_p, + /*paddingH=*/0_n, + /*paddingW=*/0_n, + /*activation=*/std::nullopt, + /*groups=*/1_p, + /*use_bias=*/false, + /*kernel_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*kernel_regularizer=*/std::nullopt, + /*name=*/"conv0.conv"); + x_conv0 = cgb.batch_norm(x_conv0, + /*affine=*/true, + /*activation=*/std::nullopt, + /*eps=*/0.001f, + /*momentum=*/0.1f, + /*name=*/"conv0.bn"); + + tensor_guid_t x_stem_0 = + create_nasnet_stem_cell_0(cgb, x_conv0, 42_p, "cell_stem_0"); + tensor_guid_t x_stem_1 = + create_nasnet_stem_cell_1(cgb, x_conv0, x_stem_0, 84_p, "cell_stem_1"); + + tensor_guid_t previous = x_stem_0; + tensor_guid_t current = + create_nasnet_first_cell(cgb, x_stem_1, previous, 84_p, 168_p, "cell_0"); + previous = x_stem_1; + + for (int idx = 1; idx <= 5; idx++) { + tensor_guid_t next = create_nasnet_normal_cell( + cgb, current, previous, 168_p, fmt::format("cell_{}", idx)); + previous = current; + current = next; + } + + tensor_guid_t reduction_0 = create_nasnet_reduction_cell( + cgb, current, previous, 336_p, "reduction_cell_0"); + current = create_nasnet_first_cell( + cgb, reduction_0, previous, 168_p, 336_p, "cell_6"); + previous = reduction_0; + + for (int idx = 7; idx <= 11; idx++) { + tensor_guid_t next = create_nasnet_normal_cell( + cgb, current, previous, 336_p, fmt::format("cell_{}", idx)); + previous = current; + current = next; + } + + tensor_guid_t reduction_1 = create_nasnet_reduction_cell( + cgb, current, previous, 672_p, "reduction_cell_1"); + current = create_nasnet_first_cell( + cgb, reduction_1, previous, 336_p, 672_p, "cell_12"); + previous = reduction_1; + + for (int idx = 13; idx <= 17; idx++) { + tensor_guid_t next = create_nasnet_normal_cell( + cgb, current, previous, 672_p, fmt::format("cell_{}", idx)); + previous = current; + current = next; + } + + tensor_guid_t output = cgb.relu(current, "act"); + output = cgb.adaptive_pool2d(output, + /*output_h=*/1_p, + /*output_w=*/1_p, + /*type=*/PoolOp::AVG, + /*activation=*/std::nullopt, + /*name=*/"global_pool"); + if (config.dropout > 0.0f) { + output = cgb.dropout(output, + /*rate=*/config.dropout, + /*seed=*/0, + /*name=*/"head_drop"); + } + output = cgb.flat(output, + /*start_dim=*/relative_ff_dim_t{1}, + /*end_dim=*/std::nullopt, + /*name=*/"flatten"); + output = cgb.dense(output, + /*outDim=*/config.num_classes, + /*activation=*/std::nullopt, + /*use_bias=*/true, + /*data_type=*/DataType::FLOAT, + /*projection_initializer=*/std::nullopt, + /*bias_initializer=*/std::nullopt, + /*name=*/"last_linear"); + + return cgb.computation_graph; +} + +} // namespace FlexFlow diff --git a/lib/models/test/src/models/nasnet/nasnet.cc b/lib/models/test/src/models/nasnet/nasnet.cc new file mode 100644 index 0000000000..d2e014c32e --- /dev/null +++ b/lib/models/test/src/models/nasnet/nasnet.cc @@ -0,0 +1,141 @@ +#include "models/nasnet/nasnet.h" +#include "op-attrs/ops/conv_2d_attrs.dtg.h" +#include "op-attrs/tensor_shape.h" +#include "pcg/computation_graph.h" + +#include +#include +#include +#include + +using namespace ::FlexFlow; + +namespace { + +TensorShape get_named_layer_output_shape(ComputationGraph const &graph, + std::string const &name) { + layer_guid_t layer = get_layer_by_name(graph, name); + tensor_guid_t output = + get_outgoing_tensors(graph, layer).at(TensorSlotName::OUTPUT); + return get_tensor_attrs(graph, output).shape; +} + +TensorShape make_4d_shape(positive_int batch_size, + positive_int channels, + positive_int height, + positive_int width) { + return TensorShape{ + TensorDims{FFOrdered{batch_size, channels, height, width}}, + DataType::FLOAT, + }; +} + +} // namespace + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("get_default_nasnet_a_large_config") { + NASNetALargeConfig config = get_default_nasnet_a_large_config(); + + CHECK(config.batch_size == 32_p); + CHECK(config.num_classes == 1000_p); + CHECK(config.input_channels == 3_p); + CHECK(config.image_height == 331_p); + CHECK(config.image_width == 331_p); + CHECK(config.dropout == 0.0f); + } + + TEST_CASE("get_nasnet_a_large_computation_graph") { + NASNetALargeConfig config = get_default_nasnet_a_large_config(); + ComputationGraph result = get_nasnet_a_large_computation_graph(config); + + std::map expected = { + {"conv0.bn", make_4d_shape(32_p, 96_p, 165_p, 165_p)}, + {"cell_stem_0.output", make_4d_shape(32_p, 168_p, 83_p, 83_p)}, + {"cell_stem_1.output", make_4d_shape(32_p, 336_p, 42_p, 42_p)}, + {"cell_5.output", make_4d_shape(32_p, 1008_p, 42_p, 42_p)}, + {"reduction_cell_0.output", make_4d_shape(32_p, 1344_p, 21_p, 21_p)}, + {"cell_11.output", make_4d_shape(32_p, 2016_p, 21_p, 21_p)}, + {"reduction_cell_1.output", make_4d_shape(32_p, 2688_p, 11_p, 11_p)}, + {"cell_17.output", make_4d_shape(32_p, 4032_p, 11_p, 11_p)}, + {"global_pool", make_4d_shape(32_p, 4032_p, 1_p, 1_p)}, + }; + + for (auto const &[name, shape] : expected) { + CHECK(get_named_layer_output_shape(result, name) == shape); + } + + std::map expected_operator_counts = { + {OperatorType::INPUT, 1_p}, + {OperatorType::WEIGHT, 1018_p}, + {OperatorType::CONV2D, 488_p}, + {OperatorType::LINEAR, 1_p}, + {OperatorType::POOL2D, 79_p}, + {OperatorType::RELU, 264_p}, + {OperatorType::FLAT, 1_p}, + {OperatorType::BATCHNORM, 264_p}, + {OperatorType::CONCAT, 26_p}, + {OperatorType::EW_ADD, 110_p}, + }; + CHECK(operator_type_counts_in_computation_graph(result) == + expected_operator_counts); + + layer_guid_t depthwise_layer = get_layer_by_name( + result, "cell_17.comb_iter_0_left.separable_1.depthwise"); + Conv2DAttrs depthwise_attrs = + get_layer_attrs(result, depthwise_layer).op_attrs.require_conv2d(); + CHECK(depthwise_attrs.groups == 672_p); + CHECK(depthwise_attrs.out_channels == 672_p); + CHECK(depthwise_attrs.kernel_h == 5_p); + CHECK(depthwise_attrs.kernel_w == 5_p); + + std::set unused_tensors = cg_get_unused_tensors(result); + REQUIRE(unused_tensors.size() == 1); + + TensorShape expected_output = TensorShape{ + TensorDims{FFOrdered{32_p, 1000_p}}, + DataType::FLOAT, + }; + CHECK(get_tensor_attrs(result, *unused_tensors.begin()).shape == + expected_output); + } + + TEST_CASE("get_nasnet_a_large_computation_graph honors custom config") { + NASNetALargeConfig config = get_default_nasnet_a_large_config(); + config.batch_size = 2_p; + config.num_classes = 17_p; + config.input_channels = 1_p; + config.image_height = 299_p; + config.image_width = 299_p; + config.dropout = 0.25f; + + ComputationGraph result = get_nasnet_a_large_computation_graph(config); + + CHECK(get_named_layer_output_shape(result, "input") == + make_4d_shape(2_p, 1_p, 299_p, 299_p)); + CHECK(get_named_layer_output_shape(result, "head_drop") == + make_4d_shape(2_p, 4032_p, 1_p, 1_p)); + + std::set unused_tensors = cg_get_unused_tensors(result); + REQUIRE(unused_tensors.size() == 1); + TensorShape expected_output = TensorShape{ + TensorDims{FFOrdered{2_p, 17_p}}, + DataType::FLOAT, + }; + CHECK(get_tensor_attrs(result, *unused_tensors.begin()).shape == + expected_output); + + std::map counts = + operator_type_counts_in_computation_graph(result); + CHECK(counts.at(OperatorType::DROPOUT) == 1_p); + } + + TEST_CASE("get_nasnet_a_large_computation_graph rejects invalid dropout") { + NASNetALargeConfig config = get_default_nasnet_a_large_config(); + + config.dropout = -0.1f; + CHECK_THROWS(get_nasnet_a_large_computation_graph(config)); + + config.dropout = 1.0f; + CHECK_THROWS(get_nasnet_a_large_computation_graph(config)); + } +} From 9210ea4c47e6377e11d8f7c13c9ad769c14c33ec Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Tue, 29 Sep 2026 15:19:28 -0400 Subject: [PATCH 2/3] Add to export-model-arch --- bin/export-model-arch/src/export-model-arch/main.cc | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/bin/export-model-arch/src/export-model-arch/main.cc b/bin/export-model-arch/src/export-model-arch/main.cc index b23f087f7e..ed5e746ec4 100644 --- a/bin/export-model-arch/src/export-model-arch/main.cc +++ b/bin/export-model-arch/src/export-model-arch/main.cc @@ -5,6 +5,7 @@ #include "models/candle_uno/candle_uno.h" #include "models/dlrm/dlrm.h" #include "models/inception_v3/inception_v3.h" +#include "models/nasnet/nasnet.h" #include "models/split_test/split_test.h" #include "models/transformer/transformer.h" #include "models/yolov10/yolov10.h" @@ -67,6 +68,9 @@ tl::expected } else if (model_name == "inception_v3") { return get_inception_v3_computation_graph( get_default_inception_v3_training_config()); + } else if (model_name == "nasnet_a_large") { + return get_nasnet_a_large_computation_graph( + get_default_nasnet_a_large_config()); } else if (model_name == "candle_uno") { return get_candle_uno_computation_graph(get_default_candle_uno_config()); } else if (model_name == "bert") { @@ -154,6 +158,7 @@ int main(int argc, char **argv) { std::vector model_options = { "transformer", "inception_v3", + "nasnet_a_large", "candle_uno", "bert", "dlrm", From e005f8303d7b30b0c27fb308e008ddcbc3f07dd9 Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Tue, 29 Sep 2026 18:17:16 -0400 Subject: [PATCH 3/3] Implement TODOs --- lib/models/src/models/nasnet/nasnet.cc | 210 ++++++++++++++++-- lib/models/test/src/models/nasnet/nasnet.cc | 5 +- lib/pcg/src/pcg/computation_graph_builder.cc | 2 +- .../test/src/pcg/computation_graph_builder.cc | 12 + lib/task-spec/src/task-spec/ops/impl/split.cc | 77 +++++-- 5 files changed, 257 insertions(+), 49 deletions(-) diff --git a/lib/models/src/models/nasnet/nasnet.cc b/lib/models/src/models/nasnet/nasnet.cc index 25249bff9d..71c0ec267b 100644 --- a/lib/models/src/models/nasnet/nasnet.cc +++ b/lib/models/src/models/nasnet/nasnet.cc @@ -8,16 +8,156 @@ #include "op-attrs/tensor_shape.h" #include "pcg/computation_graph_builder.h" +#include #include +#include #include +#include #include namespace FlexFlow { namespace { -static nonnegative_int same_padding(positive_int kernel_size) { - return nonnegative_int{kernel_size.int_from_positive_int() / 2}; +struct SpatialPadding { + int top; + int bottom; + int left; + int right; +}; + +static std::pair + get_same_padding_for_dim(int input_size, int kernel_size, int stride) { + int output_size = (input_size + stride - 1) / stride; + int total_padding = + std::max((output_size - 1) * stride + kernel_size - input_size, 0); + int padding_before = total_padding / 2; + return {padding_before, total_padding - padding_before}; +} + +static SpatialPadding get_same_padding(ComputationGraphBuilder const &cgb, + tensor_guid_t const &input, + positive_int kernel_size, + positive_int stride) { + TensorDims dims = cgb.get_shape(input).dims; + int input_height = + dim_at_idx(dims, relative_ff_dim_t{2}).int_from_positive_int(); + int input_width = + dim_at_idx(dims, relative_ff_dim_t{3}).int_from_positive_int(); + int kernel = kernel_size.int_from_positive_int(); + int stride_value = stride.int_from_positive_int(); + + auto [top, bottom] = + get_same_padding_for_dim(input_height, kernel, stride_value); + auto [left, right] = + get_same_padding_for_dim(input_width, kernel, stride_value); + return SpatialPadding{top, bottom, left, right}; +} + +static tensor_guid_t create_constant_like(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + float value, + std::string const &name) { + tensor_guid_t result = cgb.scalar_multiply(input, 0.0f, name + ".zero"); + if (value != 0.0f) { + result = cgb.scalar_add(result, value, name + ".value"); + } + return result; +} + +static tensor_guid_t + create_constant_padding_for_axis(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + relative_ff_dim_t axis, + int padding_before, + int padding_after, + float value, + std::string const &name) { + if (padding_before == 0 && padding_after == 0) { + return input; + } + + int axis_size = + dim_at_idx(cgb.get_shape(input).dims, axis).int_from_positive_int(); + ASSERT(padding_before >= 0 && padding_after >= 0); + ASSERT(padding_before + padding_after < axis_size); + + std::vector split_sizes; + if (padding_before > 0) { + split_sizes.push_back(positive_int{padding_before}); + } + split_sizes.push_back( + positive_int{axis_size - padding_before - padding_after}); + if (padding_after > 0) { + split_sizes.push_back(positive_int{padding_after}); + } + + std::vector pieces = + cgb.split(input, split_sizes, axis, name + ".split"); + std::vector concat_inputs; + if (padding_before > 0) { + concat_inputs.push_back( + create_constant_like(cgb, pieces.front(), value, name + ".before")); + } + concat_inputs.insert(concat_inputs.end(), pieces.begin(), pieces.end()); + if (padding_after > 0) { + concat_inputs.push_back( + create_constant_like(cgb, pieces.back(), value, name + ".after")); + } + return cgb.concat(concat_inputs, axis, name + ".concat"); +} + +static tensor_guid_t create_constant_padding2d(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + SpatialPadding const &padding, + float value, + std::string const &name) { + tensor_guid_t result = create_constant_padding_for_axis(cgb, + input, + relative_ff_dim_t{2}, + padding.top, + padding.bottom, + value, + name + ".height"); + return create_constant_padding_for_axis(cgb, + result, + relative_ff_dim_t{3}, + padding.left, + padding.right, + value, + name + ".width"); +} + +static tensor_guid_t + create_shifted_reduction_input(ComputationGraphBuilder &cgb, + tensor_guid_t const &input, + std::string const &name) { + TensorDims dims = cgb.get_shape(input).dims; + positive_int height = dim_at_idx(dims, relative_ff_dim_t{2}); + positive_int width = dim_at_idx(dims, relative_ff_dim_t{3}); + ASSERT(height > 1_p && width > 1_p); + + std::vector height_parts = + cgb.split(input, + {1_p, positive_int{height.int_from_positive_int() - 1}}, + relative_ff_dim_t{2}, + name + ".crop_top"); + tensor_guid_t zero_row = + create_constant_like(cgb, height_parts.at(0), 0.0f, name + ".pad_bottom"); + tensor_guid_t shifted_height = cgb.concat({height_parts.at(1), zero_row}, + relative_ff_dim_t{2}, + name + ".shift_height"); + + std::vector width_parts = + cgb.split(shifted_height, + {1_p, positive_int{width.int_from_positive_int() - 1}}, + relative_ff_dim_t{3}, + name + ".crop_left"); + tensor_guid_t zero_column = + create_constant_like(cgb, width_parts.at(0), 0.0f, name + ".pad_right"); + return cgb.concat({width_parts.at(1), zero_column}, + relative_ff_dim_t{3}, + name + ".shift_width"); } static positive_int get_num_channels(ComputationGraphBuilder const &cgb, @@ -31,18 +171,35 @@ static tensor_guid_t create_same_pool2d(ComputationGraphBuilder &cgb, positive_int stride, PoolOp type, std::string const &name) { - // TODO: timm uses dynamically computed asymmetric SAME padding when stride - // is greater than one. ComputationGraphBuilder currently exposes only - // symmetric padding, which preserves the NASNet tensor shapes but differs at - // boundary values for even-sized inputs. - nonnegative_int padding = same_padding(kernel_size); - return cgb.pool2d(input, + SpatialPadding padding = get_same_padding(cgb, input, kernel_size, stride); + bool is_symmetric = + padding.top == padding.bottom && padding.left == padding.right; + // timm materializes padding for every stride-two SAME pool. This matters for + // average pooling because those explicit zeros participate in the divisor; + // symmetric max-pool padding is equivalent to native padding. + bool use_native_padding = + is_symmetric && (stride == 1_p || type == PoolOp::MAX); + + tensor_guid_t padded_input = input; + nonnegative_int native_padding_h = 0_n; + nonnegative_int native_padding_w = 0_n; + if (use_native_padding) { + native_padding_h = nonnegative_int{padding.top}; + native_padding_w = nonnegative_int{padding.left}; + } else { + float padding_value = + type == PoolOp::MAX ? std::numeric_limits::lowest() : 0.0f; + padded_input = create_constant_padding2d( + cgb, input, padding, padding_value, name + ".same_pad"); + } + + return cgb.pool2d(padded_input, /*kernelH=*/kernel_size, /*kernelW=*/kernel_size, /*strideH=*/stride, /*strideW=*/stride, - /*paddingH=*/padding, - /*paddingW=*/padding, + /*paddingH=*/native_padding_h, + /*paddingW=*/native_padding_w, /*type=*/type, /*activation=*/std::nullopt, /*name=*/name); @@ -83,19 +240,29 @@ static tensor_guid_t create_separable_conv2d(ComputationGraphBuilder &cgb, positive_int stride, std::string const &name) { positive_int in_channels = get_num_channels(cgb, input); - nonnegative_int padding = same_padding(kernel_size); + SpatialPadding padding = get_same_padding(cgb, input, kernel_size, stride); + bool is_symmetric = + padding.top == padding.bottom && padding.left == padding.right; + + tensor_guid_t padded_input = input; + nonnegative_int native_padding_h = 0_n; + nonnegative_int native_padding_w = 0_n; + if (is_symmetric) { + native_padding_h = nonnegative_int{padding.top}; + native_padding_w = nonnegative_int{padding.left}; + } else { + padded_input = create_constant_padding2d( + cgb, input, padding, 0.0f, name + ".same_pad"); + } - // TODO: use asymmetric SAME padding for stride-two depthwise convolutions - // once the graph builder can represent it. Symmetric padding gives the same - // output dimensions for this architecture. - tensor_guid_t x = cgb.conv2d(input, + tensor_guid_t x = cgb.conv2d(padded_input, /*outChannels=*/in_channels, /*kernelH=*/kernel_size, /*kernelW=*/kernel_size, /*strideH=*/stride, /*strideW=*/stride, - /*paddingH=*/padding, - /*paddingW=*/padding, + /*paddingH=*/native_padding_h, + /*paddingW=*/native_padding_w, /*activation=*/std::nullopt, /*groups=*/in_channels, /*use_bias=*/false, @@ -187,12 +354,9 @@ static tensor_guid_t /*kernel_regularizer=*/std::nullopt, /*name=*/name + ".path_1.conv"); - // TODO: NASNet shifts path_2 by cropping one pixel from the top and left and - // padding the bottom and right before pooling. FlexFlow does not currently - // expose tensor padding/cropping, so this path temporarily uses the - // unshifted activation. It has the correct shape but samples the same spatial - // positions as path_1. - tensor_guid_t path_2 = cgb.pool2d(activated, + tensor_guid_t shifted = + create_shifted_reduction_input(cgb, activated, name + ".path_2.shift"); + tensor_guid_t path_2 = cgb.pool2d(shifted, /*kernelH=*/1_p, /*kernelW=*/1_p, /*strideH=*/2_p, diff --git a/lib/models/test/src/models/nasnet/nasnet.cc b/lib/models/test/src/models/nasnet/nasnet.cc index d2e014c32e..dea9f6bb07 100644 --- a/lib/models/test/src/models/nasnet/nasnet.cc +++ b/lib/models/test/src/models/nasnet/nasnet.cc @@ -73,8 +73,11 @@ TEST_SUITE(FF_TEST_SUITE) { {OperatorType::RELU, 264_p}, {OperatorType::FLAT, 1_p}, {OperatorType::BATCHNORM, 264_p}, - {OperatorType::CONCAT, 26_p}, + {OperatorType::CONCAT, 54_p}, {OperatorType::EW_ADD, 110_p}, + {OperatorType::SCALAR_ADD, 4_p}, + {OperatorType::SCALAR_MULTIPLY, 42_p}, + {OperatorType::SPLIT, 28_p}, }; CHECK(operator_type_counts_in_computation_graph(result) == expected_operator_counts); diff --git a/lib/pcg/src/pcg/computation_graph_builder.cc b/lib/pcg/src/pcg/computation_graph_builder.cc index eb94f1f901..216fa0a631 100644 --- a/lib/pcg/src/pcg/computation_graph_builder.cc +++ b/lib/pcg/src/pcg/computation_graph_builder.cc @@ -1071,7 +1071,7 @@ std::vector ComputationGraphBuilder::split( relative_ff_dim_t axis, std::optional const &maybe_name) { ASSERT(split.size() > 0); - ASSERT(split.size() <= get_variadic_inputs_slot_name_sequence().size()); + ASSERT(split.size() <= get_variadic_outputs_slot_name_sequence().size()); TensorShape input_shape = this->get_shape(input); diff --git a/lib/pcg/test/src/pcg/computation_graph_builder.cc b/lib/pcg/test/src/pcg/computation_graph_builder.cc index 54f61b2c46..638235601b 100644 --- a/lib/pcg/test/src/pcg/computation_graph_builder.cc +++ b/lib/pcg/test/src/pcg/computation_graph_builder.cc @@ -179,6 +179,18 @@ TEST_SUITE(FF_TEST_SUITE) { } } + TEST_CASE("ComputationGraphBuilder::split rejects too many outputs") { + ComputationGraphBuilder b; + TensorShape input_shape = TensorShape{ + TensorDims{FFOrdered{2_p, 3_p, 10_p, 10_p}}, + DataType::FLOAT, + }; + tensor_guid_t input = b.create_input(input_shape, CreateGrad::NO); + + CHECK_THROWS(b.split( + input, {1_p, 1_p, 1_p, 1_p, 1_p, 1_p, 4_p}, relative_ff_dim_t{2})); + } + TEST_CASE("ComputationGraphBuilder transpose") { ComputationGraphBuilder b; diff --git a/lib/task-spec/src/task-spec/ops/impl/split.cc b/lib/task-spec/src/task-spec/ops/impl/split.cc index 195f385c53..ae301e38bc 100644 --- a/lib/task-spec/src/task-spec/ops/impl/split.cc +++ b/lib/task-spec/src/task-spec/ops/impl/split.cc @@ -15,7 +15,10 @@ #include "task-spec/ops/impl/split.h" #include "kernels/split_kernels.h" +#include "op-attrs/tensor_slot_name.h" #include "task-spec/profiling.h" +#include "utils/containers/slice.h" +#include "utils/containers/transform.h" #include "utils/exception.h" #include "utils/hash-utils.h" #include "utils/nonnegative_int/nonnegative_range.h" @@ -24,6 +27,13 @@ namespace FlexFlow { using namespace FlexFlow::Kernels::Split; +static std::vector get_output_slots(SplitAttrs const &attrs) { + std::vector output_slots = + get_variadic_outputs_slot_name_sequence(); + ASSERT(attrs.splits.size() <= output_slots.size()); + return slice(output_slots, 0, attrs.splits.size()); +} + static std::pair calc_block_size(TensorShape const &tensor_shape, ff_dim_t axis) { positive_int num_blocks = 1_p; @@ -47,26 +57,36 @@ static std::optional SplitAttrs attrs = acc.get_op_attrs().require_split(); auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + std::vector outputs = + transform(get_output_slots(attrs), + [&](TensorSlotName slot) -> GenericTensorAccessorW { + return acc.get_tensor(slot); + }); + ASSERT(outputs.size() == attrs.splits.size()); + ASSERT(outputs.size() <= MAX_NUM_OUTPUTS); - int out_block_sizes[MAX_NUM_OUTPUTS]; auto [num_blocks, in_block_size] = calc_block_size(input.shape, attrs.axis); - - for (int i = 0; i < attrs.splits.size(); i++) { - auto [_, out_block_size] = calc_block_size(output.shape, attrs.axis); - out_block_sizes[i] = out_block_size.int_from_positive_int(); - } - float *output_float_ptr = output.get_float_ptr(); + std::vector out_block_sizes = + transform(outputs, [&](GenericTensorAccessorW const &output) -> int { + auto [output_num_blocks, out_block_size] = + calc_block_size(output.shape, attrs.axis); + ASSERT(output_num_blocks == num_blocks); + return out_block_size.int_from_positive_int(); + }); + std::vector output_ptrs = + transform(outputs, [](GenericTensorAccessorW const &output) -> float * { + return output.get_float_ptr(); + }); return profile(forward_kernel, profiling, kernel_device_type, "[Split] forward_time = {:.2lf}ms\n", - &output_float_ptr, + output_ptrs.data(), input.get_float_ptr(), - out_block_sizes, + out_block_sizes.data(), in_block_size.int_from_positive_int(), num_blocks.int_from_positive_int(), - attrs.splits.size()); + outputs.size()); } static std::optional @@ -76,29 +96,38 @@ static std::optional SplitAttrs attrs = acc.get_op_attrs().require_split(); auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output_grad = - acc.get_tensor_grad(TensorSlotName::OUTPUT); + std::vector output_grads = + transform(get_output_slots(attrs), + [&](TensorSlotName slot) -> GenericTensorAccessorR { + return acc.get_tensor_grad(slot); + }); + ASSERT(output_grads.size() == attrs.splits.size()); + ASSERT(output_grads.size() <= MAX_NUM_OUTPUTS); - int out_block_sizes[MAX_NUM_OUTPUTS]; auto [num_blocks, in_block_size] = calc_block_size(input_grad.shape, attrs.axis); - - for (int i = 0; i < attrs.splits.size(); i++) { - int out_num_blocks; - auto [_, out_block_size] = calc_block_size(output_grad.shape, attrs.axis); - out_block_sizes[i] = out_block_size.int_from_positive_int(); - } - float const *output_grad_ptr = output_grad.get_float_ptr(); + std::vector out_block_sizes = transform( + output_grads, [&](GenericTensorAccessorR const &output_grad) -> int { + auto [output_num_blocks, out_block_size] = + calc_block_size(output_grad.shape, attrs.axis); + ASSERT(output_num_blocks == num_blocks); + return out_block_size.int_from_positive_int(); + }); + std::vector output_grad_ptrs = + transform(output_grads, + [](GenericTensorAccessorR const &output_grad) -> float const * { + return output_grad.get_float_ptr(); + }); return profile(backward_kernel, profiling, kernel_device_type, "[Split] backward_time = {:.2lf}ms\n", input_grad.get_float_ptr(), - &output_grad_ptr, - out_block_sizes, + output_grad_ptrs.data(), + out_block_sizes.data(), in_block_size.int_from_positive_int(), num_blocks.int_from_positive_int(), - attrs.splits.size()); + output_grads.size()); } TaskImplFunction get_split_fwd_task_impl() {