// Copyright 2024 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <float.h>
#include <inttypes.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <vector>

#include "iree/base/api.h"
#include "iree/base/internal/cpu.h"
#include "iree/base/internal/math.h"
#include "iree/base/internal/path.h"
#include "iree/base/tooling/flags.h"
#include "iree/hal/api.h"
#include "iree/modules/hal/module.h"
#include "iree/tooling/context_util.h"
#include "iree/tooling/device_util.h"
#include "iree/vm/api.h"
#include "iree/vm/native_module_cc.h"
#include "tools/testing/e2e/test_utils.h"

//===----------------------------------------------------------------------===//
// Reference Attention
//===----------------------------------------------------------------------===//

// Helper for reference_attention.
// Function to allocate and initialize tensors
float* allocate_tensor(int dim1, int dim2, int dim3) {
  const int size = dim1 * dim2 * dim3;
  float* tensor = (float*)malloc(size * sizeof(float));
  for (int i = 0; i < size; ++i) {
    tensor[i] = 0.0f;
  }
  return tensor;
}

// Function to free allocated tensors
void free_tensor(float* tensor) {
  if (tensor != nullptr) free(tensor);
}

// Function to calculate 1D index for a 3D array
int index_3d(int i, int j, int k, int dim2, int dim3) {
  return i * dim2 * dim3 + j * dim3 + k;
}

static void reference_attention_f32_f32_f32_f32(
    iree_hal_dim_t M, iree_hal_dim_t K1, iree_hal_dim_t K2, iree_hal_dim_t N,
    iree_hal_dim_t B, float scale, const float* query_data,
    const float* key_data, const float* value_data, float* result_data,
    iree_hal_dim_t b, float* Attention, const uint8_t* mask_data) {
  // Compute Q * K^T
  for (int m = 0; m < M; ++m) {
    for (int k2 = 0; k2 < K2; ++k2) {
      float sum = 0.0;
      for (int k1 = 0; k1 < K1; ++k1) {
        int q_idx = index_3d(b, m, k1, M, K1);
        int k_idx = index_3d(b, k2, k1, K2, K1);

        sum += query_data[q_idx] * key_data[k_idx];
      }
      int att_idx = index_3d(0, m, k2, M, K2);
      Attention[att_idx] = sum * scale;
    }
  }

  // Apply mask before softmax if provided.
  // mask_data is a boolean mask where 1 = attend, 0 = mask out (-inf)
  if (mask_data != nullptr) {
    for (int m = 0; m < M; ++m) {
      for (int k2 = 0; k2 < K2; ++k2) {
        int mask_idx = index_3d(b, m, k2, M, K2);
        int att_idx = index_3d(0, m, k2, M, K2);
        if (mask_data[mask_idx] == 0) {
          Attention[att_idx] = -INFINITY;
        }
      }
    }
  }

  // Compute softmax on Attention
  for (int m = 0; m < M; ++m) {
    // Find the maximum value for the current sequence
    float max_val = -FLT_MAX;
    for (int k2 = 0; k2 < K2; ++k2) {
      int att_idx = index_3d(0, m, k2, M, K2);
      max_val = iree_max(max_val, Attention[att_idx]);
    }

    // Calculate the softmax denominator
    float sum = 0.0f;
    for (int k2 = 0; k2 < K2; ++k2) {
      int att_idx = index_3d(0, m, k2, M, K2);
      sum += exp(Attention[att_idx] - max_val);
    }

    // Apply softmax
    for (int k2 = 0; k2 < K2; ++k2) {
      int att_idx = index_3d(0, m, k2, M, K2);
      Attention[att_idx] = exp(Attention[att_idx] - max_val) / sum;
    }
  }

  // Compute Attention * V
  for (int m = 0; m < M; ++m) {
    for (int n = 0; n < N; ++n) {
      float sum = 0.0;
      for (int k2 = 0; k2 < K2; ++k2) {
        int att_idx = index_3d(0, m, k2, M, K2);
        int v_idx = index_3d(b, k2, n, K2, N);
        sum += Attention[att_idx] * value_data[v_idx];
      }
      int o_idx = index_3d(b, m, n, M, N);
      result_data[o_idx] = sum;
    }
  }
}

