| // 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; | |
| } |