ethosu: Store ethosu_tensor struct ptr in feature map

Some of the tensor info is needed at various points during lowering.
Instead of storing the tensor index and looking it up every time, store
a point to the tensor struct instead.

Signed-off-by: Rob Herring (Arm) <robh@kernel.org>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/39975>
This commit is contained in:
Rob Herring (Arm) 2026-04-22 15:26:57 -05:00 committed by Marge Bot
parent 915cd57c08
commit ed2c19a411
3 changed files with 42 additions and 52 deletions

View file

@ -139,7 +139,7 @@ emit_strides(
{
unsigned elem_size = 1;
unsigned tensor_x, tensor_y, tensor_c;
struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, feature_map->tensor_idx);
struct ethosu_tensor *tensor = feature_map->tensor;
if (tensor->layout == ETHOSU_LAYOUT_NHCWB16) {
tensor_x = 16 * elem_size;
@ -181,10 +181,9 @@ emit_ifm_precision(struct ethosu_subgraph *subgraph,
struct ethosu_feature_map *feature_map,
enum ethosu_op_to_scale op_to_scale, uint32_t precision_cmd)
{
struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, feature_map->tensor_idx);
unsigned prec = 0;
if (tensor->layout == ETHOSU_LAYOUT_NHCWB16)
if (feature_map->tensor->layout == ETHOSU_LAYOUT_NHCWB16)
prec |= NPU_SET_IFM_PRECISION_FORMAT(1);
prec |= NPU_SET_IFM_PRECISION_PRECISION(feature_map->precision);
@ -238,10 +237,9 @@ emit_ofm(struct ethosu_subgraph *subgraph, struct ethosu_feature_map *feature_ma
static void
emit_ofm_precision(struct ethosu_subgraph *subgraph, struct ethosu_operation *operation)
{
struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, operation->ofm.tensor_idx);
unsigned prec = 0;
if (tensor->layout == ETHOSU_LAYOUT_NHCWB16)
if (operation->ofm.tensor->layout == ETHOSU_LAYOUT_NHCWB16)
prec |= NPU_SET_OFM_PRECISION_FORMAT(1);
prec |= NPU_SET_OFM_PRECISION_PRECISION(operation->ofm.precision);
@ -534,13 +532,12 @@ emit_ifm2_precision(struct ethosu_subgraph *subgraph,
struct ethosu_operation *operation,
bool has_scalar)
{
struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, operation->ifm2.tensor_idx);
unsigned prec = 0;
prec |= NPU_SET_IFM2_PRECISION_ACTIVATION_TYPE(operation->ifm2.is_signed);
prec |= NPU_SET_IFM2_PRECISION_ACTIVATION_PRECISION(operation->ifm2.precision);
if (tensor->layout == ETHOSU_LAYOUT_NHCWB16)
if (operation->ifm2.tensor->layout == ETHOSU_LAYOUT_NHCWB16)
prec |= NPU_SET_IFM2_PRECISION_ACTIVATION_FORMAT(1);
/* Vela: scalar → NONE(3), non-scalar → TILE2X2(0) */
@ -1087,10 +1084,9 @@ calc_blockdep(struct ethosu_subgraph *subgraph, struct ethosu_operation *prev_op
/* Check if previous OFM matches current IFM (same tensor) */
int ifm_index = 0;
if (operation->ifm2.tensor_idx != 0 &&
operation->ifm2.tensor_idx == prev_op->ofm.tensor_idx) {
if (operation->ifm2.tensor == prev_op->ofm.tensor) {
ifm_index = 1;
} else if (operation->ifm.tensor_idx != prev_op->ofm.tensor_idx) {
} else if (operation->ifm.tensor != prev_op->ofm.tensor) {
/* Previous operation doesn't produce current operation's IFM */
return device->max_concurrent_blocks;
}

View file

@ -31,27 +31,28 @@ needed_total_padding(int input_size, int stride, int filter_size)
}
static void
set_feature_maps(struct pipe_tensor *input_tensor,
set_feature_map(struct ethosu_subgraph *subgraph,
struct pipe_tensor *tensor,
struct ethosu_feature_map *fm)
{
fm->tensor = ethosu_find_tensor(subgraph, tensor->index);
fm->shape.height = tensor->dims[1];
fm->shape.width = tensor->dims[2];
fm->shape.depth = tensor->dims[3];
fm->zero_point = tensor->zero_point;
fm->scale = tensor->scale;
fm->is_signed = tensor->is_signed;
fm->precision = log2(tensor->type_size);
}
static void
set_feature_maps(struct ethosu_subgraph *subgraph,
struct pipe_tensor *input_tensor,
struct pipe_tensor *output_tensor,
struct ethosu_operation *operation)
{
operation->ifm.tensor_idx = input_tensor->index;
operation->ifm.shape.height = input_tensor->dims[1];
operation->ifm.shape.width = input_tensor->dims[2];
operation->ifm.shape.depth = input_tensor->dims[3];
operation->ifm.zero_point = input_tensor->zero_point;
operation->ifm.scale = input_tensor->scale;
operation->ifm.is_signed = input_tensor->is_signed;
operation->ifm.precision = log2(input_tensor->type_size);
operation->ofm.tensor_idx = output_tensor->index;
operation->ofm.shape.height = output_tensor->dims[1];
operation->ofm.shape.width = output_tensor->dims[2];
operation->ofm.shape.depth = output_tensor->dims[3];
operation->ofm.zero_point = output_tensor->zero_point;
operation->ofm.scale = output_tensor->scale;
operation->ofm.is_signed = output_tensor->is_signed;
operation->ofm.precision = log2(output_tensor->type_size);
set_feature_map(subgraph, input_tensor, &operation->ifm);
set_feature_map(subgraph, output_tensor, &operation->ofm);
}
static const struct pipe_ml_operation *
@ -70,9 +71,8 @@ ethosu_find_first_consumer(const struct pipe_ml_operation *poperations,
}
static unsigned
ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, unsigned tensor_idx)
ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, struct ethosu_tensor *tensor)
{
struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, tensor_idx);
unsigned size;
if (tensor->layout == ETHOSU_LAYOUT_NHWC) {
@ -100,12 +100,12 @@ ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, unsigned tensor_id
static void
allocate_feature_maps(struct ethosu_subgraph *subgraph, struct ethosu_operation *operation)
{
operation->ofm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ofm.tensor_idx);
operation->ofm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ofm.tensor);
operation->ofm.tiles.height_0 = operation->ofm.shape.height;
operation->ofm.tiles.height_1 = operation->ofm.shape.height;
operation->ofm.tiles.width_0 = operation->ofm.shape.width;
operation->ifm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ifm.tensor_idx);
operation->ifm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ifm.tensor);
operation->ifm.tiles.height_0 = operation->ifm.shape.height;
operation->ifm.tiles.height_1 = operation->ifm.shape.height;
operation->ifm.tiles.width_0 = operation->ifm.shape.width;
@ -152,7 +152,7 @@ ethosu_lower_convolution(struct ethosu_subgraph *subgraph,
operation->conv.depthwise = is_depthwise(poperation);
set_feature_maps(input_tensor, poperation->output_tensors[0], operation);
set_feature_maps(subgraph, input_tensor, poperation->output_tensors[0], operation);
operation->kernel.height = poperation->conv.weight_tensor->dims[1];
operation->kernel.width = poperation->conv.weight_tensor->dims[2];
@ -217,7 +217,7 @@ ethosu_lower_pooling(struct ethosu_subgraph *subgraph,
assert(0 && "Unsupported pooling type");
}
set_feature_maps(poperation->input_tensors[0], poperation->output_tensors[0], operation);
set_feature_maps(subgraph, poperation->input_tensors[0], poperation->output_tensors[0], operation);
operation->kernel.height = poperation->pooling.filter_height;
operation->kernel.width = poperation->pooling.filter_width;
@ -247,16 +247,14 @@ ethosu_lower_concatenation(struct ethosu_subgraph *subgraph,
} else
operation->pooling.type = ETHOSU_POOLING_TYPE_SUM;
set_feature_maps(poperation->input_tensors[input_idx], poperation->output_tensors[0], operation);
set_feature_maps(subgraph, poperation->input_tensors[input_idx], poperation->output_tensors[0], operation);
operation->ofm.shape.depth = operation->ifm.shape.depth;
allocate_feature_maps(subgraph, operation);
for (unsigned i = 0; i < input_idx; i++) {
struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, operation->ofm.tensor_idx);
if (tensor->layout == ETHOSU_LAYOUT_NHWC)
if (operation->ofm.tensor->layout == ETHOSU_LAYOUT_NHWC)
operation->ofm.tiles.addresses[0] += poperation->input_tensors[i]->dims[3];
else if (tensor->layout == ETHOSU_LAYOUT_NHCWB16)
else if (operation->ofm.tensor->layout == ETHOSU_LAYOUT_NHCWB16)
operation->ofm.tiles.addresses[0] += poperation->input_tensors[i]->dims[2] * align(poperation->input_tensors[i]->dims[3], 16);
else
assert(0 && "Unsupported layout");
@ -273,7 +271,7 @@ ethosu_lower_resize(struct ethosu_subgraph *subgraph,
operation->type = ETHOSU_OPERATION_TYPE_POOLING;
operation->pooling.type = ETHOSU_POOLING_TYPE_AVG;
set_feature_maps(poperation->input_tensors[0], poperation->output_tensors[0], operation);
set_feature_maps(subgraph, poperation->input_tensors[0], poperation->output_tensors[0], operation);
operation->upscale = ETHOSU_UPSCALE_NEAREST;
@ -289,7 +287,7 @@ ethosu_lower_strided_slice(struct ethosu_subgraph *subgraph,
operation->type = ETHOSU_OPERATION_TYPE_POOLING;
operation->pooling.type = ETHOSU_POOLING_TYPE_AVG;
set_feature_maps(poperation->input_tensors[0], poperation->output_tensors[0], operation);
set_feature_maps(subgraph, poperation->input_tensors[0], poperation->output_tensors[0], operation);
operation->ifm.shape = operation->ofm.shape;
allocate_feature_maps(subgraph, operation);
@ -340,16 +338,10 @@ ethosu_lower_add(struct ethosu_subgraph *subgraph,
operation->eltwise.ifm_reversed = true;
}
set_feature_maps(poperation->input_tensors[ifm_idx], poperation->output_tensors[0], operation);
set_feature_maps(subgraph, poperation->input_tensors[ifm_idx], poperation->output_tensors[0], operation);
set_feature_map(subgraph, poperation->input_tensors[ifm2_idx], &operation->ifm2);
operation->ifm2.tensor_idx = poperation->input_tensors[ifm2_idx]->index;
operation->ifm2.shape.height = poperation->input_tensors[ifm2_idx]->dims[1];
operation->ifm2.shape.width = poperation->input_tensors[ifm2_idx]->dims[2];
operation->ifm2.shape.depth = poperation->input_tensors[ifm2_idx]->dims[3];
operation->ifm2.zero_point = poperation->input_tensors[ifm2_idx]->zero_point;
operation->ifm2.scale = poperation->input_tensors[ifm2_idx]->scale;
operation->ifm2.is_signed = poperation->input_tensors[ifm2_idx]->is_signed;
operation->ifm2.precision = log2(poperation->input_tensors[ifm2_idx]->type_size);
if (poperation->input_tensors[ifm2_idx]->data &&
operation->ifm2.shape.width == 1 &&
operation->ifm2.shape.height == 1 &&
@ -361,7 +353,7 @@ ethosu_lower_add(struct ethosu_subgraph *subgraph,
allocate_feature_maps(subgraph, operation);
operation->ifm2.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ifm2.tensor_idx);
operation->ifm2.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ifm2.tensor);
operation->ifm2.tiles.height_0 = operation->ifm2.shape.height;
operation->ifm2.tiles.height_1 = operation->ifm2.shape.height;
operation->ifm2.tiles.width_0 = operation->ifm2.shape.width;

View file

@ -66,8 +66,10 @@ enum ethosu_upscale_mode {
ETHOSU_UPSCALE_ZEROS = 2,
};
struct ethosu_tensor;
struct ethosu_feature_map {
unsigned tensor_idx;
struct ethosu_tensor *tensor;
struct ethosu_block shape;
bool is_signed;
uint8_t precision;