text: // Copyright (c) Microsoft Corporation. All rights reserved.
  // Licensed under the MIT License.
  
  #include "gtest/gtest.h"
  #include "test/common/tensor_op_test_utils.h"
  #include "test/common/cuda_op_test_utils.h"
  #include "test/providers/provider_test_utils.h"
  
  namespace onnxruntime {
  namespace test {
  
  template <typename Integer, typename = typename std::enable_if<std::is_integral<Integer>::value, Integer>::type>
  inline std::vector<Integer> ToInteger(const std::vector<float>& data, float scale) {
    std::vector<Integer> result;
    result.reserve(data.size());
    for (size_t i = 0; i < data.size(); i++) {
      result.push_back(static_cast<Integer>(std::round(data[i] / scale)));
    }
    return result;
  }
  
  static void RunAttentionTest(
      const std::vector<float>& input_data,         // input:      [batch_size, sequence_length, hidden_size]
      const std::vector<float>& weights_data,       // weights:    [hidden_size, 3 * hidden_size]
      const std::vector<float>& bias_data,          // bias:       [3 * hidden_size]
      const std::vector<int32_t>& mask_index_data,  // mask_index: [batch_size]
      const std::vector<float>& output_data,        // output:     [batch_size, sequence_length, hidden_size]
      int batch_size,
      int sequence_length,
      int hidden_size,
      int number_of_heads,
      bool use_float16 = false) {
    int min_cuda_architecture = 530;
  
    bool enable_cuda = HasCudaEnvironment(min_cuda_architecture);
  
    if (enable_cuda) {
      OpTester tester("QAttention", 1, onnxruntime::kMSDomain);
      tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
  
      std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
      std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};
      std::vector<int64_t> bias_dims = {3 * hidden_size};
      std::vector<int64_t> mask_index_dims = {batch_size};
      std::vector<int64_t> output_dims = input_dims;
  
      float input_scale = 0.1f;
      float weight_scale = 0.1f;
      tester.AddInput<int8_t>("input", input_dims, ToInteger<int8_t>(input_data, input_scale));
      tester.AddInput<int8_t>("weight", weights_dims, ToInteger<int8_t>(weights_data, weight_scale));
      if (use_float16) {
        tester.AddInput<MLFloat16>("bias", bias_dims, ToFloat16(bias_data));
        tester.AddInput<MLFloat16>("input_scale", {1}, ToFloat16({input_scale}));
        tester.AddInput<MLFloat16>("weight_scale", {1}, ToFloat16({weight_scale}));
        tester.AddOutput<MLFloat16>("output", output_dims, ToFloat16(output_data));
      } else {
        tester.AddInput<float>("bias", bias_dims, bias_data);
        tester.AddInput<float>("input_scale", {1}, {input_scale});
        tester.AddInput<float>("weight_scale", {1}, {weight_scale});
        tester.AddOutput<float>("output", output_dims, output_data);
      }
  
      if (mask_index_data.size() > 0) {  // mask index is optional.
        tester.AddInput<int32_t>("mask_index", mask_index_dims, mask_index_data);
      }
  
      tester.Run();
    }
  }
  
  TEST(QAttentionTest, AttentionBatch1) {
    int batch_size = 1;
    int sequence_length = 2;
    int hidden_size = 4;
    int number_of_heads = 2;
  
    std::vector<float> input_data = {
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f};
  
    std::vector<float> weight_data = {
        0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
        0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
        0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
        0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
  
    std::vector<float> bias_data = {
        -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
  
    std::vector<int32_t> mask_index_data = {2L};
  
    std::vector<float> output_data = {
        3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
        3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f};
  
    RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
                     batch_size, sequence_length, hidden_size, number_of_heads);
  }
  
  TEST(QAttentionTest, AttentionBatch1_Float16) {
    int batch_size = 1;
    int sequence_length = 2;
    int hidden_size = 4;
    int number_of_heads = 2;
  
    std::vector<float> input_data = {
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f};
  
    std::vector<float> weight_data = {
        0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
        0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
        0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
        0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
  
    std::vector<float> bias_data = {
        -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
  
    std::vector<int32_t> mask_index_data = {2L};
  
    std::vector<float> output_data = {
        3.15039f, 0.1082763671875f, 4.24609375f, 5.6484375f,
        3.96679f, 0.072998046875f, 4.24609f, 5.6484375f};
  
    RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
                     batch_size, sequence_length, hidden_size, number_of_heads, true);
  }
  
  TEST(QAttentionTest, AttentionBatch2) {
    int batch_size = 2;
    int sequence_length = 2;
    int hidden_size = 4;
    int number_of_heads = 2;
  
    std::vector<float> input_data = {
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f,
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f};
  
    std::vector<float> weight_data = {
        0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
        0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
        0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
        0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
  
    std::vector<float> bias_data = {
        -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
  
    std::vector<int32_t> mask_index_data = {2L, 2L};
  
    std::vector<float> output_data = {
        3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
        3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f,
        3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
        3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f};
  
    RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
                     batch_size, sequence_length, hidden_size, number_of_heads);
  }
  
  TEST(QAttentionTest, AttentionMaskPartialSequence) {
    int batch_size = 1;
    int sequence_length = 2;
    int hidden_size = 4;
    int number_of_heads = 2;
  
    std::vector<float> input_data = {
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f};
  
    std::vector<float> weight_data = {
        0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
        0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
        0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
        0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
  
    std::vector<float> bias_data = {
        -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
  
    // Test mask_index < sequence_length
    std::vector<int32_t> mask_index_data = {1L};
  
    std::vector<float> output_data = {
        8.6899995803833008f, -0.13000002503395081f, 4.25f, 5.6499996185302734f,
        8.6899995803833008f, -0.13000002503395081f, 4.2499995231628418f, 5.6499991416931152f};
  
    RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
                     batch_size, sequence_length, hidden_size, number_of_heads);
  }
  
  TEST(QAttentionTest, AttentionMaskExceedSequence) {
    int batch_size = 1;
    int sequence_length = 2;
    int hidden_size = 4;
    int number_of_heads = 2;
  
    std::vector<float> input_data = {
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f};
  
    std::vector<float> weight_data = {
        0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
        0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
        0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
        0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
  
    std::vector<float> bias_data = {
        -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
  
    // Test mask_index > sequence_length
    std::vector<int32_t> mask_index_data = {3L};
  
    std::vector<float> output_data = {
        3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
        3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f};
  
    RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
                     batch_size, sequence_length, hidden_size, number_of_heads);
  }
  
  TEST(QAttentionTest, AttentionNoMaskIndex) {
    int batch_size = 1;
    int sequence_length = 2;
    int hidden_size = 4;
    int number_of_heads = 2;
  
    std::vector<float> input_data = {
        0.8f, -0.5f, 0.0f, 1.f,
        0.5f, 0.2f, 0.3f, -0.6f};
  
    std::vector<float> weight_data = {
        0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
        0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
        0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
        0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
  
    std::vector<float> bias_data = {
        -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
  
    // No mask_index
    std::vector<int32_t> mask_index_data = {};
  
    std::vector<float> output_data = {
        3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f,
        3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f};
  
    RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
                     batch_size, sequence_length, hidden_size, number_of_heads);
  }
  }  // namespace test
  }  // namespace onnxruntime
  