static iree_status_t reference_attention_element(
    iree_hal_dim_t M, iree_hal_dim_t K1, iree_hal_dim_t K2, iree_hal_dim_t N,
    iree_hal_dim_t B, float scale, iree_hal_element_type_t query_elem_type,
    iree_hal_element_type_t key_elem_type,
    iree_hal_element_type_t value_elem_type, void* query_data, void* key_data,
    void* value_data, void* result_data, iree_hal_dim_t b, float* Attention,
    const uint8_t* mask_data) {
  if (query_elem_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 &&
      key_elem_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 &&
      value_elem_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) {
    reference_attention_f32_f32_f32_f32(
        M, K1, K2, N, B, scale, (const float*)query_data,
        (const float*)key_data, (const float*)value_data, (float*)result_data,
        b, Attention, mask_data);

  } else {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "unhandled combination of element types in attention");
  }
  return iree_ok_status();
}

// Reference attention implementation, used to compare attention results
// against.
static iree_status_t reference_attention(
    iree_hal_dim_t B, iree_hal_dim_t M, iree_hal_dim_t K1, iree_hal_dim_t K2,
    iree_hal_dim_t N, float scale, iree_hal_element_type_t query_elem_type,
    iree_hal_element_type_t key_elem_type,
    iree_hal_element_type_t value_elem_type, iree_byte_span_t query_contents,
    iree_byte_span_t key_contents, iree_byte_span_t value_contents,
    iree_byte_span_t result_contents, iree_byte_span_t mask_contents,
    int compute_every) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, B);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, M);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, K1);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, K2);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, N);

  iree_host_size_t count = 0;
  float* Attention = allocate_tensor(1, M, K2);
  // mask_data is nullptr if no mask is provided.
  const uint8_t* mask_data = mask_contents.data_length > 0
                                 ? (const uint8_t*)mask_contents.data
                                 : nullptr;
  for (iree_hal_dim_t b = 0; b < B; ++b) {
    if (++count < compute_every) continue;
    count = 0;
    iree_status_t status = reference_attention_element(
        M, K1, K2, N, B, scale, query_elem_type, key_elem_type, value_elem_type,
        query_contents.data, key_contents.data, value_contents.data,
        result_contents.data, b, Attention, mask_data);
    if (!iree_status_is_ok(status)) {
      free_tensor(Attention);
      IREE_TRACE_ZONE_END(z0);
      return status;
    }
  }
  free_tensor(Attention);

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}
//===----------------------------------------------------------------------===//
// Attention comparison/logging
//===----------------------------------------------------------------------===//

typedef struct {
  iree_allocator_t host_allocator;
  iree_hal_dim_t b;
  iree_hal_dim_t m;
  iree_hal_dim_t k1;
  iree_hal_dim_t k2;
  iree_hal_dim_t n;
  // Scale applied to Q*K^T by the attention op under test. The reference must
  // use the exact same value or the softmax distributions will not match.
  float scale;
  iree_hal_element_type_t query_elem_type;
  iree_hal_element_type_t key_elem_type;
  iree_hal_element_type_t value_elem_type;
  iree_hal_element_type_t result_elem_type;
  iree_byte_span_t query_contents;
  iree_byte_span_t key_contents;
  iree_byte_span_t value_contents;
  iree_byte_span_t actual_contents;
  iree_byte_span_t expected_contents;
  iree_byte_span_t mask_contents;  // Optional: empty if no mask.
} attention_results_t;

static void attention_results_deinitialize(attention_results_t* results);

