risp4ml/isp_stages/wbg_rvv.c

/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <riscv_vector.h>

#include "risp4ml/common/utils.h"
#include "risp4ml/isp_stages/wbg.h"

#define MAX(a, b) (((a) > (b)) ? (a) : (b))

static const uint16_t kWbgFractional = kRawPipelineFraction;
static const uint16_t kWbgUnityGain = 1 << kWbgFractional;
static WbgParams wbg_params = {
    .enable = true,
    .fixed = false,
    .gains = {kWbgUnityGain, kWbgUnityGain, kWbgUnityGain, kWbgUnityGain}};

void set_wbg_params(WbgParams* params) { wbg_params = *params; }

static void compute_wbg_gain(Image* img) {
  // Calculate the white-balance gain values using the "gray world" algorithm
  uint32_t size = img->num_channels * img->height * img->width;
  uint64_t sum_of_reds = 0;
  // will use only one of the greens for scaling, since the difference between
  // the two green sensor pixels is negligible
  uint64_t sum_of_greens = 0;
  uint64_t sum_of_blues = 0;
  size_t vl;
  // auxiliary variables
  vuint16m8_t vx;
  vuint32m1_t vy;

  for (uint16_t y = 0; y < img->height; ++y) {
    pixel_type_t* line = image_row(img, 0, y);
    for (uint8_t n = 0; n < 2; n++) {
      for (uint16_t x = n; x < img->width; x += 2 * vl) {
        size_t avl = (img->width + 1 - x) / 2;
        vl = __riscv_vsetvl_e16m8(avl);
        vx =
            __riscv_vlse16_v_u16m8(line + x, 2 * sizeof(uint16_t), vl);  // load

        vy = __riscv_vmv_v_x_u32m1(0, vl);   // init
        vy = __riscv_vwredsumu(vx, vy, vl);  // sum
        uint32_t sum = __riscv_vmv_x(vy);
        if ((y & 0x1) == 0 && n == 0) {
          sum_of_reds += sum;
        } else if ((y & 0x1) == 1 && n == 1) {
          sum_of_blues += sum;
        } else {
          sum_of_greens += sum;
        }
      }
    }
  }

  // scale values to green channel
  float average_red = 4.0 * sum_of_reds / size;
  float average_green = 2.0 * sum_of_greens / size;
  float average_blue = 4.0 * sum_of_blues / size;

  float max_average = MAX(MAX(average_red, average_green), average_blue);

  // Convert the float value to fixed point representation, i.e. 0xFF.FF
  uint32_t red_wb = FloatToFixedPoint(max_average / average_red,
                                      kRawPipelineInteger, kRawPipelineFraction,
                                      /*bool is_signed*/ false);
  uint32_t green_wb = FloatToFixedPoint(
      max_average / average_green, kRawPipelineInteger, kRawPipelineFraction,
      /*bool is_signed*/ false);
  uint32_t blue_wb = FloatToFixedPoint(
      max_average / average_blue, kRawPipelineInteger, kRawPipelineFraction,
      /*bool is_signed*/ false);

  wbg_params.gains[0] = red_wb;
  wbg_params.gains[1] = green_wb;
  wbg_params.gains[2] = green_wb;
  wbg_params.gains[3] = blue_wb;
}

void wbg_process(Image* img) {
  if (!wbg_params.enable) return;
  if (!wbg_params.fixed) {
    compute_wbg_gain(img);
  }

  size_t vl;
  uint32_t offset = 1 << (kWbgFractional - 1);
  uint32_t max_val = kRawPipelineMaxVal << kWbgFractional;
  uint16_t gain;
  // auxiliary variables
  vuint16m4_t vx;
  vuint32m8_t vy;
  for (uint16_t y = 0; y < img->height; ++y) {
    pixel_type_t* line = image_row(img, 0, y);
    for (uint8_t n = 0; n < 2; n++) {
      gain = (y & 0x1) ? wbg_params.gains[2 + n] : wbg_params.gains[n];
      for (uint16_t x = n; x < img->width; x += 2 * vl) {
        size_t avl = (img->width + 1 - x) / 2;
        vl = __riscv_vsetvl_e16m4(avl);
        vx =
            __riscv_vlse16_v_u16m4(line + x, 2 * sizeof(uint16_t), vl);  // load
        vy = __riscv_vwmulu(vx, gain, vl);                       // multiply
        vy = __riscv_vadd(vy, offset, vl);                       // add
        vy = __riscv_vminu(vy, max_val, vl);                     // clamp
        vx = __riscv_vnsrl(vy, kWbgFractional, vl);              // bit shift
        __riscv_vsse16(line + x, 2 * sizeof(uint16_t), vx, vl);  // save
      }
    }
  }
}