[WASM] Add output_min, output_max to cache key for conv2d. (#2507)

BUG

Previously we weren't using this in the cache key, which could give incorrect results at runtime.

Author: Nikhil Thorat

tfjs-backend-wasm/src/cc/conv2d_impl.cc

@@ -19,7 +19,7 @@
 #include "src/cc/conv2d_impl.h"
 
 #include <xnnpack.h>
-#include <array>
+#include <tuple>
 #include <cmath>
 #include <limits>
 #include <map>
@@ -34,10 +34,11 @@
 #include "src/cc/util.h"
 
 namespace {
-// These integer values are keys to creating the conv2d operator. We use
-// std::array instead of a vanilla array as it implements the compare operator
+// We use std::tuple as the cache key as it implements the compare operator
 // needed for std::map.
-typedef std::array<int, 19> OperatorCacheKey;
+typedef std::tuple<int, int, int, int, int, int, int, int, int, int, int, int,
+                   int, int, int, int, int, int, int, float, float>
+    OperatorCacheKey;
 
 struct CachedInfo {
   xnn_operator_t op;
@@ -162,6 +163,16 @@ void conv2d(const int x_id, const int batch_size, const int input_height,
     clamp_method = tfjs::wasm::FusableActivation::LINEAR;
   }
 
+  float output_min = -std::numeric_limits<float>::infinity();
+  float output_max = std::numeric_limits<float>::infinity();
+
+  if (activation == FusableActivation::RELU) {
+    output_min = 0;
+  } else if (activation == FusableActivation::RELU6) {
+    output_min = 0;
+    output_max = 6;
+  }
+
   OperatorCacheKey cache_key = {pad_top,
                                 pad_right,
                                 pad_bottom,
@@ -180,20 +191,12 @@ void conv2d(const int x_id, const int batch_size, const int input_height,
                                 clamp_method,
                                 filter_id,
                                 bias_id,
-                                flags};
+                                flags,
+                                output_min,
+                                output_max};
 
   auto operator_cache_idx = operator_cache.find(cache_key);
   if (operator_cache_idx == operator_cache.end()) {
-    float output_min = -std::numeric_limits<float>::infinity();
-    float output_max = std::numeric_limits<float>::infinity();
-
-    if (activation == FusableActivation::RELU) {
-      output_min = 0;
-    } else if (activation == FusableActivation::RELU6) {
-      output_min = 0;
-      output_max = 6;
-    }
-
     // This lives outside the if statement so the data survives the scope.
     std::vector<float> transposed_filter;
 

tfjs-backend-wasm/src/cc/kernels/AvgPool.cc

@@ -65,8 +65,8 @@ void AvgPool(const int x_id, const int batch_size, const int input_height,
   auto operator_cache_idx = operator_cache.find(cache_key);
 
   if (operator_cache_idx == operator_cache.end()) {
-    float output_min = -std::numeric_limits<float>::infinity();
-    float output_max = std::numeric_limits<float>::infinity();
+    const float output_min = -std::numeric_limits<float>::infinity();
+    const float output_max = std::numeric_limits<float>::infinity();
 
     xnn_status status = xnn_create_average_pooling2d_nhwc_f32(
         pad_top, pad_right, pad_bottom, pad_left, filter_height, filter_width,

tfjs-backend-wasm/src/cc/kernels/FusedConv2D_test.cc

@@ -168,9 +168,20 @@ TEST(FUSEDCONV2D, xnn_operator_lifetime) {
       prelu_weights_id, out_id);
   ASSERT_EQ(6, tfjs::backend::xnn_operator_count);
 
+  // One new XNN operator should be created for the next call to conv2d with a
+  // different activation.
+  const int activation2 = tfjs::wasm::FusableActivation::RELU6;
+  tfjs::wasm::FusedConv2D(
+      x1_id, batch_size, input_height, input_width, weights1_id, filter_height,
+      filter_width, bias1_id, pad_top1, pad_right, pad_bottom1, pad_left,
+      is_same_pad1, dilation_height, dilation_width, stride_height,
+      stride_width, input_channels, output_channels, activation2,
+      prelu_weights_id, out_id);
+  ASSERT_EQ(7, tfjs::backend::xnn_operator_count);
+
   // Disposing the first weights should remove 2 operators.
   tfjs::wasm::dispose_data(weights0_id);
-  ASSERT_EQ(4, tfjs::backend::xnn_operator_count);
+  ASSERT_EQ(5, tfjs::backend::xnn_operator_count);
 
   // Disposing the second bias should remove 2 operators it's associated with.
   tfjs::wasm::dispose_data(bias1_id);

tfjs-backend-wasm/src/cc/kernels/FusedDepthwiseConv2D_test.cc

@@ -184,9 +184,20 @@ TEST(FUSEDDEPTHWISECONV2D, xnn_operator_lifetime) {
       prelu_weights_id, out_id);
   ASSERT_EQ(7, tfjs::backend::xnn_operator_count);
 
+  // One new XNN operator should be created for the next call to conv2d with a
+  // different activation.
+  const int activation2 = tfjs::wasm::FusableActivation::RELU6;
+  tfjs::wasm::FusedDepthwiseConv2D(
+      x1_id, batch_size, input_height, input_width, weights1_id, filter_height,
+      filter_width, bias1_id, pad_top1, pad_right, pad_bottom1, pad_left,
+      is_same_pad1, dilation_height, dilation_width, stride_height,
+      stride_width, input_channels, output_channels, activation2,
+      prelu_weights_id, out_id);
+  ASSERT_EQ(8, tfjs::backend::xnn_operator_count);
+
   // Disposing the first weights should remove 2 operators.
   tfjs::wasm::dispose_data(weights0_id);
-  ASSERT_EQ(5, tfjs::backend::xnn_operator_count);
+  ASSERT_EQ(6, tfjs::backend::xnn_operator_count);
 
   // Disposing the second bias should remove 2 operators it's associated with.
   tfjs::wasm::dispose_data(bias1_id);

tfjs-backend-wasm/src/cc/kernels/MaxPool.cc

@@ -67,8 +67,8 @@ void MaxPool(const int x_id, const int batch_size, const int input_height,
   auto operator_cache_idx = operator_cache.find(cache_key);
 
   if (operator_cache_idx == operator_cache.end()) {
-    float output_min = -std::numeric_limits<float>::infinity();
-    float output_max = std::numeric_limits<float>::infinity();
+    const float output_min = -std::numeric_limits<float>::infinity();
+    const float output_max = std::numeric_limits<float>::infinity();
 
     xnn_status status = xnn_create_max_pooling2d_nhwc_f32(
         pad_top, pad_right, pad_bottom, pad_left, filter_height, filter_width,