// 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 <springbok.h>
#include <stdio.h>
#include <stdlib.h>

#include "pw_unit_test/framework.h"
#include "test_v_helpers/include/test_v_helpers.h"

namespace vsetvl_test {
namespace {

using namespace test_v_helpers;

const uint64_t VLEN = 512u;
const uint32_t VILL = 0x80000000;

#define MIN(X, Y) ((X) < (Y) ? (X) : (Y))

static inline uint32_t read_vtype() {
  uint32_t vtype;
  __asm__ volatile("csrr %[VTYPE], vtype" : [VTYPE] "=r"(vtype) :);
  return vtype;
}

static inline bool is_illegal_vtype() {
  uint32_t vtype = read_vtype();
  if (vtype & VILL) {
    return true;
  }
  return false;
}

static uint32_t calculate_vl(uint32_t sew, uint32_t avl, float lmul) {
  uint32_t vlmax = (uint32_t)(VLEN * lmul / sew);
  return MIN(avl, vlmax);
}

static void test_vsetvl(VSEW sew, VLMUL vlmul, uint32_t width, float lmul,
                        bool expect_illegal = false) {
  size_t vl = 0;
  for (int i = 0; i < AVL_COUNT; i++) {
    vl = set_vsetvl_intrinsic(sew, vlmul, AVLS[i]);
    if (!expect_illegal) {
      EXPECT_EQ(vl, calculate_vl(width, AVLS[i], lmul));
    } else {
      EXPECT_TRUE(is_illegal_vtype());
    }
  }
  for (int i = 0; i < AVL_COUNT; i++) {
    vl = set_vsetvl(sew, vlmul, AVLS[i], false, false);
    if (!expect_illegal) {
      EXPECT_EQ(vl, calculate_vl(width, AVLS[i], lmul));
    } else {
      EXPECT_TRUE(is_illegal_vtype());
    }
  }
  for (int i = 0; i < AVL_COUNT; i++) {
    vl = set_vsetvli(sew, vlmul, AVLS[i]);
    if (!expect_illegal) {
      EXPECT_EQ(vl, calculate_vl(width, AVLS[i], lmul));
    } else {
      EXPECT_TRUE(is_illegal_vtype());
    }
  }
  vl = set_vsetivli(sew, vlmul);
  if (!expect_illegal) {
    EXPECT_EQ(vl, calculate_vl(width, AVL_CONST, lmul));
  } else {
    EXPECT_TRUE(is_illegal_vtype());
  }
}

static void test_vsetvlmax(VSEW sew, VLMUL vlmul, uint32_t width, float lmul,
                           bool expect_illegal = false) {
  if (!expect_illegal) {
    EXPECT_EQ(get_vsetvlmax_intrinsic(sew, vlmul), (int)VLEN / width * lmul);
  } else {
    EXPECT_EQ(get_vsetvlmax_intrinsic(sew, vlmul), 0);
  }
}

TEST(VsetvlTest, vsetvl_e8m1) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M1, 8, 1.0);
}

TEST(VsetvlTest, vsetvl_e16m1) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M1, 16, 1.0);
}

TEST(VsetvlTest, vsetvl_e32m1) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M1, 32, 1.0);
}

TEST(VsetvlTest, vsetvl_e8m2) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M2, 8, 2.0);
}

TEST(VsetvlTest, vsetvl_e16m2) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M2, 16, 2.0);
}

TEST(VsetvlTest, vsetvl_e32m2) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M2, 32, 2.0);
}

TEST(VsetvlTest, vsetvl_e8m4) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M4, 8, 4.0);
}

TEST(VsetvlTest, vsetvl_e16m4) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M4, 16, 4.0);
}

TEST(VsetvlTest, vsetvl_e32m4) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M4, 32, 4.0);
}

TEST(VsetvlTest, vsetvl_e8m8) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M8, 8, 8.0);
}

TEST(VsetvlTest, vsetvl_e16m8) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M8, 16, 8);
}

TEST(VsetvlTest, vsetvl_e32m8) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M8, 32, 8.0);
}

TEST(VsetvlTest, vsetvl_e8mf2) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_MF2, 8, 0.5);
}

TEST(VsetvlTest, vsetvl_e8mf4) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_MF4, 8, 0.25);
}

TEST(VsetvlTest, vsetvl_e16mf2) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_MF2, 16, 0.5);
}

TEST(VsetvlTest, vsetvl_e8mf8_vill) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_MF8, 8, 0.125, true);
}

TEST(VsetvlTest, vsetvl_e16mf4_vill) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_MF4, 16, 0.25, true);
}

TEST(VsetvlTest, vsetvl_e32mf2_vill) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_MF2, 32, 0.5, true);
}

TEST(VsetvlTest, vsetvlmax_e8m1) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M1, 8, 1.0);
}

TEST(VsetvlTest, vsetvlmax_e8m2) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M2, 8, 2.0);
}

TEST(VsetvlTest, vsetvlmax_e8m4) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M4, 8, 4.0);
}

TEST(VsetvlTest, vsetvlmax_e8m8) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M8, 8, 8.0);
}

TEST(VsetvlTest, vsetvlmax_e16m1) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M1, 16, 1.0);
}

TEST(VsetvlTest, vsetvlmax_e16m2) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M2, 16, 2.0);
}

TEST(VsetvlTest, vsetvlmax_e16m4) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M4, 16, 4.0);
}

TEST(VsetvlTest, vsetvlmax_e16m8) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M8, 16, 8.0);
}

TEST(VsetvlTest, vsetvlmax_e32m1) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M1, 32, 1.0);
}

TEST(VsetvlTest, vsetvlmax_e32m2) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M2, 32, 2.0);
}

TEST(VsetvlTest, vsetvlmax_e32m4) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M4, 32, 4.0);
}

TEST(VsetvlTest, vsetvlmax_e32m8) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M8, 32, 8.0);
}

TEST(VsetvlTest, vsetvlmax_e8mf2) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_MF2, 8, 0.5);
}

TEST(VsetvlTest, vsetvlmax_e8mf4) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_MF4, 8, 0.25);
}

TEST(VsetvlTest, vsetvlmax_e16mf2) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_MF2, 16, 0.5);
}

TEST(VsetvlTest, vsetvlmax_e8mf8_vill) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_MF8, 8, 0.125, true);
}

TEST(VsetvlTest, vsetvlmax_e16mf4_vill) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_MF4, 16, 0.25, true);
}

TEST(VsetvlTest, vsetvlmax_e32mf2_vill) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_MF2, 32, 0.5, true);
}

}  // namespace
}  // namespace vsetvl_test