From 3d8f10851455da130a15c02e97699277149b9515 Mon Sep 17 00:00:00 2001
From: Tomeu Vizoso <tomeu.vizoso@ideasonboard.com>
Date: Fri, 8 Nov 2024 18:59:07 +0100
Subject: [PATCH] teflon: Add support for FullyConnected

Reviewed-by: Philipp Zabel <p.zabel@pengutronix.de>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32510>
---
 src/gallium/frontends/teflon/tfl_device.c | 12 ++++++++++++
 src/gallium/include/pipe/p_state.h        | 18 ++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/gallium/frontends/teflon/tfl_device.c b/src/gallium/frontends/teflon/tfl_device.c
index 22411795306..25592f3040f 100644
--- a/src/gallium/frontends/teflon/tfl_device.c
+++ b/src/gallium/frontends/teflon/tfl_device.c
@@ -171,6 +171,12 @@ fill_operation(struct teflon_delegate *delegate, TfLiteContext *tf_context, TfLi
          operation->pad.after_y = paddings[5];
          break;
       }
+      case kTfLiteBuiltinFullyConnected: {
+         operation->type = PIPE_ML_OPERATION_TYPE_FULLY_CONNECTED;
+         operation->fcon.weight_tensor = &tensors[node->inputs->data[1]];
+         operation->fcon.bias_tensor = &tensors[node->inputs->data[2]];
+         break;
+      }
       default:
          unreachable("Unsupported ML operation type");
    }
@@ -252,6 +258,9 @@ dump_graph(struct pipe_tensor *tensors, unsigned tensor_count, struct pipe_ml_op
          case PIPE_ML_OPERATION_TYPE_PAD:
             teflon_debug("%-6s ", "PAD");
             break;
+         case PIPE_ML_OPERATION_TYPE_FULLY_CONNECTED:
+            teflon_debug("%-6s ", "FCON");
+            break;
       }
 
       for (unsigned j = 0; j < operations[i].input_count; j++) {
@@ -585,6 +594,9 @@ PrepareDelegate(TfLiteContext *context, TfLiteDelegate *delegate)
                         padding[7] == 0;
             break;
          }
+         case kTfLiteBuiltinFullyConnected:
+            supported = true;
+            break;
       }
 
       if (supported)
diff --git a/src/gallium/include/pipe/p_state.h b/src/gallium/include/pipe/p_state.h
index 8d59b5decfd..717d9fd7fa0 100644
--- a/src/gallium/include/pipe/p_state.h
+++ b/src/gallium/include/pipe/p_state.h
@@ -1062,6 +1062,7 @@ enum pipe_ml_operation_type {
    PIPE_ML_OPERATION_TYPE_CONCATENATION,
    PIPE_ML_OPERATION_TYPE_SPLIT,
    PIPE_ML_OPERATION_TYPE_PAD,
+   PIPE_ML_OPERATION_TYPE_FULLY_CONNECTED,
 };
 
 /**
@@ -1092,6 +1093,7 @@ struct pipe_ml_operation
           * For convolutions, tensor containing the weights.
           */
          struct pipe_tensor *weight_tensor;
+
          /**
           * For convolutions, tensor containing the biases.
           */
@@ -1173,6 +1175,22 @@ struct pipe_ml_operation
           */
          unsigned after_y;
       } pad;
+
+      struct {
+         /**
+          * Tensor containing the weights.
+          */
+         struct pipe_tensor *weight_tensor;
+         /**
+          * Tensor containing the biases.
+          */
+         struct pipe_tensor *bias_tensor;
+
+         /**
+          * Whether a ReLU activation should be applied to the output.
+          */
+         bool relu;
+      } fcon;
    };
 };