risp4ml/isp_stages/demosaic_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/demosaic.h"

static DemosaicParams demosaic_params = {.enable = true};

void set_demosaic_params(DemosaicParams* params) { demosaic_params = *params; }

// Basic bilinear demosaic
void demosaic_process(Image* input, Image* output) {
  if (!demosaic_params.enable) {
    return;
  }

  const pixel_type_t* line_in[RGB_COLOR_CHANNELS];
  pixel_type_t* line_out[RGB_COLOR_CHANNELS];
  int x_offset[RGB_COLOR_CHANNELS];
  const uint16_t boundary[RGB_COLOR_CHANNELS] = {0, input->width - 2,
                                                 input->width - 1};

  size_t vl;
  // auxiliary variables
  vuint16m4_t vx, vy;
  vuint32m8_t vz;

  for (uint16_t y = 0; y < input->height; ++y) {
    line_in[0] = (y) ? image_row(input, 0, y - 1) : image_row(input, 0, 1);
    line_in[1] = image_row(input, 0, y);
    line_in[2] = (y < input->height - 1)
                     ? image_row(input, 0, y + 1)
                     : image_row(input, 0, input->height - 2);

    line_out[1] = image_row(output, 1, y);
    if ((y & 0x1) == 0) {
      line_out[0] = image_row(output, 0, y);
      line_out[2] = image_row(output, 2, y);
    } else {
      line_out[0] = image_row(output, 2, y);
      line_out[2] = image_row(output, 0, y);
    }

    // x at boundary
    for (uint8_t i = 0; i < 3; ++i) {
      uint16_t x = boundary[i];
      for (uint16_t c = 0; c < RGB_COLOR_CHANNELS; ++c) {
        x_offset[c] = BayerMirrorBoundary(x - 1 + c, input->width);
      }

      BayerIndex bayer_index = GetBayerIndex(kBayerType, x, y);
      switch (bayer_index) {
        case (kB):
        case (kR): {
          line_out[0][x] = line_in[1][x_offset[1]];
          line_out[1][x] = (line_in[0][x_offset[1]] + line_in[2][x_offset[1]] +
                            line_in[1][x_offset[0]] + line_in[1][x_offset[2]]) /
                           4;
          line_out[2][x] = (line_in[0][x_offset[0]] + line_in[0][x_offset[2]] +
                            line_in[2][x_offset[0]] + line_in[2][x_offset[2]]) /
                           4;
        }; break;
        case (kGb):
        case (kGr): {
          line_out[0][x] =
              (line_in[1][x_offset[0]] + line_in[1][x_offset[2]]) / 2;
          line_out[1][x] = line_in[1][x_offset[1]];
          line_out[2][x] =
              (line_in[0][x_offset[1]] + line_in[2][x_offset[1]]) / 2;
        }; break;
        default:
          break;
      }
    }

    // x not at boundary: vector instructions
    for (uint8_t n = 1; n <= 2; n++) {
      for (uint16_t x = n; x < input->width - 2; x += 2 * vl) {
        x_offset[0] = x - 1;
        x_offset[1] = x;
        x_offset[2] = x + 1;
        ptrdiff_t stride = 2 * sizeof(uint16_t);
        size_t avl = (input->width - 1 - x) / 2;
        vl = __riscv_vsetvl_e16m4(avl);

        if (n + (y & 0x1) == 2) {  // kR or kB
          // ch0
          vx = __riscv_vlse16_v_u16m4(line_in[1] + x_offset[1], stride,
                                      vl);                  // load
          __riscv_vsse16(line_out[0] + x, stride, vx, vl);  // save
          // ch1
          vx = __riscv_vlse16_v_u16m4(line_in[0] + x_offset[1], stride,
                                      vl);  // load
          vy = __riscv_vlse16_v_u16m4(line_in[2] + x_offset[1], stride,
                                      vl);     // load
          vz = __riscv_vwaddu_vv(vx, vy, vl);  // add
          vy = __riscv_vlse16_v_u16m4(line_in[1] + x_offset[0], stride,
                                      vl);     // load
          vz = __riscv_vwaddu_wv(vz, vy, vl);  // add
          vy = __riscv_vlse16_v_u16m4(line_in[1] + x_offset[2], stride,
                                      vl);                  // load
          vz = __riscv_vwaddu_wv(vz, vy, vl);               // add
          vx = __riscv_vnsrl(vz, 2, vl);                    // 1/4
          __riscv_vsse16(line_out[1] + x, stride, vx, vl);  // save
          // ch2
          vx = __riscv_vlse16_v_u16m4(line_in[0] + x_offset[0], stride,
                                      vl);  // load
          vy = __riscv_vlse16_v_u16m4(line_in[0] + x_offset[2], stride,
                                      vl);     // load
          vz = __riscv_vwaddu_vv(vx, vy, vl);  // add
          vy = __riscv_vlse16_v_u16m4(line_in[2] + x_offset[0], stride,
                                      vl);     // load
          vz = __riscv_vwaddu_wv(vz, vy, vl);  // add
          vy = __riscv_vlse16_v_u16m4(line_in[2] + x_offset[2], stride,
                                      vl);                  // load
          vz = __riscv_vwaddu_wv(vz, vy, vl);               // add
          vx = __riscv_vnsrl(vz, 2, vl);                    // 1/4
          __riscv_vsse16(line_out[2] + x, stride, vx, vl);  // save
        } else {  // kGr or kRb
          // ch0
          vx = __riscv_vlse16_v_u16m4(line_in[1] + x_offset[0], stride,
                                      vl);  // load
          vy = __riscv_vlse16_v_u16m4(line_in[1] + x_offset[2], stride,
                                      vl);                  // load
          vz = __riscv_vwaddu_vv(vx, vy, vl);               // add
          vx = __riscv_vnsrl(vz, 1, vl);                    // 1/2
          __riscv_vsse16(line_out[0] + x, stride, vx, vl);  // save
          // ch1
          vx = __riscv_vlse16_v_u16m4(line_in[1] + x_offset[1], stride,
                                      vl);                  // load
          __riscv_vsse16(line_out[1] + x, stride, vx, vl);  // save
          // ch2
          vx = __riscv_vlse16_v_u16m4(line_in[0] + x_offset[1], stride,
                                      vl);  // load
          vy = __riscv_vlse16_v_u16m4(line_in[2] + x_offset[1], stride,
                                      vl);                  // load
          vz = __riscv_vwaddu_vv(vx, vy, vl);               // add
          vx = __riscv_vnsrl(vz, 1, vl);                    // 1/2
          __riscv_vsse16(line_out[2] + x, stride, vx, vl);  // save
        }
      }
    }
  }
}