From 01dc82bc633583540da6d7d730a462d29b0ddb6e Mon Sep 17 00:00:00 2001
From: Philipp Zabel <p.zabel@pengutronix.de>
Date: Wed, 18 Sep 2024 12:59:01 +0200
Subject: [PATCH] teflon: Support fused ReLU activation

Enable support for fused ReLU activation on Conv2D and
DepthwiseConv2D operations.

Reviewed-by: Tomeu Vizoso <tomeu@tomeuvizoso.net>
Signed-off-by: Philipp Zabel <p.zabel@pengutronix.de>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31981>
---
 src/gallium/frontends/teflon/tfl_device.c | 18 ++++++++++++------
 src/gallium/include/pipe/p_state.h        |  5 +++++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/gallium/frontends/teflon/tfl_device.c b/src/gallium/frontends/teflon/tfl_device.c
index b5bffc2fe38..80bd7ee5fcb 100644
--- a/src/gallium/frontends/teflon/tfl_device.c
+++ b/src/gallium/frontends/teflon/tfl_device.c
@@ -115,7 +115,8 @@ fill_operation(struct teflon_delegate *delegate, TfLiteContext *tf_context, TfLi
          if (node_registration->builtin_code == kTfLiteBuiltinConv2d) {
             TfLiteConvParams* params = (TfLiteConvParams*)node->builtin_data;
 
-            assert(params->activation == kTfLiteActNone);
+            assert(params->activation == kTfLiteActNone ||
+                   params->activation == kTfLiteActRelu);
             if (node_registration->version >= 2) {
                assert(params->dilation_width_factor == 1);
                assert(params->dilation_height_factor == 1);
@@ -124,10 +125,12 @@ fill_operation(struct teflon_delegate *delegate, TfLiteContext *tf_context, TfLi
             operation->conv.stride_y = params->stride_height;
             operation->conv.padding_same = params->padding == kTfLitePaddingSame;
             operation->conv.depthwise = false;
+            operation->conv.relu = params->activation == kTfLiteActRelu;
          } else {
             TfLiteDepthwiseConvParams* params = (TfLiteDepthwiseConvParams*)node->builtin_data;
 
-            assert(params->activation == kTfLiteActNone);
+            assert(params->activation == kTfLiteActNone ||
+                   params->activation == kTfLiteActRelu);
             if (node_registration->version >= 2) {
                assert(params->dilation_width_factor == 1);
                assert(params->dilation_height_factor == 1);
@@ -136,6 +139,7 @@ fill_operation(struct teflon_delegate *delegate, TfLiteContext *tf_context, TfLi
             operation->conv.stride_y = params->stride_height;
             operation->conv.padding_same = params->padding == kTfLitePaddingSame;
             operation->conv.depthwise = true;
+            operation->conv.relu = params->activation == kTfLiteActRelu;
          }
          operation->conv.pointwise = operation->conv.weight_tensor->dims[1] == 1 && \
                                      operation->conv.weight_tensor->dims[2] == 1;
@@ -390,8 +394,9 @@ PrepareDelegate(TfLiteContext *context, TfLiteDelegate *delegate)
          case kTfLiteBuiltinConv2d: {
             TfLiteConvParams* params = (TfLiteConvParams*)node->builtin_data;
 
-            // Fused activation and dilation not yet implemented
-            if (params->activation == kTfLiteActNone &&
+            // Dilation not yet implemented
+            if ((params->activation == kTfLiteActNone ||
+                 params->activation == kTfLiteActRelu) &&
                 (registration->version < 2 ||
                  (params->dilation_width_factor == 1 &&
                   params->dilation_height_factor == 1))) {
@@ -402,8 +407,9 @@ PrepareDelegate(TfLiteContext *context, TfLiteDelegate *delegate)
          case kTfLiteBuiltinDepthwiseConv2d: {
             TfLiteDepthwiseConvParams* params = (TfLiteDepthwiseConvParams*)node->builtin_data;
 
-            // Fused activation and dilation not yet implemented
-            if (params->activation == kTfLiteActNone &&
+            // Dilation not yet implemented
+            if ((params->activation == kTfLiteActNone ||
+                 params->activation == kTfLiteActRelu) &&
                 (registration->version < 2 ||
                  (params->dilation_width_factor == 1 &&
                   params->dilation_height_factor == 1))) {
diff --git a/src/gallium/include/pipe/p_state.h b/src/gallium/include/pipe/p_state.h
index 01182a59246..8a808a530e5 100644
--- a/src/gallium/include/pipe/p_state.h
+++ b/src/gallium/include/pipe/p_state.h
@@ -1110,6 +1110,11 @@ struct pipe_ml_operation
           * Whether this is a depthwise convolution.
           */
          bool depthwise;
+
+         /**
+          * Whether this convolution has fused ReLU activation.
+          */
+         bool relu;
       } conv;
       struct {
          /**