teflon: Add support for Reshape operations

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/34629>
This commit is contained in:
Tomeu Vizoso 2024-12-11 17:28:05 +01:00 committed by Marge Bot
parent 1420f57ec7
commit dae0af20ab
2 changed files with 18 additions and 0 deletions

View file

@ -218,6 +218,13 @@ fill_operation(struct teflon_delegate *delegate, TfLiteContext *tf_context, TfLi
operation->fcon.bias_tensor = &tensors[node->inputs->data[2]];
break;
}
case kTfLiteBuiltinReshape: {
int32_t *shape = tf_context->tensors[node->inputs->data[1]].data.data;
operation->type = PIPE_ML_OPERATION_TYPE_RESHAPE;
memcpy(operation->reshape.shape, shape, 4 * sizeof(*operation->reshape.shape));
break;
}
default:
return false;
}
@ -348,6 +355,9 @@ dump_graph(struct pipe_tensor *tensors, unsigned tensor_count, struct pipe_ml_op
case PIPE_ML_OPERATION_TYPE_FULLY_CONNECTED:
teflon_debug("%-6s ", "FCON");
break;
case PIPE_ML_OPERATION_TYPE_RESHAPE:
teflon_debug("%-6s ", "RESHAPE");
break;
}
for (unsigned j = 0; j < operations[i].input_count; j++) {

View file

@ -1048,6 +1048,7 @@ enum pipe_ml_operation_type {
PIPE_ML_OPERATION_TYPE_SPLIT,
PIPE_ML_OPERATION_TYPE_PAD,
PIPE_ML_OPERATION_TYPE_FULLY_CONNECTED,
PIPE_ML_OPERATION_TYPE_RESHAPE,
};
/**
@ -1204,6 +1205,13 @@ struct pipe_ml_operation
*/
int axis;
} split;
struct {
/**
* Shape of the output tensor.
*/
unsigned shape[4];
} reshape;
};
};