Add unit tests for RISP4ML

This CL is Part 3 of the comprehensive 3-part changes of integrating RISP4ML into Shodan ML toolchain. This change focuses on adding unit tests for RISP4ML based on pw_unit_test.
I have verifed that all unit tests passed.

Details of the 3-part changes:

(Part 1) Completely re-wrote RISP4ML in plain C (merged)
- The original RISP4ML on google3 was written in C++. The memory usage with C++ RISP4ML was significantly larger. Thus, I have completely re-written RISP4ML in plain C. The memory increase with the plain-C RISP4ML is much smaller.
- Made the input/output shape to be configurable arguments instead of constants.
- Made low-level optimizations and simplifications.

(Part 2) Integrated plain-C RISP4ML into Shodan ML toolchain (merged)
- Use fssd_25_8bit_v2 as an example to load the binary (raw bayer) file and go through RISP4ML toolchain, and feed the RISP4ML output into IREE flow.
- The fssd_25_8bit_v2 example is up and running, and the correctness of output has been verified (compare output@plain-C risp4ml against output@google3).

(Part 3) Added unit tests for plain-C RISP4ML (this CL)
- Added unit tests based on pw_unit_test
- Build unit tests in iree_cc_binary
- Run and check unit tests via lit_test

Change-Id: I9c01792dd8e35176cbc16e18bf2f9c17d5f801b7
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 5f0802f..52f6a55 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -88,3 +88,10 @@
 include(iree_model_input)
 # Add the included directory here.
 add_subdirectory(samples)
+
+# Add pigweed support
+include($ENV{ROOTDIR}/sw/pigweed/pw_build/pigweed.cmake)
+pw_set_backend(pw_log pw_log_basic)
+pw_set_backend(pw_assert pw_assert_log)
+pw_set_backend(pw_sys_io pw_sys_io_stdio)
+add_subdirectory($ENV{ROOTDIR}/sw/pigweed pw)
diff --git a/samples/risp4ml/common/CMakeLists.txt b/samples/risp4ml/common/CMakeLists.txt
index 49346c1..793f933 100644
--- a/samples/risp4ml/common/CMakeLists.txt
+++ b/samples/risp4ml/common/CMakeLists.txt
@@ -24,3 +24,16 @@
   SRCS
     "utils.c"
 )