static iree_status_t attention_results_initialize(
    iree_hal_device_t* device, iree_hal_dim_t b_size, iree_hal_dim_t m_size,
    iree_hal_dim_t k1_size, iree_hal_dim_t k2_size, iree_hal_dim_t n_size,
    float scale, iree_hal_buffer_view_t* query, iree_hal_buffer_view_t* key,
    iree_hal_buffer_view_t* value, iree_hal_buffer_view_t* result,
    iree_hal_buffer_view_t* mask,  // Optional: can be nullptr.
    iree_allocator_t host_allocator, attention_results_t* out_results) {
  IREE_TRACE_ZONE_BEGIN(z0);

  memset(out_results, 0, sizeof(*out_results));
  out_results->host_allocator = host_allocator;

  out_results->b = b_size;
  out_results->m = m_size;
  out_results->k1 = k1_size;
  out_results->k2 = k2_size;
  out_results->n = n_size;
  out_results->scale = scale;

  out_results->query_elem_type = iree_hal_buffer_view_element_type(query);
  out_results->key_elem_type = iree_hal_buffer_view_element_type(key);
  out_results->value_elem_type = iree_hal_buffer_view_element_type(value);
  out_results->result_elem_type = iree_hal_buffer_view_element_type(result);

  iree_hal_buffer_t* query_buffer = iree_hal_buffer_view_buffer(query);
  iree_hal_buffer_t* key_buffer = iree_hal_buffer_view_buffer(key);
  iree_hal_buffer_t* value_buffer = iree_hal_buffer_view_buffer(value);
  iree_hal_buffer_t* result_buffer = iree_hal_buffer_view_buffer(result);
  iree_hal_buffer_t* mask_buffer =
      mask ? iree_hal_buffer_view_buffer(mask) : nullptr;

  iree_status_t status = iree_ok_status();

  if (iree_status_is_ok(status)) {
    out_results->query_contents.data_length =
        iree_hal_buffer_byte_length(query_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->query_contents.data_length,
                                   (void**)&out_results->query_contents.data);
  }
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, query_buffer, 0, out_results->query_contents.data,
        out_results->query_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }
  if (iree_status_is_ok(status)) {
    out_results->key_contents.data_length =
        iree_hal_buffer_byte_length(key_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->key_contents.data_length,
                                   (void**)&out_results->key_contents.data);
  }
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, key_buffer, 0, out_results->key_contents.data,
        out_results->key_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }
  if (iree_status_is_ok(status)) {
    out_results->value_contents.data_length =
        iree_hal_buffer_byte_length(value_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->value_contents.data_length,
                                   (void**)&out_results->value_contents.data);
  }

  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, value_buffer, 0, out_results->value_contents.data,
        out_results->value_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }
  if (iree_status_is_ok(status)) {
    out_results->actual_contents.data_length =
        iree_hal_buffer_byte_length(result_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->actual_contents.data_length,
                                   (void**)&out_results->actual_contents.data);
  }
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, result_buffer, 0, out_results->actual_contents.data,
        out_results->actual_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }
  if (iree_status_is_ok(status)) {
    out_results->expected_contents.data_length =
        iree_hal_buffer_byte_length(result_buffer);
    status = iree_allocator_malloc(
        host_allocator, out_results->expected_contents.data_length,
        (void**)&out_results->expected_contents.data);
  }
  // Transfer mask data if mask is provided.
  if (iree_status_is_ok(status) && mask_buffer != nullptr) {
    out_results->mask_contents.data_length =
        iree_hal_buffer_byte_length(mask_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->mask_contents.data_length,
                                   (void**)&out_results->mask_contents.data);
    if (iree_status_is_ok(status)) {
      status = iree_hal_device_transfer_d2h(
          device, mask_buffer, 0, out_results->mask_contents.data,
          out_results->mask_contents.data_length,
          IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
    }
  }
  if (!iree_status_is_ok(status)) {
    attention_results_deinitialize(out_results);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static void attention_results_deinitialize(attention_results_t* results) {
  IREE_TRACE_ZONE_BEGIN(z0);
  iree_allocator_free(results->host_allocator, results->query_contents.data);
  iree_allocator_free(results->host_allocator, results->key_contents.data);
  iree_allocator_free(results->host_allocator, results->value_contents.data);
  iree_allocator_free(results->host_allocator, results->actual_contents.data);
  iree_allocator_free(results->host_allocator, results->expected_contents.data);
  iree_allocator_free(results->host_allocator, results->mask_contents.data);

  IREE_TRACE_ZONE_END(z0);
}

// Helper for check_attention_results: the actual interesting part once we've
// obtained and validated the {b,m,k1,k2,n}_size values. On error, detailed
// logging is written to |file| if it is not NULL.
static iree_status_t check_attention_results_impl(
    FILE* file, const attention_results_t* results, int check_every) {
  IREE_TRACE_ZONE_BEGIN(z0);

  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, reference_attention(
              results->b, results->m, results->k1, results->k2, results->n,
              results->scale, results->query_elem_type, results->key_elem_type,
              results->value_elem_type, results->query_contents,
              results->key_contents, results->value_contents,
              results->expected_contents, results->mask_contents, check_every));

  iree_host_size_t count = 0;
  for (iree_hal_dim_t b = 0; b < results->b; ++b) {
    if (++count < check_every) continue;
    count = 0;
    for (iree_hal_dim_t m = 0; m < results->m; ++m) {
      for (iree_hal_dim_t n = 0; n < results->n; ++n) {
        iree_hal_dim_t idx = (b * results->m + m) * results->n + n;
        iree_test_utils_e2e_value_t actual_value =
            iree_test_utils_read_buffer_element(idx, results->result_elem_type,
                                                results->actual_contents.data);
        iree_test_utils_e2e_value_t expected_value =
            iree_test_utils_read_buffer_element(
                idx, results->result_elem_type,
                results->expected_contents.data);
        if (!iree_test_utils_result_elements_agree(expected_value,
                                                   actual_value)) {
          if (file) {
            char actual_str[32];
            char expected_str[32];
            iree_test_utils_snprintf_value(actual_str, sizeof(actual_str),
                                           actual_value);
            iree_test_utils_snprintf_value(expected_str, sizeof(expected_str),
                                           expected_value);
            fprintf(file,
                    "\n\nerror: the actual and expected attention results "
                    "disagree at b %" PRIdim ", m %" PRIdim ", n %" PRIdim
                    ": actual %s, expected %s\n\n",
                    b, m, n, actual_str, expected_str);
          }
          IREE_TRACE_ZONE_END(z0);
          return iree_make_status(IREE_STATUS_ABORTED);
        }
      }
    }
  }

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

// Given an actual attention's inputs and output (all host-local), uses a
// reference attention implementation on the same inputs to check if the output
// is correct. On error, detailed logging is written to |file| if it is not
// NULL.
static iree_status_t check_attention_results(
    FILE* file, const attention_results_t* results) {
  IREE_TRACE_ZONE_BEGIN(z0);
  // TODO: Increase the check every param to reduce the number of comparisons.
  int check_every = 1;
  iree_status_t status =
      check_attention_results_impl(file, results, check_every);
  if (!iree_status_is_ok(status) && check_every > 1) {
    // If we got a failure with check_every>1, that didn't log a useful
    // numerical summary, as most of the reference matrix entries hadn't been
    // computed. Rerun now with check_every=1 to get that numerical logging.
    iree_status_ignore(status);
    status = check_attention_results_impl(file, results, 1);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// `attention_test` custom module
//===----------------------------------------------------------------------===//
// This uses the C++ wrapper to keep things simple. Though easier to use it's
// got additional overhead/code-size bloat that doesn't matter in a test like
// this. Making a C module builder API that removes the boilerplate there is TBD
// so this file is written in C besides this module so that we can swap it back
// to being pure C in the future.

namespace iree {

class AttentionTestModuleState final {
 public:
  explicit AttentionTestModuleState(iree_allocator_t host_allocator)
      : host_allocator_(host_allocator) {}
  ~AttentionTestModuleState() = default;

  // Fills the destination span with pseudorandom values of the given
  // |element_type|. The given |seed| is passed to the pseudorandom generator.
  // The pseudorandom values are reproducible both across runs and across
  // machines.
  StatusOr<vm::ref<iree_hal_buffer_view_t>> GenerateRandom3dTensor(
      iree_hal_device_t* device, int64_t dim0, int64_t dim1, int64_t dim2,
      iree_hal_element_type_t element_type, int32_t seed) {
    iree_hal_dim_t dims[3] = {
        (iree_hal_dim_t)dim0,
        (iree_hal_dim_t)dim1,
        (iree_hal_dim_t)dim2,
    };
    iree_hal_buffer_params_t buffer_params = {0};
    buffer_params.usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
    buffer_params.access = IREE_HAL_MEMORY_ACCESS_ALL;
    buffer_params.type = IREE_HAL_MEMORY_TYPE_OPTIMAL_FOR_DEVICE;
    vm::ref<iree_hal_buffer_view_t> result_view;
    struct callback_state_t {
      iree_hal_element_type_t element_type;
      int32_t seed;
    } callback_state = {
        element_type,
        seed,
    };
    IREE_RETURN_IF_ERROR(iree_hal_buffer_view_generate_buffer(
        device, iree_hal_device_allocator(device), IREE_ARRAYSIZE(dims), dims,
        element_type, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, buffer_params,
        +[](iree_hal_buffer_mapping_t* mapping, void* user_data) {
          callback_state_t callback_state = *(callback_state_t*)user_data;
          iree_byte_span_t span = mapping->contents;
          // Generate "uniform" integer-valued numbers in the range [min, max].
          int32_t min = 0;
          int32_t max = 0;
          iree_test_utils_get_min_max_for_element_type(
              callback_state.element_type, &min, &max);
          uint32_t range = (max - min + 1);
          iree_host_size_t element_byte_count =
              iree_hal_element_dense_byte_count(callback_state.element_type);
          uint8_t* data_end = span.data + span.data_length;
          uint32_t state = callback_state.seed;
          for (uint8_t* data = span.data; data < data_end;
               data += element_byte_count) {
            int32_t value =
                (int32_t)iree_test_utils_pseudorandom_range(&state, range) +
                min;
            iree_test_utils_write_element(callback_state.element_type, value,
                                          data);
          }
          return iree_ok_status();
        },
        &callback_state, &result_view));
    return std::move(result_view);
  }

  StatusOr<vm::ref<iree_hal_buffer_view_t>> GenerateRandom4dTensor(
      iree_hal_device_t* device, int64_t dim0, int64_t dim1, int64_t dim2,
      int64_t dim3, iree_hal_element_type_t element_type, int32_t seed) {
    iree_hal_dim_t dims[4] = {
        (iree_hal_dim_t)dim0,
        (iree_hal_dim_t)dim1,
        (iree_hal_dim_t)dim2,
        (iree_hal_dim_t)dim3,
    };
    iree_hal_buffer_params_t buffer_params = {0};
    buffer_params.usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
    buffer_params.access = IREE_HAL_MEMORY_ACCESS_ALL;
    buffer_params.type = IREE_HAL_MEMORY_TYPE_OPTIMAL_FOR_DEVICE;
    vm::ref<iree_hal_buffer_view_t> result_view;
    struct callback_state_t {
      iree_hal_element_type_t element_type;
      int32_t seed;
    } callback_state = {element_type, seed};
    IREE_RETURN_IF_ERROR(iree_hal_buffer_view_generate_buffer(
        device, iree_hal_device_allocator(device), IREE_ARRAYSIZE(dims), dims,
        element_type, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, buffer_params,
        +[](iree_hal_buffer_mapping_t* mapping, void* user_data) {
          callback_state_t callback_state = *(callback_state_t*)user_data;
          iree_byte_span_t span = mapping->contents;
          int32_t min = 0;
          int32_t max = 0;
          iree_test_utils_get_min_max_for_element_type(
              callback_state.element_type, &min, &max);
          uint32_t range = (max - min + 1);
          iree_host_size_t element_byte_count =
              iree_hal_element_dense_byte_count(callback_state.element_type);
          uint8_t* data_end = span.data + span.data_length;
          uint32_t state = callback_state.seed;
          for (uint8_t* data = span.data; data < data_end;
               data += element_byte_count) {
            int32_t value =
                (int32_t)iree_test_utils_pseudorandom_range(&state, range) +
                min;
            iree_test_utils_write_element(callback_state.element_type, value,
                                          data);
          }
          return iree_ok_status();
        },
        &callback_state, &result_view));
    return std::move(result_view);
  }

  // Gathers the paged KV cache into dense key/value tensors and then defers to
  // the shared attention reference and result checker.
  Status CheckPagedAttentionResults(iree_hal_device_t* device, int64_t batch,
                                    int64_t num_pages, int64_t page_size,
                                    int64_t head_dim, float scale,
                                    iree_hal_buffer_view_t* query,
                                    iree_hal_buffer_view_t* kv_storage,
                                    iree_hal_buffer_view_t* key_page_table,
                                    iree_hal_buffer_view_t* value_page_table,
                                    iree_hal_buffer_view_t* actual_result) {
    iree_hal_buffer_t* query_buffer = iree_hal_buffer_view_buffer(query);
    iree_hal_buffer_t* kv_buffer = iree_hal_buffer_view_buffer(kv_storage);
    iree_hal_buffer_t* key_table_buffer =
        iree_hal_buffer_view_buffer(key_page_table);
    iree_hal_buffer_t* value_table_buffer =
        iree_hal_buffer_view_buffer(value_page_table);
    iree_hal_buffer_t* result_buffer =
        iree_hal_buffer_view_buffer(actual_result);

    const iree_host_size_t query_bytes =
        iree_hal_buffer_byte_length(query_buffer);
    const iree_host_size_t kv_bytes = iree_hal_buffer_byte_length(kv_buffer);
    const iree_host_size_t key_table_bytes =
        iree_hal_buffer_byte_length(key_table_buffer);
    const iree_host_size_t value_table_bytes =
        iree_hal_buffer_byte_length(value_table_buffer);
    const iree_host_size_t result_bytes =
        iree_hal_buffer_byte_length(result_buffer);

    std::vector<uint8_t> query_data(query_bytes);
    std::vector<uint8_t> kv_data(kv_bytes);
    std::vector<uint8_t> key_table_data(key_table_bytes);
    std::vector<uint8_t> value_table_data(value_table_bytes);
    std::vector<uint8_t> result_data(result_bytes);

    IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h(
        device, query_buffer, 0, query_data.data(), query_data.size(),
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()));
    IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h(
        device, kv_buffer, 0, kv_data.data(), kv_data.size(),
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()));
    IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h(
        device, key_table_buffer, 0, key_table_data.data(),
        key_table_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
        iree_infinite_timeout()));
    IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h(
        device, value_table_buffer, 0, value_table_data.data(),
        value_table_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
        iree_infinite_timeout()));
    IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h(
        device, result_buffer, 0, result_data.data(), result_data.size(),
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()));

    if (iree_hal_buffer_view_element_type(query) !=
            IREE_HAL_ELEMENT_TYPE_FLOAT_16 ||
        iree_hal_buffer_view_element_type(kv_storage) !=
            IREE_HAL_ELEMENT_TYPE_FLOAT_16 ||
        iree_hal_buffer_view_element_type(key_page_table) !=
            IREE_HAL_ELEMENT_TYPE_INT_64 ||
        iree_hal_buffer_view_element_type(value_page_table) !=
            IREE_HAL_ELEMENT_TYPE_INT_64 ||
        iree_hal_buffer_view_element_type(actual_result) !=
            IREE_HAL_ELEMENT_TYPE_FLOAT_32) {
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "unexpected paged attention element types");
    }

    const int64_t pool_pages =
        kv_bytes / (sizeof(uint16_t) * page_size * head_dim);
    const uint16_t* query_f16 =
        reinterpret_cast<const uint16_t*>(query_data.data());
    const uint16_t* kv_f16 = reinterpret_cast<const uint16_t*>(kv_data.data());
    const int64_t* key_pages =
        reinterpret_cast<const int64_t*>(key_table_data.data());
    const int64_t* value_pages =
        reinterpret_cast<const int64_t*>(value_table_data.data());

    if (query_bytes != batch * head_dim * sizeof(uint16_t) ||
        key_table_bytes != batch * num_pages * sizeof(int64_t) ||
        value_table_bytes != batch * num_pages * sizeof(int64_t) ||
        result_bytes != batch * head_dim * sizeof(float) || pool_pages <= 0) {
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "unexpected paged attention buffer sizes");
    }

    const int64_t sequence_length = num_pages * page_size;
    std::vector<float> query_f32(batch * head_dim);
    std::vector<float> key_f32(batch * sequence_length * head_dim);
    std::vector<float> value_f32(batch * sequence_length * head_dim);
    std::vector<float> expected(batch * head_dim);
    for (int64_t b = 0; b < batch; ++b) {
      for (int64_t d = 0; d < head_dim; ++d) {
        query_f32[b * head_dim + d] =
            iree_math_f16_to_f32(query_f16[b * head_dim + d]);
      }
      for (int64_t position = 0; position < sequence_length; ++position) {
        const int64_t logical_page = position / page_size;
        const int64_t page_offset = position % page_size;
        const int64_t key_page = key_pages[b * num_pages + logical_page];
        const int64_t value_page = value_pages[b * num_pages + logical_page];
        if (key_page < 0 || key_page >= pool_pages || value_page < 0 ||
            value_page >= pool_pages) {
          return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                  "paged attention page is out of bounds");
        }
        for (int64_t d = 0; d < head_dim; ++d) {
          const int64_t dense_index =
              (b * sequence_length + position) * head_dim + d;
          key_f32[dense_index] = iree_math_f16_to_f32(
              kv_f16[(key_page * page_size + page_offset) * head_dim + d]);
          value_f32[dense_index] = iree_math_f16_to_f32(
              kv_f16[(value_page * page_size + page_offset) * head_dim + d]);
        }
      }
    }

    // All spans below point at buffers owned by this function, so this struct
    // must not be passed to attention_results_deinitialize.
    attention_results_t results = {};
    results.b = batch;
    results.m = 1;
    results.k1 = head_dim;
    results.k2 = sequence_length;
    results.n = head_dim;
    results.scale = scale;
    results.query_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32;
    results.key_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32;
    results.value_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32;
    results.result_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32;
    results.query_contents =
        iree_make_byte_span(query_f32.data(), query_f32.size() * sizeof(float));
    results.key_contents =
        iree_make_byte_span(key_f32.data(), key_f32.size() * sizeof(float));
    results.value_contents =
        iree_make_byte_span(value_f32.data(), value_f32.size() * sizeof(float));
    results.actual_contents =
        iree_make_byte_span(result_data.data(), result_data.size());
    results.expected_contents =
        iree_make_byte_span(expected.data(), expected.size() * sizeof(float));
    return check_attention_results(stderr, &results);
  }

  Status CheckAttentionResults(iree_hal_device_t* device, int64_t b, int64_t m,
                               int64_t k1, int64_t k2, int64_t n, float scale,
                               iree_hal_buffer_view_t* query,
                               iree_hal_buffer_view_t* key,
                               iree_hal_buffer_view_t* value,
                               iree_hal_buffer_view_t* actual_result) {
    attention_results_t results = {};
    IREE_RETURN_IF_ERROR(attention_results_initialize(
        device, (iree_hal_dim_t)b, (iree_hal_dim_t)m, (iree_hal_dim_t)k1,
        (iree_hal_dim_t)k2, (iree_hal_dim_t)n, scale, query, key, value,
        actual_result,
        nullptr,  // No mask
        host_allocator_, &results));
    iree_status_t status = check_attention_results(stderr, &results);
    attention_results_deinitialize(&results);
    return status;
  }

  Status CheckAttentionResultsWithMask(
      iree_hal_device_t* device, int64_t b, int64_t m, int64_t k1, int64_t k2,
      int64_t n, float scale, iree_hal_buffer_view_t* query,
      iree_hal_buffer_view_t* key, iree_hal_buffer_view_t* value,
      iree_hal_buffer_view_t* mask, iree_hal_buffer_view_t* actual_result) {
    attention_results_t results = {};
    IREE_RETURN_IF_ERROR(attention_results_initialize(
        device, (iree_hal_dim_t)b, (iree_hal_dim_t)m, (iree_hal_dim_t)k1,
        (iree_hal_dim_t)k2, (iree_hal_dim_t)n, scale, query, key, value,
        actual_result, mask, host_allocator_, &results));
    iree_status_t status = check_attention_results(stderr, &results);
    attention_results_deinitialize(&results);
    return status;
  }

 private:
  iree_allocator_t host_allocator_;
};

