Skip to content

Commit

Permalink
Inlined the helper function
Browse files Browse the repository at this point in the history
  • Loading branch information
RahulSundarMCW committed Jan 22, 2025
1 parent 9f04837 commit 9324d34
Showing 1 changed file with 59 additions and 5 deletions.
64 changes: 59 additions & 5 deletions src/subgraph/even-split.c
Original file line number Diff line number Diff line change
Expand Up @@ -342,8 +342,8 @@ enum xnn_status xnn_define_even_split_impl(
return xnn_status_invalid_parameter;
}

for(int i = 0; i < num_outputs; ++i){
check_datatype_copyable(subgraph, input_id, output_ids[i], "Nth", node_type);
for (int i = 0; i < num_outputs; ++i) {
check_datatype_copyable(subgraph, input_id, output_ids[i], "Nth", node_type);
}

struct xnn_node* node = xnn_subgraph_new_node(subgraph);
Expand All @@ -367,14 +367,68 @@ enum xnn_status xnn_define_even_split_impl(
return xnn_status_success;
};

inline enum xnn_status xnn_define_even_split(
enum xnn_status xnn_define_even_split(
xnn_subgraph_t subgraph,
int32_t split_dim,
uint32_t input_id,
uint32_t num_outputs,
const uint32_t* output_ids,
uint32_t flags)
{
return xnn_define_even_split_impl(
xnn_node_type_even_split, subgraph, split_dim, input_id, num_outputs, output_ids, flags);
assert(num_outputs > 1);
assert(num_outputs <= XNN_MAX_OUTPUTS);

enum xnn_node_type node_type = xnn_node_type_even_split;
enum xnn_status status;
if ((status = xnn_subgraph_check_xnnpack_initialized(node_type)) != xnn_status_success) {
return status;
}

if ((status = xnn_subgraph_check_input_node_id(node_type, input_id, subgraph->num_values)) != xnn_status_success) {
return status;
}

const struct xnn_value* input_value = &subgraph->values[input_id];
status = xnn_subgraph_check_input_type_dense(node_type, input_id, input_value);
if (status != xnn_status_success) {
return status;
}

for (int i = 0; i < num_outputs; ++i) {
status = check_output_value(subgraph, split_dim, input_id, output_ids[i], "Nth", node_type);
if (status != xnn_status_success) {
return status;
}
}

if (num_outputs > XNN_MAX_OUTPUTS) {
xnn_log_error(
"failed to define %s operator with %zu inputs: number of inputs (%zu) exceeds the supported maximum (%zu)",
xnn_node_type_to_string(node_type), num_outputs, num_outputs, (size_t) XNN_MAX_OUTPUTS);
return xnn_status_invalid_parameter;
}

for (int i = 0; i < num_outputs; ++i) {
check_datatype_copyable(subgraph, input_id, output_ids[i], "Nth", node_type);
}

struct xnn_node* node = xnn_subgraph_new_node(subgraph);
if (node == NULL) {
return xnn_status_out_of_memory;
}

node->params.even_split.axis = split_dim;
node->type = node_type;
node->num_inputs = 1;
node->inputs[0] = input_id;
node->num_outputs = num_outputs;
for(int i=0;i<num_outputs;++i){
node->outputs[i]=output_ids[i];
}
node->create = create_even_split_operator;
node->reshape = reshape_even_split_operator;
node->setup = setup_even_split_operator;
node->flags = flags;

return xnn_status_success;
}

0 comments on commit 9324d34

Please sign in to comment.