text: /**
   * \file
   * dnn/test/common/dct_ref.cpp
   * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
   *
   * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
   * implied.
   */
  #include "test/common/dct_ref.h"
  namespace megdnn {
  namespace test {
  struct FixCase {
      std::vector<int> mask_offset;
      std::vector<int> mask_val;
  };
  using Param = DctChannelSelectForward::Param;
  
  static inline FixCase get_fix_mask(Param::FastImpl impl) {
      std::vector<int> fix_32_mask_offset{0, 16, 24, 32};
      std::vector<int> fix_32_mask_val{0,  1,  8,  16, 9, 2,  3, 10, 17, 24, 32,
                                       25, 18, 11, 4,  5, 0,  1, 8,  16, 9,  2,
                                       3,  10, 0,  1,  8, 16, 9, 2,  3,  10};
      megdnn_assert(impl == Param::FastImpl::FIX_32_MASK, "only support gen FIX_32_MASK");
      return {fix_32_mask_offset, fix_32_mask_val};
  }
  
  CheckerHelper::TensorsConstriant gen_dct_constriant(
          const size_t /* n */, const size_t ic, const size_t ih, const size_t iw,
          const size_t oc, Param param) {
      auto constraint = [=](CheckerHelper::TensorValueArray& tensors_orig) {
          const size_t block = param.dct_block_size;
          const int block_c = param.format == Param::Format::NCHW4 ? 4 : 1;
          megdnn_assert(oc % block_c == 0, "oc mod block_c must == 0");
          std::shared_ptr<DctTestcase> test_case_ptr = DctTestcase::make();
          DctTestcase& test_case = *test_case_ptr.get();
          UniformIntRNG rng(0, 255);
          UniformIntRNG mask_rng(0, 64 / block_c - 1);
          const size_t no_mask_oc = ic * block * block;
          megdnn_assert(ih % block == 0, "%zu mod %zu == 0", ih, block);
          megdnn_assert(iw % block == 0, "%zu mod %zu == 0", iw, block);
  
          TensorND mask_offset;
          TensorND mask_val;
          std::vector<int>& mask_offset_vec = test_case.mask_offset_vec;
          std::vector<int>& mask_val_vec = test_case.mask_val_vec;
          UniformIntRNG rng_oc(0, oc);
          if (param.fastImpl == Param::FastImpl::FIX_32_MASK) {
              auto fix_32_mask = get_fix_mask(Param::FastImpl::FIX_32_MASK);
              mask_offset_vec = fix_32_mask.mask_offset;
              mask_val_vec = fix_32_mask.mask_val;
              megdnn_assert(oc == 32, "oc must eq 32");
          } else if (no_mask_oc > oc) {
              size_t remain_oc = oc;
              mask_offset_vec.resize(ic + 1);
              mask_val_vec.resize(oc);
              mask_offset_vec[0] = 0;
              for (size_t ic_idx = 0; ic_idx < ic; ++ic_idx) {
                  size_t random_len = (int)rng_oc.gen_single_val() * block_c;
                  size_t mask_len = (ic_idx == ic - 1) || (remain_oc == 0)
                                          ? remain_oc
                                          : random_len % remain_oc;
                  megdnn_assert(
                          mask_len % block_c == 0,
                          "mask_len mod block_c == 0, but %zu mod %d ", mask_len,
                          block_c);
                  const size_t oc_idx = mask_offset_vec[ic_idx];
                  remain_oc -= mask_len;
                  mask_offset_vec[ic_idx + 1] = oc_idx + mask_len;
                  for (size_t mask_idx = 0; mask_idx < mask_len; ++mask_idx) {
                      mask_val_vec[oc_idx + mask_idx] = (int)mask_rng.gen_single_val();
                  }
              }
          }
          mask_offset = TensorND(
                  mask_offset_vec.data(), {{mask_offset_vec.size()}, dtype::Int32()});
          mask_val =
                  TensorND(mask_val_vec.data(), {{mask_val_vec.size()}, dtype::Int32()});
          if (tensors_orig.size() > 1) {
              megdnn_assert(tensors_orig.size() == 4, "tensors_orig.size() == 4");
              megdnn_assert(mask_offset_vec.size() >= 2, "mask_offset_vec.size() >= 2");
              megdnn_assert(
                      tensors_orig[1].layout == mask_offset.layout,
                      "tensors_orig[1].layout == mask_offset.layout");
              megdnn_assert(
                      tensors_orig[2].layout == mask_val.layout,
                      "tensors_orig[2].layout == mask_val.layout");
              auto naive_handle = create_cpu_handle(2, false);
              megdnn_memcpy_D2D(
                      naive_handle.get(), tensors_orig[1].raw_ptr, mask_offset.raw_ptr,
                      mask_offset.layout.span().dist_byte());
              megdnn_memcpy_D2D(
                      naive_handle.get(), tensors_orig[2].raw_ptr, mask_val.raw_ptr,
                      mask_val.layout.span().dist_byte());
          }
      };
      return constraint;
  }
  
  std::shared_ptr<DctTestcase> gen_dct_case(
          const size_t n, const size_t ic, const size_t ih, const size_t iw,
          const size_t oc, Param param, DType dst_dtype, bool correct_result) {
      const size_t block = param.dct_block_size;
      const int block_c = param.format == Param::Format::NCHW4 ? 4 : 1;
      megdnn_assert(oc % block_c == 0, "oc mod block_c must == 0");
      std::shared_ptr<DctTestcase> test_case_ptr = DctTestcase::make();
      DctTestcase& test_case = *test_case_ptr.get();
      UniformIntRNG rng(0, 255);
      UniformIntRNG mask_rng(0, 64 / block_c - 1);
      const size_t input_elements = n * ic * ih * iw;
      const size_t no_mask_oc = ic * block * block;
      megdnn_assert(ih % block == 0, "%zu mod %zu == 0", ih, block);
      megdnn_assert(iw % block == 0, "%zu mod %zu == 0", iw, block);
      std::vector<uint8_t>& inp_vec = test_case.inp_vec;
      inp_vec.resize(input_elements);
      TensorShape input_shape{n, ic, ih, iw};
      for (auto& elm : inp_vec) {
          elm = (uint8_t)rng.gen_single_val();
      }
      auto src = TensorND(inp_vec.data(), {input_shape, dtype::Uint8()});
      TensorND mask_offset;
      TensorND mask_val;
      std::vector<int>& mask_offset_vec = test_case.mask_offset_vec;
      std::vector<int>& mask_val_vec = test_case.mask_val_vec;
      UniformIntRNG rng_oc(0, oc);
      if (param.fastImpl == Param::FastImpl::FIX_32_MASK) {
          auto fix_32_mask = get_fix_mask(Param::FastImpl::FIX_32_MASK);
          mask_offset_vec = fix_32_mask.mask_offset;
          mask_val_vec = fix_32_mask.mask_val;
          megdnn_assert(oc == 32, "oc must eq 32");
      } else if (no_mask_oc > oc) {
          size_t remain_oc = oc;
          mask_offset_vec.resize(ic + 1);
          mask_val_vec.resize(oc);
          mask_offset_vec[0] = 0;
          for (size_t ic_idx = 0; ic_idx < ic; ++ic_idx) {
              size_t random_len = (int)rng_oc.gen_single_val() * block_c;
              size_t mask_len = (ic_idx == ic - 1) || (remain_oc == 0)
                                      ? remain_oc
                                      : random_len % remain_oc;
              megdnn_assert(
                      mask_len % block_c == 0,
                      "mask_len mod block_c == 0, but %zu mod %d ", mask_len, block_c);
              const size_t oc_idx = mask_offset_vec[ic_idx];
              remain_oc -= mask_len;
              mask_offset_vec[ic_idx + 1] = oc_idx + mask_len;
              for (size_t mask_idx = 0; mask_idx < mask_len; ++mask_idx) {
                  mask_val_vec[oc_idx + mask_idx] = (int)mask_rng.gen_single_val();
              }
          }
      }
      mask_offset = TensorND(
              mask_offset_vec.data(), {{mask_offset_vec.size()}, dtype::Int32()});
      mask_val = TensorND(mask_val_vec.data(), {{mask_val_vec.size()}, dtype::Int32()});
      if (mask_offset_vec.size() >= 2) {
          test_case.testcase_in = {
                  src, mask_offset, mask_val, {nullptr, {{}, dst_dtype}}};
      } else {
          test_case.testcase_in = {src, {}, {}, {nullptr, {{}, dst_dtype}}};
      }
  
      auto naive_handle = create_cpu_handle(2, false);
      auto opr_naive = naive_handle->create_operator<DctChannelSelectForward>();
      opr_naive->param() = param;
      using Proxy = OprProxy<DctChannelSelectForward>;
      Proxy naive_proxy;
      TensorLayout temp_dst_layout;
      temp_dst_layout.dtype = dst_dtype;
  
      TensorLayoutArray layouts{
              src.layout, mask_offset.layout, mask_val.layout, temp_dst_layout};
      naive_proxy.deduce_layout(opr_naive.get(), layouts);
      const size_t output_elements = layouts[3].total_nr_elems();
      std::vector<float>& output_vec = test_case.output_vec;
      output_vec.resize(output_elements);
      auto dst = TensorND(output_vec.data(), layouts[3]);
      DctTestcase::TensorValueArray testcase_naive;
      testcase_naive.emplace_back(test_case.testcase_in[0]);
      testcase_naive.emplace_back(test_case.testcase_in[1]);
      testcase_naive.emplace_back(test_case.testcase_in[2]);
      testcase_naive.emplace_back(dst);
      if (correct_result) {
          naive_proxy.exec(opr_naive.get(), testcase_naive);
      }
      test_case.testcase_out = {{}, {}, {}, dst};
      return test_case_ptr;
  }
  
  }  // namespace test
  }  // namespace megdnn
     // vim: syntax=cpp.doxygen