static const vm::NativeFunction<AttentionTestModuleState>
    kAttentionTestModuleFunctions[] = {
        vm::MakeNativeFunction(
            "generate_random_tensor",
            &AttentionTestModuleState::GenerateRandom3dTensor),
        vm::MakeNativeFunction(
            "generate_random_4d_tensor",
            &AttentionTestModuleState::GenerateRandom4dTensor),
        vm::MakeNativeFunction(
            "check_paged_attention_results",
            &AttentionTestModuleState::CheckPagedAttentionResults),
        vm::MakeNativeFunction(
            "check_attention_results",
            &AttentionTestModuleState::CheckAttentionResults),
        vm::MakeNativeFunction(
            "check_attention_results_with_mask",
            &AttentionTestModuleState::CheckAttentionResultsWithMask),
};

struct AttentionTestModule final
    : public vm::NativeModule<AttentionTestModuleState> {
  using vm::NativeModule<AttentionTestModuleState>::NativeModule;
  StatusOr<std::unique_ptr<AttentionTestModuleState>> CreateState(
      iree_allocator_t host_allocator) override {
    return std::make_unique<AttentionTestModuleState>(host_allocator);
  }
  StatusOr<std::unique_ptr<AttentionTestModuleState>> ForkState(
      AttentionTestModuleState* parent_state,
      iree_allocator_t host_allocator) override {
    return std::make_unique<AttentionTestModuleState>(host_allocator);
  }
};

}  // namespace iree

