From dae0af20abbc445a79108cbdb00ff7e738610861 Mon Sep 17 00:00:00 2001
From: Tomeu Vizoso <tomeu.vizoso@collabora.com>
Date: Wed, 11 Dec 2024 17:28:05 +0100
Subject: [PATCH] teflon: Add support for Reshape operations

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/34629>
---
 src/gallium/frontends/teflon/tfl_device.c | 10 ++++++++++
 src/gallium/include/pipe/p_state.h        |  8 ++++++++
 2 files changed, 18 insertions(+)

diff --git a/src/gallium/frontends/teflon/tfl_device.c b/src/gallium/frontends/teflon/tfl_device.c
index 0c7ecb63567..cde30ffab69 100644
--- a/src/gallium/frontends/teflon/tfl_device.c
+++ b/src/gallium/frontends/teflon/tfl_device.c
@@ -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++) {
diff --git a/src/gallium/include/pipe/p_state.h b/src/gallium/include/pipe/p_state.h
index 5cc3490c9b5..ef7e8ad2cdb 100644
--- a/src/gallium/include/pipe/p_state.h
+++ b/src/gallium/include/pipe/p_state.h
@@ -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;
    };
 };