+
+iree_cc_binary(
+  NAME
+    image_test
+  SRCS
+    "image_test.cc"
+  DEPS
+    ::image
+    ::test_utils
+    pw_unit_test
+    pw_unit_test.main
+    pw_assert_basic
+)
diff --git a/samples/risp4ml/common/image_test.cc b/samples/risp4ml/common/image_test.cc
new file mode 100644
index 0000000..a2fb28d
--- /dev/null
+++ b/samples/risp4ml/common/image_test.cc
@@ -0,0 +1,31 @@
+#include "pw_unit_test/framework.h"
+#include "samples/risp4ml/common/image.h"
+#include "samples/risp4ml/common/test_utils.h"
+
+class ImageTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    img_ = image_new(3, 32, 64);
+    FillImage(img_);
+  }
+  void TearDown() override { image_delete(img_); }
+
+  Image* img_;
+};
+
+TEST_F(ImageTest, PixelAccess) {
+  for (uint16_t c = 0; c < img_->num_channels; ++c) {
+    for (uint16_t y = 0; y < img_->height; ++y) {
+      for (uint16_t x = 0; x < img_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(img_, c, y, x), Pattern(c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(ImageTest, RowAccess) {
+  pixel_type_t* row_4 = image_row(img_, 1, 4);
+  for (uint16_t i = 0; i < 16; ++i) {
+    ASSERT_EQ(row_4[i], Pattern(1, 4, i));
+  }
+}
diff --git a/samples/risp4ml/common/image_test.txt b/samples/risp4ml/common/image_test.txt
new file mode 100644
index 0000000..beeed75
--- /dev/null
+++ b/samples/risp4ml/common/image_test.txt
@@ -0,0 +1,3 @@
+// RUN: ${TEST_RUNNER_CMD} ${OUT}/springbok_iree/samples/risp4ml/common/image_test
+// RUN: cat %t | FileCheck %s
+// CHECK: {{PASSED}}
diff --git a/samples/risp4ml/isp_stages/CMakeLists.txt b/samples/risp4ml/isp_stages/CMakeLists.txt
index 54022b7..bc2c798 100644
--- a/samples/risp4ml/isp_stages/CMakeLists.txt
+++ b/samples/risp4ml/isp_stages/CMakeLists.txt
@@ -69,3 +69,65 @@
     samples::risp4ml::common::image
     samples::risp4ml::common::utils
 )
+
+iree_cc_binary(
+  NAME
+    blc_test
+  SRCS
+    "blc_test.cc"
+  DEPS
+    ::blc
+    samples::risp4ml::common::test_utils
+    pw_unit_test
+    pw_unit_test.main
+    pw_assert_basic
+)
+
+iree_cc_binary(
+  NAME
+    dg_test
+  SRCS
+    "dg_test.cc"
+  DEPS
+    ::dg
+    samples::risp4ml::common::test_utils
+    pw_unit_test
+    pw_unit_test.main
+    pw_assert_basic
+)
+
+iree_cc_binary(
+  NAME
+    downscale_test
+  SRCS
+    "downscale_test.cc"
+  DEPS
+    ::downscale
+    pw_unit_test
+    pw_unit_test.main
+    pw_assert_basic
+)
+
+iree_cc_binary(
+  NAME
+    gamma_test
+  SRCS
+    "gamma_test.cc"
+  DEPS
+    ::gamma
+    pw_unit_test
+    pw_unit_test.main
+    pw_assert_basic
+)
+
+iree_cc_binary(
+  NAME
+    wbg_test
+  SRCS
+    "wbg_test.cc"
+  DEPS
+    ::wbg
+    pw_unit_test
+    pw_unit_test.main
+    pw_assert_basic
+)
diff --git a/samples/risp4ml/isp_stages/blc_test.cc b/samples/risp4ml/isp_stages/blc_test.cc
new file mode 100644
index 0000000..1935048
--- /dev/null
+++ b/samples/risp4ml/isp_stages/blc_test.cc
@@ -0,0 +1,80 @@
+#include "pw_unit_test/framework.h"
+#include "samples/risp4ml/common/constants.h"
+#include "samples/risp4ml/common/test_utils.h"
+#include "samples/risp4ml/isp_stages/blc.h"
+
+static constexpr int kFrameSize = 16;
+static constexpr int kNumPatterns = 4;
+
+class BlcTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    in_ = image_new(1, kFrameSize, kFrameSize);
+    out_ = image_new(1, kFrameSize, kFrameSize);
+    InitImageRandom(in_, 0, kRawPipelineMaxVal);
+    // Force max/min values to be included.
+    *image_pixel(in_, 0, 0, 0) = 0;
+    *image_pixel(in_, 0, 0, 1) = kRawPipelineMaxVal;
+  }
+  void TearDown() override {
+    image_delete(in_);
+    image_delete(out_);
+  }
+
+  Image* in_;
+  Image* out_;
+};
+
+TEST_F(BlcTest, Bypass) {
+  BlcParams params = {.enable = false, .offsets = {20, 20, 20, 20}};
+  set_blc_params(&params);
+
+  blc_process(in_, out_);
+
+  // Expect no change
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(in_, c, y, x),
+                  image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(BlcTest, NoChangeRandomPixel) {
+  // Set all 4 Bayer channel offsets to 0.
+  BlcParams params = {.enable = true, .offsets = {0, 0, 0, 0}};
+  set_blc_params(&params);
+
+  blc_process(in_, out_);
+
+  // Expect no change
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(in_, c, y, x),
+                  image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(BlcTest, MaxOffsetClampToZero) {
+  // Set all 4 Bayer channel offsets to max allowed value.
+  BlcParams params = {.enable = true,
+                      .offsets = {kRawPipelineMaxVal, kRawPipelineMaxVal,
+                                  kRawPipelineMaxVal, kRawPipelineMaxVal}};
+  set_blc_params(&params);
+
+  blc_process(in_, out_);
+
+  // Expect correct subtraction
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(0, image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
diff --git a/samples/risp4ml/isp_stages/blc_test.txt b/samples/risp4ml/isp_stages/blc_test.txt
new file mode 100644
index 0000000..608d615
--- /dev/null
+++ b/samples/risp4ml/isp_stages/blc_test.txt
@@ -0,0 +1,3 @@
+// RUN: ${TEST_RUNNER_CMD} ${OUT}/springbok_iree/samples/risp4ml/isp_stages/blc_test
+// RUN: cat %t | FileCheck %s
+// CHECK: {{PASSED}}
diff --git a/samples/risp4ml/isp_stages/dg_test.cc b/samples/risp4ml/isp_stages/dg_test.cc
new file mode 100644
index 0000000..cf3b7dc
--- /dev/null
+++ b/samples/risp4ml/isp_stages/dg_test.cc
@@ -0,0 +1,250 @@
+#include "pw_unit_test/framework.h"
+#include "samples/risp4ml/common/constants.h"
+#include "samples/risp4ml/common/test_utils.h"
+#include "samples/risp4ml/common/utils.h"
+#include "samples/risp4ml/isp_stages/dg.h"
+
+static constexpr uint16_t kDgFractional = kRawPipelineFraction;
+static constexpr uint16_t kFrameSize = 16;
+
+class DgTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    in_ = image_new(1, kFrameSize, kFrameSize);
+    out_ = image_new(1, kFrameSize, kFrameSize);
+  }
+  void TearDown() override {
+    image_delete(in_);
+    image_delete(out_);
+  }
+
+  Image* in_;
+  Image* out_;
+};
+
+TEST_F(DgTest, Bypass) {
+  InitImage(in_, 1);
+
+  // Set gain to 2x, gain is in 8.16 format.
+  uint16_t gain = 2 << kDgFractional;
+  DgParams params = {.enable = false, .gains = {gain, gain, gain, gain}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect no change.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(in_, c, y, x),
+                  image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, NoChangeRandomPixel) {
+  InitImageRandom(in_, kRawPipelineMinVal, kRawPipelineMaxVal);
+
+  // Force max/min values to be included.
+  *image_pixel(in_, 0, 0, 0) = 0;
+  *image_pixel(in_, 0, 0, 1) = kRawPipelineMaxVal;
+
+  // Set gain to 1x, gain is in_ 8.8 format.
+  uint16_t gain = 1 << kDgFractional;
+  DgParams params = {.enable = true, .gains = {gain, gain, gain, gain}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect no change.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(in_, c, y, x),
+                  image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, TwoTimesGainRandomPixel) {
+  // Set pixel values inside 0.5x max range, to multiply by 2x without clipping.
+  InitImageRandom(in_, kRawPipelineMinVal / 2, kRawPipelineMaxVal / 2);
+
+  // Force max/min values to be included.
+  *image_pixel(in_, 0, 0, 0) = kRawPipelineMinVal / 2;
+  *image_pixel(in_, 0, 0, 1) = kRawPipelineMaxVal / 2;
+
+  // Set gain to 2x, gain is in_ 8.16 format.
+  uint16_t gain = 2 << kDgFractional;
+  DgParams params = {.enable = true, .gains = {gain, gain, gain, gain}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect all pixel values to be doubled, as no clipping.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(in_, c, y, x) * 2,
+                  image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, ClampHighRandomPixel) {
+  // Init image with range of values that will clamp to max with 2x gain.
+  InitImageRandom(in_, kRawPipelineMaxVal / 2, kRawPipelineMaxVal);
+
+  // Set gain to 2x, gain is in_ 8.16 format.
+  uint16_t gain = 2 << kDgFractional;
+  DgParams params = {.enable = true, .gains = {gain, gain, gain, gain}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect all pixel values to be clamped high.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(out_, c, y, x), kRawPipelineMaxVal);
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, ClampLowRandomPixel) {
+  // Init image with range of values that will clamp to min with 2x gain.
+  InitImageRandom(in_, kRawPipelineMinVal / 2, kRawPipelineMinVal);
+
+  // Set gain to 2x, gain is in_ 8.8 format.
+  uint16_t gain = 2 << kDgFractional;
+  DgParams params = {.enable = true, .gains = {gain, gain, gain, gain}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect all pixel values to be clamped low.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(out_, c, y, x), kRawPipelineMinVal);
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, MaxGainLinearRampInput) {
+  // Set pixel values from left to right with progressively larger values.
+  // 1, 2, 4, ... 2^15
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        *image_pixel(in_, c, y, x) =
+            (pixel_type_t)Clamp(1 << y, 0, kRawPipelineMaxVal);
+      }
+    }
+  }
+
+  // Set gain to max gain, 0xFFFF in_ 8.8 format, approximately 256.
+  constexpr uint32_t kMaxGain =
+      (0xFF << kDgFractional) + ((1 << kDgFractional) - 1);
+  DgParams params = {.enable = true,
+                     .gains = {kMaxGain, kMaxGain, kMaxGain, kMaxGain}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect values 1, 2, ... 64 to be gained by 256, 128 and above to clip.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ((pixel_type_t)Clamp((1 << 8) << y, 0, kRawPipelineMaxVal),
+                  image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, MinValidGainLinearRampInput) {
+  // Set pixel values from left to right with progressively larger values.
+  // 1, 2, 4, ... 2^15
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        *image_pixel(in_, c, y, x) =
+            (pixel_type_t)Clamp(1 << y, 0, kRawPipelineMaxVal);
+      }
+    }
+  }
+
+  // Set gain to min non-zero gain, 0x000001 in_ 8.8 format, 1/256.
+  DgParams params = {.enable = true, .gains = {1, 1, 1, 1}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  constexpr uint32_t kExpectedOutput[] = {0, 0, 0, 0, 0,  0,  0,  1,
+                                          1, 2, 4, 8, 16, 32, 64, 128};
+
+  // Expect all output values to be zero, as 32767/65536=0.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(kExpectedOutput[y], image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, MinUsefulGainLinearRampInput) {
+  // Set pixel values from left to right with progressively larger values.
+  // 1, 2, 4, ... kRawPipelineMaxVal
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        *image_pixel(in_, c, y, x) =
+            (pixel_type_t)Clamp(1 << y, 0, kRawPipelineMaxVal);
+      }
+    }
+  }
+
+  // Set gain to min useful gain, 0x000002 in_ 8.8 format, 1/256.
+  DgParams params = {.enable = true, .gains = {2, 2, 2, 2}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect all output values < 128 to be zero, 128->1, 256->2.
+  constexpr uint32_t kExpectedOutput[] = {0, 0, 0, 0,  0,  0,  1,   1,
+                                          2, 4, 8, 16, 32, 64, 128, 256};
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(kExpectedOutput[y], image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(DgTest, ZeroGainRandomInput) {
+  // Init image with range of values.
+  InitImageRandom(in_, kRawPipelineMinVal, kRawPipelineMinVal);
+
+  // Set gain to 0x, gain is in_ 8.16 format.
+  DgParams params = {.enable = true, .gains = {0, 0, 0, 0}};
+  set_dg_params(&params);
+
+  dg_process(in_, out_);
+
+  // Expect all output values to be zero.
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(0, image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
diff --git a/samples/risp4ml/isp_stages/dg_test.txt b/samples/risp4ml/isp_stages/dg_test.txt
new file mode 100644
index 0000000..7913734
--- /dev/null
+++ b/samples/risp4ml/isp_stages/dg_test.txt
@@ -0,0 +1,3 @@
+// RUN: ${TEST_RUNNER_CMD} ${OUT}/springbok_iree/samples/risp4ml/isp_stages/dg_test
+// RUN: cat %t | FileCheck %s
+// CHECK: {{PASSED}}
diff --git a/samples/risp4ml/isp_stages/downscale_test.cc b/samples/risp4ml/isp_stages/downscale_test.cc
new file mode 100644
index 0000000..f679b0b
--- /dev/null
+++ b/samples/risp4ml/isp_stages/downscale_test.cc
@@ -0,0 +1,251 @@
+#include <climits>
+#include <cmath>
+
+#include "pw_unit_test/framework.h"
+#include "samples/risp4ml/common/constants.h"
+#include "samples/risp4ml/common/test_utils.h"
+#include "samples/risp4ml/isp_stages/downscale.h"
+
+static constexpr uint16_t kScalePrecision = 8;
+static constexpr uint16_t kScaleFixedOne = (1 << kScalePrecision);
+static constexpr uint16_t kInterpolatePrecision = 8;
+static constexpr uint16_t kInterpolateShift = 1;
+static constexpr float kOutBitsShift = 1 << (kRawPipelineBpp - kPipeOutputBpp);
+
+class DownscaleTest : public ::testing::Test {
+ protected:
+  void setup(uint16_t in_ch, uint16_t in_height, uint16_t in_width,
+             uint16_t out_ch, uint16_t out_height, uint16_t out_width) {
+    in_ = image_new(in_ch, in_height, in_width);
+    out_ = imageu8_new(out_ch, out_height, out_width);
+    params_.enable = true;
+    params_.scale_precision = kScalePrecision;
+    params_.interpolate_precision = kScalePrecision;
+    params_.interpolate_shift = kInterpolateShift;
+    params_.scale_fixed_one = kScaleFixedOne;
+    params_.scale_fraction_mask = kScaleFixedOne - 1;
+    params_.weight_shift = 0;
+  }
+  void TearDown() override {
+    image_delete(in_);
+    imageu8_delete(out_);
+  }
+  ImageU8* imageu8_new(uint16_t num_channels, uint16_t height, uint16_t width);
+  void imageu8_delete(ImageU8* image) {
+    if (image) {
+      if (image->data) free(image->data);
+      free(image);
+    }
+  }
+  pixel_type_t imageu8_pixel_val(ImageU8* image, uint16_t c, uint16_t y,
+                                 uint16_t x) {
+    const uint32_t stride_c = image->width * image->height;
+    const uint16_t stride_y = image->width;
+    const uint16_t stride_x = 1;
+    return *(image->data + c * stride_c + y * stride_y + x * stride_x);
+  }
+  float ExpectedOut(Image* in, uint16_t y, uint16_t x, float hor_scale,
+                    float ver_scale, float hor_initial_offset,
+                    float ver_initial_offset);
+  void ScaleRampImageTest(float hor_scale, float ver_scale,
+                          float hor_initial_offset, float ver_initial_offset,
+                          uint16_t input_width = 640,
+                          uint16_t input_height = 480);
+  struct BilinearScaleTestValues {
+    float hor_scale;
+    float ver_scale;
+    float hor_initial_offset;
+    float ver_initial_offset;
+    uint16_t input_width;
+    uint16_t input_height;
+  };
+
+  Image* in_;
+  ImageU8* out_;
+  DownscaleParams params_;
+};
+
+ImageU8* DownscaleTest::imageu8_new(uint16_t num_channels, uint16_t height,
+                                    uint16_t width) {
+  ImageU8* image = (ImageU8*)malloc(sizeof(ImageU8));
+  if (image) {
+    image->num_channels = num_channels;
+    image->height = height;
+    image->width = width;
+    uint32_t num_pixels = width * height * num_channels;
+    image->data = (uint8_t*)malloc(num_pixels * sizeof(uint8_t));
+  }
+  return image;
+}
+
+float DownscaleTest::ExpectedOut(Image* in, uint16_t y, uint16_t x,
+                                 float hor_scale, float ver_scale,
+                                 float hor_initial_offset,
+                                 float ver_initial_offset) {
+  float input_x = x * hor_scale + hor_initial_offset;
+  float input_y = y * ver_scale + ver_initial_offset;
+  uint16_t int_input_x_pre = static_cast<uint16_t>(floorf(input_x));
+  uint16_t int_input_x_post = int_input_x_pre + 1;
+  uint16_t int_input_y_pre = static_cast<uint16_t>(floorf(input_y));
+  uint16_t int_input_y_post = int_input_y_pre + 1;
+
+  float p0 = (input_y - int_input_y_pre) *
+                 image_pixel_val(in, 0, int_input_y_post, int_input_x_pre) +
+             (int_input_y_post - input_y) *
+                 image_pixel_val(in, 0, int_input_y_pre, int_input_x_pre);
+  float p1 = (input_y - int_input_y_pre) *
+                 image_pixel_val(in, 0, int_input_y_post, int_input_x_post) +
+             (int_input_y_post - input_y) *
+                 image_pixel_val(in, 0, int_input_y_pre, int_input_x_post);
+  float expected_out =
+      (input_x - int_input_x_pre) * p1 + (int_input_x_post - input_x) * p0;
+  // Downscaler out_ is kPipeOutputBpp bits wide.
+  // Shift right by kRawPipelineBpp - kPipeOutputBpp and round.
+
+  expected_out = floorf(expected_out / kOutBitsShift);
+  return expected_out;
+}
+
+// Helper function for 2D ramp tests. image is downscaled successfully.
+void DownscaleTest::ScaleRampImageTest(float hor_scale, float ver_scale,
+                                       float hor_initial_offset,
+                                       float ver_initial_offset,
+                                       uint16_t input_width,
+                                       uint16_t input_height) {
+  constexpr int kTolerance = 1;  // Tolerance for rounding error.
+  uint16_t output_width =
+      static_cast<uint16_t>((input_width - hor_initial_offset) / hor_scale);
+  uint16_t output_height =
+      static_cast<uint16_t>((input_height - ver_initial_offset) / ver_scale);
+  setup(1, input_height, input_width, 1, output_height, output_width);
+
+  // Fill in_ images as 2D ramp whose values are increased from the
+  // top-left corner to the bottom-right corner.
+  for (uint16_t y = 0; y < input_height; ++y) {
+    for (uint16_t x = 0; x < input_width; ++x) {
+      *image_pixel(in_, 0, y, x) = (y * input_width + x) % (1024);
+    }
+  }
+
+  params_.ver_initial_offset = ver_initial_offset * kScaleFixedOne;
+  params_.hor_initial_offset = hor_initial_offset * kScaleFixedOne;
+  params_.hor_scale_factor = hor_scale * kScaleFixedOne;
+  params_.ver_scale_factor = ver_scale * kScaleFixedOne;
+  set_downscale_param(&params_);
+
+  downscale_process(in_, out_);
+
+  for (uint16_t y = 0; y < output_height; ++y) {
+    for (uint16_t x = 0; x < output_width; ++x) {
+      float expected_out = ExpectedOut(in_, y, x, hor_scale, ver_scale,
+                                       hor_initial_offset, ver_initial_offset);
+      float diff =
+          std::abs((float)imageu8_pixel_val(out_, 0, y, x) - expected_out);
+      ASSERT_LE(diff, kTolerance);
+    }
+  }
+}
+
+TEST_F(DownscaleTest, NoScaleTest) {
+  constexpr uint16_t kOutputWidth = 128;
+  constexpr uint16_t kInputHeight = 96;
+  setup(1, kInputHeight, kOutputWidth, 1, kInputHeight, kOutputWidth);
+
+  // Generate random image.
+  InitImageRandom(in_, 0, USHRT_MAX);
+
+  params_.ver_initial_offset = 0;
+  params_.hor_initial_offset = 0;
+  params_.hor_scale_factor = 1 << kScalePrecision;
+  params_.ver_scale_factor = 1 << kScalePrecision;
+  set_downscale_param(&params_);
+
+  downscale_process(in_, out_);
+
+  // Verify the out_ image is identical to the in_ image.
+  for (uint16_t y = 0; y < kInputHeight; ++y) {
+    for (uint16_t x = 0; x < kOutputWidth; ++x) {
+      ASSERT_EQ(imageu8_pixel_val(out_, 0, y, x),
+                static_cast<pixel_type_t>(
+                    floorf(image_pixel_val(in_, 0, y, x) >>
+                           (kRawPipelineBpp - kPipeOutputBpp))));
+    }
+  }
+}
+
+TEST_F(DownscaleTest, VGAZeroOffsetTest) {
+  std::vector<BilinearScaleTestValues> tests = {
+      {2.f, 2.f, 0, 0},
+      {4.f, 4.f, 320, 240},  // Offset from center
+      {2.f, 4.f, 200, 100},
+      {1.5f, 3.4f, 17, 55}};
+
+  for (const auto& test : tests) {
+    ScaleRampImageTest(test.hor_scale, test.ver_scale, test.hor_initial_offset,
+                       test.ver_initial_offset);
+  }
+}
+
+// TODO(alexkaplan): parametrize this test for different in_ size and scale
+// parameters.
+TEST_F(DownscaleTest, DownscaleTest) {
+  std::vector<BilinearScaleTestValues> tests = {
+      {7.3f, 5.1f, 5, 2, 64, 64},  // Test odd values
+      {4, 4, 320, 240, 640, 480},  // Offset from center
+      {2, 4, 200, 100, 640, 480},
+      {1.5f, 3.4f, 17, 55, 200, 100}};
+
+  for (const auto& test : tests) {
+    ScaleRampImageTest(test.hor_scale, test.ver_scale, test.hor_initial_offset,
+                       test.ver_initial_offset, test.input_width,
+                       test.input_height);
+  }
+}
+
+TEST_F(DownscaleTest, Trivial3DTest) {
+  constexpr uint16_t kChannels = 3;
+  constexpr uint16_t kInputHeight = 4;
+  constexpr uint16_t kInputWidth = 4;
+
+  constexpr uint16_t kVerScale = 2;
+  constexpr uint16_t kHorScale = 2;
+
+  constexpr uint16_t kOutputHeight = kInputHeight / kVerScale;
+  constexpr uint16_t kOutputWidth = kInputWidth / kHorScale;
+
+  setup(kChannels, kInputHeight, kInputWidth, kChannels, kOutputHeight,
+        kOutputWidth);
+
+  for (uint16_t c = 0; c < kChannels; ++c) {
+    for (uint16_t y = 0; y < kInputHeight; ++y) {
+      for (uint16_t x = 0; x < kInputWidth; ++x) {
+        *image_pixel(in_, c, y, x) = ((y * kInputWidth + x) * 10 + c) << 8;
+      }
+    }
+  }
+
+  params_.ver_initial_offset = 0;
+  params_.hor_initial_offset = 0;
+  params_.ver_scale_factor = kVerScale * kScaleFixedOne;
+  params_.hor_scale_factor = kHorScale * kScaleFixedOne;
+  set_downscale_param(&params_);
+
+  downscale_process(in_, out_);
+
+  // out_ is in interleaved order, so strides are different than defined in
+  // g_image.h
+  const uint16_t kInterleavedChStride = 1;
+  const uint16_t kInterleavedXStride = kChannels;
+  const uint32_t kInterleavedYStride = kChannels * kOutputWidth;
+
+  // for exact integer ratios out_ is just downsampled in_
+  for (uint16_t c = 0; c < kChannels; ++c) {
+    for (uint16_t y = 0; y < kOutputHeight; ++y) {
+      for (uint16_t x = 0; x < kOutputWidth; ++x) {
+        ASSERT_EQ(out_->data[y * kInterleavedYStride + x * kInterleavedXStride +
+                             c * kInterleavedChStride],
+                  image_pixel_val(in_, c, y * kVerScale, x * kHorScale) >> 8);
+      }
+    }
+  }
+}
diff --git a/samples/risp4ml/isp_stages/downscale_test.txt b/samples/risp4ml/isp_stages/downscale_test.txt
new file mode 100644
index 0000000..5cd708f
--- /dev/null
+++ b/samples/risp4ml/isp_stages/downscale_test.txt
@@ -0,0 +1,3 @@
+// RUN: ${TEST_RUNNER_CMD} ${OUT}/springbok_iree/samples/risp4ml/isp_stages/downscale_test
+// RUN: cat %t | FileCheck %s
+// CHECK: {{PASSED}}
diff --git a/samples/risp4ml/isp_stages/gamma_test.cc b/samples/risp4ml/isp_stages/gamma_test.cc
new file mode 100644
index 0000000..9adc32c
--- /dev/null
+++ b/samples/risp4ml/isp_stages/gamma_test.cc
@@ -0,0 +1,119 @@
+#include <cmath>
+
+#include "pw_unit_test/framework.h"
+#include "samples/risp4ml/isp_stages/gamma.h"
+
+static constexpr uint16_t kRgbPipelineBpp = 16;
+static constexpr uint16_t kPostGammaPipelineBpp = 16;
+GammaParams linear_params = {
+    .enable = true,
+    .lut = {0,     256,   512,   768,   1024,  1280,  1536,  1792,  2048,
+            2304,  2560,  2816,  3072,  3328,  3584,  3840,  4096,  4352,
+            4608,  4864,  5120,  5376,  5632,  5888,  6144,  6400,  6656,
+            6912,  7168,  7424,  7680,  7936,  8192,  8704,  9216,  9728,
+            10240, 10752, 11264, 11776, 12288, 12800, 13312, 13824, 14336,
+            14848, 15360, 15872, 16384, 17408, 18432, 19456, 20480, 21504,
+            22528, 23552, 24576, 25600, 26624, 27648, 28672, 29696, 30720,
+            31744, 32768, 34816, 36864, 38912, 40960, 43008, 45056, 47104,
+            49152, 51200, 53248, 55296, 57344, 59392, 61440, 63488, 65535}};
+
+GammaParams rgb_params = {
+    .enable = true,
+    .lut = {0,     3255,  5552,  7237,  8618,  9809,  10868, 11828, 12710,
+            13531, 14300, 15026, 15713, 16368, 16995, 17596, 18173, 18731,
+            19269, 19790, 20295, 20786, 21264, 21728, 22182, 22624, 23056,
+            23479, 23892, 24297, 24694, 25083, 25466, 26209, 26928, 27623,
+            28298, 28953, 29590, 30211, 30816, 31406, 31983, 32547, 33099,
+            33640, 34170, 34689, 35199, 36192, 37151, 38080, 38980, 39855,
+            40705, 41534, 42341, 43129, 43899, 44652, 45389, 46111, 46818,
+            47512, 48192, 49517, 50798, 52037, 53239, 54407, 55542, 56648,
+            57726, 58778, 59806, 60811, 61794, 62757, 63702, 64627, 65535}};
+
+class GammaTest : public ::testing::Test {
+ protected:
+  void setup(uint16_t width) {
+    in_ = image_new(3, 2, width);
+    out_ = image_new(3, 2, width);
+    for (uint16_t c = 0; c < in_->num_channels; ++c) {
+      for (uint16_t y = 0; y < in_->height; ++y) {
+        for (uint16_t x = 0; x < in_->width; ++x) {
+          *image_pixel(in_, c, y, x) = x;
+        }
+      }
+    }
+  }
+  void TearDown() override {
+    image_delete(in_);
+    image_delete(out_);
+  }
+  float sRgb_gamma(float in_) {
+    return (in_ < 0.0031308f) ? 12.92f * in_
+                              : 1.055f * std::pow(in_, 1.0f / 2.4f) - 0.055f;
+  }
+
+  Image* in_;
+  Image* out_;
+};
+
+TEST_F(GammaTest, Bypass) {
+  setup((1 << 15) - 1);
+
+  GammaParams params = rgb_params;
+  params.enable = false;
+
+  set_gamma_params(&params);
+
+  gamma_process(in_, out_);
+
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        pixel_type_t expected_val =
+            x >> (kRgbPipelineBpp - kPostGammaPipelineBpp);
+        ASSERT_EQ(expected_val, image_pixel_val(out_, c, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(GammaTest, Linear) {
+  setup((1 << 15) - 2);
+
+  set_gamma_params(&linear_params);
+
+  gamma_process(in_, out_);
+
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        ASSERT_EQ(image_pixel_val(out_, 0, y, x),
+                  image_pixel_val(in_, 0, y, x));
+      }
+    }
+  }
+}
+
+TEST_F(GammaTest, sRgbLUT) {
+  setup((1 << 15) - 1);
+
+  constexpr float kToleranceRatio = 0.03;
+
+  set_gamma_params(&rgb_params);
+
+  gamma_process(in_, out_);
+
+  for (uint16_t c = 0; c < in_->num_channels; ++c) {
+    for (uint16_t y = 0; y < in_->height; ++y) {
+      for (uint16_t x = 0; x < in_->width; ++x) {
+        pixel_type_t expected_val =
+            (pixel_type_t)((1 << kRgbPipelineBpp) *
+                           sRgb_gamma(static_cast<float>(x) /
+                                      (1 << kRgbPipelineBpp)));
+        float tolerance = ceilf(kToleranceRatio * expected_val);
+        float diff = std::abs((float)expected_val -
+                              (float)image_pixel_val(out_, c, y, x));
+        ASSERT_LE(diff, tolerance);
+      }
+    }
+  }
+}
diff --git a/samples/risp4ml/isp_stages/gamma_test.txt b/samples/risp4ml/isp_stages/gamma_test.txt
new file mode 100644
index 0000000..453890d
--- /dev/null
+++ b/samples/risp4ml/isp_stages/gamma_test.txt
@@ -0,0 +1,3 @@
+// RUN: ${TEST_RUNNER_CMD} ${OUT}/springbok_iree/samples/risp4ml/isp_stages/gamma_test
+// RUN: cat %t | FileCheck %s
+// CHECK: {{PASSED}}
diff --git a/samples/risp4ml/isp_stages/wbg_test.cc b/samples/risp4ml/isp_stages/wbg_test.cc
new file mode 100644
index 0000000..012a5c7
--- /dev/null
+++ b/samples/risp4ml/isp_stages/wbg_test.cc
@@ -0,0 +1,86 @@
+#include "pw_unit_test/framework.h"
+#include "samples/risp4ml/common/constants.h"
+#include "samples/risp4ml/common/test_utils.h"
+#include "samples/risp4ml/isp_stages/wbg.h"
+
+static constexpr uint16_t kBayerColorChannels = 4;
+static constexpr uint16_t kWbgFractional = kRawPipelineFraction;
+static constexpr uint16_t kFrameWidth = 4;
+static constexpr uint16_t kFrameHeight = 4;
+
+class WbgTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    in_ = image_new(1, kFrameHeight, kFrameWidth);
+    out_ = image_new(1, kFrameHeight, kFrameWidth);
+  }
+  void TearDown() override {
+    image_delete(in_);
+    image_delete(out_);
+  }
+
+  Image* in_;
+  Image* out_;
+};
+
+TEST_F(WbgTest, IdentityTest) {
+  // Use a grey input image.
+  InitImageRandom(in_, kRawPipelineMinVal, kRawPipelineMaxVal);
+
+  // set the params to something boring.
+  uint32_t gain = 1 << kWbgFractional;
+  WbgParams params = {
+      .enable = true, .fixed = true, .gains = {gain, gain, gain, gain}};
+  set_wbg_params(&params);
+
+  wbg_process(in_, out_);
+  for (uint16_t y = 0; y < kFrameHeight; y++) {
+    for (uint16_t x = 0; x < kFrameWidth; x++) {
+      ASSERT_EQ(image_pixel_val(in_, 0, y, x), image_pixel_val(out_, 0, y, x));
+    }
+  }
+}
+
+TEST_F(WbgTest, BayerTest) {
+  // Use a grey input image.
+  constexpr pixel_type_t kPixelVal = 1 << kWbgFractional;
+  InitImage(in_, kPixelVal);
+
+  // set the params to something boring.
+  WbgParams params = {.enable = true,
+                      .fixed = true,
+                      .gains = {kPixelVal, kPixelVal, kPixelVal, kPixelVal}};
+  set_wbg_params(&params);
+
+  wbg_process(in_, out_);
+
+  for (uint16_t c1 = 0; c1 < 2; ++c1) {
+    for (uint16_t c2 = 0; c2 < 2; ++c2) {
+      for (uint16_t y = 0; y < kFrameHeight / 2; y++) {
+        for (uint16_t x = 0; x < kFrameWidth / 2; x++) {
+          ASSERT_EQ(kPixelVal,
+                    image_pixel_val(out_, 0, 2 * y + c1, 2 * x + c2));
+        }
+      }
+    }
+  }
+}
+
+TEST_F(WbgTest, SaturationTest) {
+  // Use a grey input image.
+  constexpr pixel_type_t kPixelValHi = kRawPipelineMaxVal - 10;
+  InitImage(in_, kPixelValHi);
+
+  uint32_t gain = 2 * (1 << kWbgFractional);
+  WbgParams params = {
+      .enable = true, .fixed = true, .gains = {gain, gain, gain, gain}};
+  set_wbg_params(&params);
+
+  wbg_process(in_, out_);
+
+  for (uint16_t y = 0; y < kFrameHeight; y++) {
+    for (uint16_t x = 0; x < kFrameWidth; x++) {
+      ASSERT_EQ(kRawPipelineMaxVal, image_pixel_val(out_, 0, y, x));
+    }
+  }
+}
diff --git a/samples/risp4ml/isp_stages/wbg_test.txt b/samples/risp4ml/isp_stages/wbg_test.txt
new file mode 100644
index 0000000..7a35387
--- /dev/null
+++ b/samples/risp4ml/isp_stages/wbg_test.txt
@@ -0,0 +1,3 @@
+// RUN: ${TEST_RUNNER_CMD} ${OUT}/springbok_iree/samples/risp4ml/isp_stages/wbg_test
+// RUN: cat %t | FileCheck %s
+// CHECK: {{PASSED}}