static iree_status_t attention_test_module_create(
    iree_vm_instance_t* instance, iree_allocator_t host_allocator,
    iree_vm_module_t** out_module) {
  IREE_ASSERT_ARGUMENT(out_module);
  *out_module = NULL;
  auto module = std::make_unique<iree::AttentionTestModule>(
      "attention_test", /*version=*/0, instance, host_allocator,
      iree::span<
          const iree::vm::NativeFunction<iree::AttentionTestModuleState>>(
          iree::kAttentionTestModuleFunctions));
  *out_module = module.release()->interface();
  return iree_ok_status();
}

int main(int argc, char** argv) {
  IREE_TRACE_APP_ENTER();

  iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
  if (argc != 1) {
    fprintf(stderr, "use --module= flags to specify the modules to run\n");
    IREE_TRACE_APP_EXIT(EXIT_FAILURE);
    return EXIT_FAILURE;
  }

  iree_status_t status = iree_test_utils_load_and_run_e2e_tests(
      iree_allocator_system(), attention_test_module_create);
  int exit_code = EXIT_SUCCESS;
  if (!iree_status_is_ok(status)) {
    iree_status_fprint(stderr, status);
    bool is_unavailable = iree_status_is_unavailable(status);
    iree_status_free(status);
    exit_code = is_unavailable ? EXIT_SUCCESS : EXIT_FAILURE;
  }

  IREE_TRACE_APP_EXIT(exit_code);
  return exit_code;
}