diff --git a/src/gallium/drivers/ethosu/ethosu_cmd.c b/src/gallium/drivers/ethosu/ethosu_cmd.c
index 36e1a781aaf..b859307a3fd 100644
--- a/src/gallium/drivers/ethosu/ethosu_cmd.c
+++ b/src/gallium/drivers/ethosu/ethosu_cmd.c
@@ -397,12 +397,12 @@ emit_common(struct ethosu_subgraph *subgraph, struct ethosu_operation *operation
 static void
 emit_convolution(struct ethosu_subgraph *subgraph, struct ethosu_operation *operation)
 {
-   ethosu_allocate_feature_map(subgraph, &operation->ifm);
+   operation->ifm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ifm.tensor_idx);
    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;
 
-   ethosu_allocate_feature_map(subgraph, &operation->ofm);
+   operation->ofm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ofm.tensor_idx);
    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;
diff --git a/src/gallium/drivers/ethosu/ethosu_lower.c b/src/gallium/drivers/ethosu/ethosu_lower.c
index 1006ce09dee..b5116fcad51 100644
--- a/src/gallium/drivers/ethosu/ethosu_lower.c
+++ b/src/gallium/drivers/ethosu/ethosu_lower.c
@@ -8,6 +8,7 @@
 #include "ethosu_device.h"
 #include "ethosu_lower.h"
 #include "ethosu_coefs.h"
+#include "ethosu_ml.h"
 #include "ethosu_sched.h"
 
 static bool
@@ -71,15 +72,15 @@ ethosu_find_first_consumer(const struct pipe_ml_operation *poperations,
 static void
 allocate_feature_maps(struct ethosu_subgraph *subgraph, struct ethosu_operation *operation)
 {
-   ethosu_allocate_feature_map(subgraph, &operation->ifm);
-   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;
-
-   ethosu_allocate_feature_map(subgraph, &operation->ofm);
+   operation->ofm.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ofm.tensor_idx);
    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.height_0 = operation->ifm.shape.height;
+   operation->ifm.tiles.height_1 = operation->ifm.shape.height;
+   operation->ifm.tiles.width_0 = operation->ifm.shape.width;
 }
 
 static const struct pipe_ml_operation *
@@ -368,7 +369,7 @@ ethosu_lower_add(struct ethosu_subgraph *subgraph,
 
    allocate_feature_maps(subgraph, operation);
 
-   ethosu_allocate_feature_map(subgraph, &operation->ifm2);
+   operation->ifm2.tiles.addresses[0] = ethosu_allocate_feature_map(subgraph, operation->ifm2.tensor_idx);
    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;
diff --git a/src/gallium/drivers/ethosu/ethosu_ml.c b/src/gallium/drivers/ethosu/ethosu_ml.c
index e3b24238e8e..7643623d21d 100644
--- a/src/gallium/drivers/ethosu/ethosu_ml.c
+++ b/src/gallium/drivers/ethosu/ethosu_ml.c
@@ -59,10 +59,10 @@ ethosu_register_tensor(struct ethosu_subgraph *subgraph,
    util_dynarray_append(&subgraph->tensors, new_tensor);
 }
 
-void
-ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, struct ethosu_feature_map *feature_map)
+unsigned
+ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, unsigned tensor_idx)
 {
-   struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, feature_map->tensor_idx);
+   struct ethosu_tensor *tensor = ethosu_find_tensor(subgraph, tensor_idx);
    unsigned size;
 
    if (tensor->layout == ETHOSU_LAYOUT_NHWC) {
@@ -77,16 +77,14 @@ ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, struct ethosu_feat
 
    assert(tensor);
 
-   if (tensor->size > 0) {
-      feature_map->tiles.addresses[0] = tensor->offset;
-      return;
-   }
+   if (tensor->size > 0)
+      return tensor->offset;
 
    tensor->offset = subgraph->io_used;
    tensor->size = size;
    subgraph->io_used += ALIGN_POT(size, 16);
 
-   feature_map->tiles.addresses[0] = tensor->offset;
+   return tensor->offset;
 }
 
 struct ethosu_tensor *
diff --git a/src/gallium/drivers/ethosu/ethosu_ml.h b/src/gallium/drivers/ethosu/ethosu_ml.h
index e582ac08efa..e1609ccb6a2 100644
--- a/src/gallium/drivers/ethosu/ethosu_ml.h
+++ b/src/gallium/drivers/ethosu/ethosu_ml.h
@@ -242,7 +242,7 @@ void ethosu_ml_subgraph_read_outputs(struct pipe_context *pcontext,
 void ethosu_ml_subgraph_destroy(struct pipe_context *context,
                                 struct pipe_ml_subgraph *psubgraph);
 
-void ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, struct ethosu_feature_map *feature_map);
+unsigned ethosu_allocate_feature_map(struct ethosu_subgraph *subgraph, unsigned tensor_idx);
 
 void ethosu_register_tensor(struct ethosu_subgraph *subgraph, const struct pipe_tensor *ptensor);