| // 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 <stdio.h> | |
| #include <stdlib.h> | |
| #include <string.h> | |
| #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 matmul | |
| //===----------------------------------------------------------------------===// | |
| #define REFERENCE_MATMUL(LHSTYPE, RHSTYPE, RESTYPE, ACCTYPE) \ | |
| static void reference_matmul_##LHSTYPE##_##RHSTYPE##_##RESTYPE##_##ACCTYPE( \ | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, \ | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, \ | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, \ | |
| const LHSTYPE* lhs_data, const RHSTYPE* rhs_data, \ | |
| const ACCTYPE* acc_data, RESTYPE* result_data, iree_hal_dim_t m, \ | |
| iree_hal_dim_t n) { \ | |
| ACCTYPE acc = acc_data ? acc_data[n + m * n_size] : 0; \ | |
| for (iree_hal_dim_t k = 0; k < k_size; ++k) { \ | |
| LHSTYPE lhs_value = lhs_data[k + m * k_size]; \ | |
| RHSTYPE rhs_value = \ | |
| transpose_rhs ? rhs_data[k + n * k_size] : rhs_data[n + k * n_size]; \ | |
| acc += (ACCTYPE)lhs_value * (ACCTYPE)rhs_value; \ | |
| } \ | |
| result_data[n + m * n_size] = acc; \ | |
| } | |
| // Reference matmul instantiations | |
| REFERENCE_MATMUL(float, float, float, float) | |
| REFERENCE_MATMUL(double, double, double, double) | |
| REFERENCE_MATMUL(int8_t, int8_t, int32_t, int32_t) | |
| REFERENCE_MATMUL(uint8_t, int8_t, int32_t, int32_t) | |
| REFERENCE_MATMUL(int32_t, int32_t, int32_t, int32_t) | |
| // Reference matmul for the f16 input, f16 accumulation, and f16 result. | |
| // [f16 <= f16 * f16 + f16] | |
| static void reference_matmul_f16_f16_f16_f16( | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, | |
| const uint16_t* lhs_data, const uint16_t* rhs_data, | |
| const uint16_t* acc_data, uint16_t* result_data, iree_hal_dim_t m, | |
| iree_hal_dim_t n) { | |
| float acc = acc_data ? iree_math_f16_to_f32(acc_data[n + m * n_size]) : 0.f; | |
| for (iree_hal_dim_t k = 0; k < k_size; ++k) { | |
| int64_t rhs_index = transpose_rhs ? k + n * k_size : n + k * n_size; | |
| acc += iree_math_f16_to_f32(lhs_data[k + m * k_size]) * | |
| iree_math_f16_to_f32(rhs_data[rhs_index]); | |
| } | |
| result_data[n + m * n_size] = iree_math_f32_to_f16(acc); | |
| } | |
| // Reference matmul for the f16 input, f32 accumulation, and f32 result. | |
| // [f32 <= f16 * f16 + f32] | |
| static void reference_matmul_f16_f16_f32_f32( | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, | |
| const uint16_t* lhs_data, const uint16_t* rhs_data, const float* acc_data, | |
| float* result_data, iree_hal_dim_t m, iree_hal_dim_t n) { | |
| float acc = acc_data ? acc_data[n + m * n_size] : 0.f; | |
| for (iree_hal_dim_t k = 0; k < k_size; ++k) { | |
| int64_t rhs_index = transpose_rhs ? k + n * k_size : n + k * n_size; | |
| acc += iree_math_f16_to_f32(lhs_data[k + m * k_size]) * | |
| iree_math_f16_to_f32(rhs_data[rhs_index]); | |
| } | |
| result_data[n + m * n_size] = acc; | |
| } | |
| // Reference matmul for the bf16 input, bf16 accumulation, and bf16 result. | |
| // [bf16 <= bf16 * bf16 + bf16] | |
| static void reference_matmul_bf16_bf16_bf16_bf16( | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, | |
| const uint16_t* lhs_data, const uint16_t* rhs_data, | |
| const uint16_t* acc_data, uint16_t* result_data, iree_hal_dim_t m, | |
| iree_hal_dim_t n) { | |
| float acc = acc_data ? iree_math_bf16_to_f32(acc_data[n + m * n_size]) : 0.f; | |
| for (iree_hal_dim_t k = 0; k < k_size; ++k) { | |
| int64_t rhs_index = transpose_rhs ? k + n * k_size : n + k * n_size; | |
| acc += iree_math_bf16_to_f32(lhs_data[k + m * k_size]) * | |
| iree_math_bf16_to_f32(rhs_data[rhs_index]); | |
| } | |
| result_data[n + m * n_size] = iree_math_f32_to_bf16(acc); | |
| } | |
| // Reference matmul for the bf16 input, f32 accumulation, and f32 result. | |
| // [f32 <= bf16 * bf16 + f32] | |
| static void reference_matmul_bf16_bf16_f32_f32( | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, | |
| const uint16_t* lhs_data, const uint16_t* rhs_data, const float* acc_data, | |
| float* result_data, iree_hal_dim_t m, iree_hal_dim_t n) { | |
| float acc = acc_data ? acc_data[n + m * n_size] : 0.f; | |
| for (iree_hal_dim_t k = 0; k < k_size; ++k) { | |
| int64_t rhs_index = transpose_rhs ? k + n * k_size : n + k * n_size; | |
| acc += iree_math_bf16_to_f32(lhs_data[k + m * k_size]) * | |
| iree_math_bf16_to_f32(rhs_data[rhs_index]); | |
| } | |
| result_data[n + m * n_size] = acc; | |
| } | |
| #define REFERENCE_MATMUL_F8(LHSTYPE, RHSTYPE) \ | |
| static void reference_matmul_##LHSTYPE##_##RHSTYPE##_f32_f32( \ | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, \ | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, \ | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, \ | |
| const uint8_t* lhs_data, const uint8_t* rhs_data, const float* acc_data, \ | |
| float* result_data, iree_hal_dim_t m, iree_hal_dim_t n) { \ | |
| float acc = acc_data ? acc_data[n + m * n_size] : 0; \ | |
| for (iree_hal_dim_t k = 0; k < k_size; ++k) { \ | |
| float lhs_float = \ | |
| iree_math_##LHSTYPE##_to_f32(lhs_data[k + m * k_size]); \ | |
| float rhs_float = iree_math_##RHSTYPE##_to_f32( \ | |
| rhs_data[transpose_rhs ? k + n * k_size : n + k * n_size]); \ | |
| acc += lhs_float * rhs_float; \ | |
| } \ | |
| result_data[n + m * n_size] = acc; \ | |
| } | |
| REFERENCE_MATMUL_F8(f8e5m2, f8e5m2) | |
| REFERENCE_MATMUL_F8(f8e4m3fn, f8e4m3fn) | |
| REFERENCE_MATMUL_F8(f8e5m2fnuz, f8e5m2fnuz) | |
| REFERENCE_MATMUL_F8(f8e4m3fnuz, f8e4m3fnuz) | |
| // Helper for reference_matmul. | |
| // Computes one element in the result matrix. | |
| static iree_status_t reference_matmul_element( | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, | |
| iree_hal_element_type_t lhs_type, iree_hal_element_type_t rhs_type, | |
| iree_hal_element_type_t acc_type, bool transpose_rhs, void* lhs_data, | |
| void* rhs_data, void* acc_data, void* result_data, iree_hal_dim_t m, | |
| iree_hal_dim_t n) { | |
| if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_float_float_float_float( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const float*)lhs_data, (const float*)rhs_data, (const float*)acc_data, | |
| (float*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_64 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_64 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_64) { | |
| reference_matmul_double_double_double_double( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const double*)lhs_data, (const double*)rhs_data, | |
| (const double*)acc_data, (double*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_UINT_8 && | |
| iree_hal_element_type_is_integer(rhs_type, 8) && | |
| iree_hal_element_type_is_integer(acc_type, 32)) { | |
| // Mixed unsigned-LHS / signed-RHS i8 matmul (e.g. x86 vpdpbusd). IR | |
| // storage is signless i8; the buffer-view's UINT_8 LHS carries the | |
| // signedness. | |
| reference_matmul_uint8_t_int8_t_int32_t_int32_t( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint8_t*)lhs_data, (const int8_t*)rhs_data, | |
| (const int32_t*)acc_data, (int32_t*)result_data, m, n); | |
| } else if (iree_hal_element_type_is_integer(lhs_type, 8) && | |
| iree_hal_element_type_is_integer(rhs_type, 8) && | |
| iree_hal_element_type_is_integer(acc_type, 32)) { | |
| reference_matmul_int8_t_int8_t_int32_t_int32_t( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const int8_t*)lhs_data, (const int8_t*)rhs_data, | |
| (const int32_t*)acc_data, (int32_t*)result_data, m, n); | |
| } else if (iree_hal_element_type_is_integer(lhs_type, 32) && | |
| iree_hal_element_type_is_integer(rhs_type, 32) && | |
| iree_hal_element_type_is_integer(acc_type, 32)) { | |
| reference_matmul_int32_t_int32_t_int32_t_int32_t( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const int32_t*)lhs_data, (const int32_t*)rhs_data, | |
| (const int32_t*)acc_data, (int32_t*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16) { | |
| reference_matmul_f16_f16_f16_f16( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint16_t*)lhs_data, (const uint16_t*)rhs_data, | |
| (const uint16_t*)acc_data, (uint16_t*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_f16_f16_f32_f32( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint16_t*)lhs_data, (const uint16_t*)rhs_data, | |
| (const float*)acc_data, (float*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_BFLOAT_16 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_BFLOAT_16 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_BFLOAT_16) { | |
| reference_matmul_bf16_bf16_bf16_bf16( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint16_t*)lhs_data, (const uint16_t*)rhs_data, | |
| (const uint16_t*)acc_data, (uint16_t*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_BFLOAT_16 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_BFLOAT_16 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_bf16_bf16_f32_f32( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint16_t*)lhs_data, (const uint16_t*)rhs_data, | |
| (const float*)acc_data, (float*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E5M2 && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E5M2 && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_f8e5m2_f8e5m2_f32_f32( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint8_t*)lhs_data, (const uint8_t*)rhs_data, | |
| (const float*)acc_data, (float*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E4M3_FN && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E4M3_FN && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_f8e4m3fn_f8e4m3fn_f32_f32( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint8_t*)lhs_data, (const uint8_t*)rhs_data, | |
| (const float*)acc_data, (float*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E5M2_FNUZ && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E5M2_FNUZ && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_f8e5m2fnuz_f8e5m2fnuz_f32_f32( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint8_t*)lhs_data, (const uint8_t*)rhs_data, | |
| (const float*)acc_data, (float*)result_data, m, n); | |
| } else if (lhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E4M3_FNUZ && | |
| rhs_type == IREE_HAL_ELEMENT_TYPE_FLOAT_8_E4M3_FNUZ && | |
| acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { | |
| reference_matmul_f8e4m3fnuz_f8e4m3fnuz_f32_f32( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, transpose_rhs, | |
| (const uint8_t*)lhs_data, (const uint8_t*)rhs_data, | |
| (const float*)acc_data, (float*)result_data, m, n); | |
| } else { | |
| return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "unhandled combination of element types in matmul"); | |
| } | |
| return iree_ok_status(); | |
| } | |
| // Reference matmul implementation, used to compare matmul results against. | |
| static iree_status_t reference_matmul( | |
| iree_hal_dim_t m_size, iree_hal_dim_t k_size, iree_hal_dim_t n_size, | |
| iree_hal_dim_t m_start, iree_hal_dim_t m_end, iree_hal_dim_t n_start, | |
| iree_hal_dim_t n_end, iree_hal_element_type_t lhs_type, | |
| iree_hal_element_type_t rhs_type, iree_hal_element_type_t acc_type, | |
| bool transpose_rhs, iree_byte_span_t lhs_contents, | |
| iree_byte_span_t rhs_contents, iree_byte_span_t acc_contents, | |
| iree_byte_span_t result_contents, int compute_every) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, m_size); | |
| IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, k_size); | |
| IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, n_size); | |
| iree_host_size_t count = 0; | |
| for (iree_hal_dim_t m = m_start; m < m_end; ++m) { | |
| for (iree_hal_dim_t n = n_start; n < n_end; ++n) { | |
| if (++count < compute_every) continue; | |
| count = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, reference_matmul_element( | |
| m_size, k_size, n_size, lhs_type, rhs_type, acc_type, | |
| transpose_rhs, lhs_contents.data, rhs_contents.data, | |
| acc_contents.data, result_contents.data, m, n)); | |
| } | |
| } | |
| IREE_TRACE_ZONE_END(z0); | |
| return iree_ok_status(); | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Matmul comparison/logging | |
| //===----------------------------------------------------------------------===// | |
| typedef struct { | |
| iree_allocator_t host_allocator; | |
| iree_hal_dim_t m; | |
| iree_hal_dim_t k; | |
| iree_hal_dim_t n; | |
| iree_hal_element_type_t lhs_type; | |
| iree_hal_element_type_t rhs_type; | |
| iree_hal_element_type_t acc_type; | |
| iree_hal_element_type_t result_type; | |
| bool transpose_rhs; | |
| iree_byte_span_t lhs_contents; | |
| iree_byte_span_t rhs_contents; | |
| iree_byte_span_t acc_contents; | |
| iree_byte_span_t actual_contents; | |
| iree_byte_span_t expected_contents; | |
| } matmul_results_t; | |
| static void matmul_results_deinitialize(matmul_results_t* results); | |
| static iree_status_t matmul_results_initialize( | |
| iree_hal_device_t* device, iree_hal_dim_t m_size, iree_hal_dim_t k_size, | |
| iree_hal_dim_t n_size, uint32_t transpose_rhs, iree_hal_buffer_view_t* lhs, | |
| iree_hal_buffer_view_t* rhs, iree_hal_buffer_view_t* acc, | |
| iree_hal_buffer_view_t* result, iree_allocator_t host_allocator, | |
| matmul_results_t* out_results) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| memset(out_results, 0, sizeof(*out_results)); | |
| out_results->host_allocator = host_allocator; | |
| out_results->m = m_size; | |
| out_results->k = k_size; | |
| out_results->n = n_size; | |
| out_results->lhs_type = iree_hal_buffer_view_element_type(lhs); | |
| out_results->rhs_type = iree_hal_buffer_view_element_type(rhs); | |
| out_results->acc_type = iree_hal_buffer_view_element_type(result); | |
| out_results->result_type = iree_hal_buffer_view_element_type(result); | |
| out_results->transpose_rhs = transpose_rhs != 0; | |
| iree_hal_buffer_t* lhs_buffer = iree_hal_buffer_view_buffer(lhs); | |
| iree_hal_buffer_t* rhs_buffer = iree_hal_buffer_view_buffer(rhs); | |
| iree_hal_buffer_t* acc_buffer = acc ? iree_hal_buffer_view_buffer(acc) : NULL; | |
| iree_hal_buffer_t* result_buffer = iree_hal_buffer_view_buffer(result); | |
| iree_status_t status = iree_ok_status(); | |
| if (iree_status_is_ok(status)) { | |
| out_results->lhs_contents.data_length = | |
| iree_hal_buffer_byte_length(lhs_buffer); | |
| status = iree_allocator_malloc(host_allocator, | |
| out_results->lhs_contents.data_length, | |
| (void**)&out_results->lhs_contents.data); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, lhs_buffer, 0, out_results->lhs_contents.data, | |
| out_results->lhs_contents.data_length, | |
| IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| out_results->rhs_contents.data_length = | |
| iree_hal_buffer_byte_length(rhs_buffer); | |
| status = iree_allocator_malloc(host_allocator, | |
| out_results->rhs_contents.data_length, | |
| (void**)&out_results->rhs_contents.data); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, rhs_buffer, 0, out_results->rhs_contents.data, | |
| out_results->rhs_contents.data_length, | |
| IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); | |
| } | |
| if (acc_buffer) { | |
| if (iree_status_is_ok(status)) { | |
| out_results->acc_contents.data_length = | |
| iree_hal_buffer_byte_length(acc_buffer); | |
| status = iree_allocator_malloc(host_allocator, | |
| out_results->acc_contents.data_length, | |
| (void**)&out_results->acc_contents.data); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, acc_buffer, 0, out_results->acc_contents.data, | |
| out_results->acc_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); | |
| } | |
| if (!iree_status_is_ok(status)) { | |
| matmul_results_deinitialize(out_results); | |
| } | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| static void matmul_results_deinitialize(matmul_results_t* results) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| iree_allocator_free(results->host_allocator, results->lhs_contents.data); | |
| iree_allocator_free(results->host_allocator, results->rhs_contents.data); | |
| if (!iree_byte_span_is_empty(results->acc_contents)) { | |
| iree_allocator_free(results->host_allocator, results->acc_contents.data); | |
| } | |
| iree_allocator_free(results->host_allocator, results->actual_contents.data); | |
| iree_allocator_free(results->host_allocator, results->expected_contents.data); | |
| IREE_TRACE_ZONE_END(z0); | |
| } | |
| // Returns the largest number of characters to print any matrix element. | |
| static int get_max_elem_width(iree_hal_dim_t rows, iree_hal_dim_t row_start, | |
| iree_hal_dim_t row_end, iree_hal_dim_t cols, | |
| iree_hal_dim_t col_start, iree_hal_dim_t col_end, | |
| iree_hal_element_type_t element_type, | |
| const uint8_t* matrix) { | |
| int max_elem_width = 0; | |
| for (int row = row_start; row < row_end; row++) { | |
| for (int col = col_start; col < col_end; col++) { | |
| iree_hal_dim_t idx = col + row * cols; | |
| iree_test_utils_e2e_value_t elem = | |
| iree_test_utils_read_buffer_element(idx, element_type, matrix); | |
| // NOTE: iree_max is a macro and may evaluate its args twice. | |
| char buf[64]; | |
| int this_elem_width = | |
| iree_test_utils_snprintf_value(buf, sizeof(buf), elem); | |
| max_elem_width = iree_max(max_elem_width, this_elem_width); | |
| } | |
| } | |
| return max_elem_width; | |
| } | |
| // Prints |matrix| to |file|, with |label| as caption. | |
| // | |
| // If |other_matrix| is not NULL, then any matrix entries that disagree | |
| // between |matrix| and |other_matrix| (according to | |
| // matmul_result_elements_agree) are highlighted. | |
| // | |
| // |highlight| is either NULL or is a UTF-8 string that will be printed next to | |
| // any entry of |matrix| that disagrees with the corresponding entry of | |
| // |other_matrix|. | |
| // | |
| // |highlight| should be NULL if and only if |other_matrix| is NULL. | |
| // | |
| // In order for matrix columns to be properly laid out, the rendering of | |
| // |highlight| in a fixed-width font should have the width of two regular Latin | |
| // characters. According to | |
| // https://www.unicode.org/reports/tr11/#Recommendations, a single emoji | |
| // character should meet that requirement. | |
| static void print_matrix(FILE* file, const char* label, iree_hal_dim_t rows, | |
| iree_hal_dim_t row_start, iree_hal_dim_t row_end, | |
| iree_hal_dim_t cols, iree_hal_dim_t col_start, | |
| iree_hal_dim_t col_end, | |
| iree_hal_element_type_t element_type, | |
| const uint8_t* matrix, const uint8_t* other_matrix, | |
| const char* highlight) { | |
| IREE_ASSERT((other_matrix == NULL) == (highlight == NULL)); | |
| int max_elem_width = get_max_elem_width( | |
| rows, row_start, row_end, cols, col_start, col_end, element_type, matrix); | |
| if (other_matrix) { | |
| // NOTE: iree_max is a macro and may evaluate its args twice. | |
| int other_matrix_max_elem_width = | |
| get_max_elem_width(rows, row_start, row_end, cols, col_start, col_end, | |
| element_type, other_matrix); | |
| max_elem_width = iree_max(max_elem_width, other_matrix_max_elem_width); | |
| } | |
| fprintf(file, | |
| "%s (rows %" PRIdsz "..%" PRIdsz " out of 0..%" PRIdsz | |
| ", columns %" PRIdsz "..%" PRIdsz " out of 0..%" PRIdsz ")\n", | |
| label, row_start, row_end - 1, rows - 1, col_start, col_end - 1, | |
| cols - 1); | |
| for (int row = row_start; row < row_end; row++) { | |
| for (int col = col_start; col < col_end; col++) { | |
| iree_hal_dim_t idx = col + row * cols; | |
| iree_test_utils_e2e_value_t element = | |
| iree_test_utils_read_buffer_element(idx, element_type, matrix); | |
| bool disagree = false; | |
| if (other_matrix) { | |
| iree_test_utils_e2e_value_t other_element = | |
| iree_test_utils_read_buffer_element(idx, element_type, | |
| other_matrix); | |
| disagree = | |
| !iree_test_utils_result_elements_agree(element, other_element); | |
| } | |
| char buf[64]; | |
| iree_test_utils_snprintf_value(buf, sizeof(buf), element); | |
| fprintf(file, "%*s", max_elem_width, buf); | |
| // See comment on |highlight| function parameter for why 2 spaces. | |
| // A 3rd space is added unconditionally to make it clear that a highlight | |
| // concerns the matrix entry to its left. | |
| fprintf(file, "%s ", disagree ? highlight : " "); | |
| } | |
| fprintf(file, "\n"); | |
| } | |
| } | |
| // Coordinates of regions of matrices to dump on numerical diagnostics. | |
| typedef struct { | |
| iree_hal_dim_t m_start; | |
| iree_hal_dim_t m_end; | |
| iree_hal_dim_t n_start; | |
| iree_hal_dim_t n_end; | |
| iree_hal_dim_t k_start; | |
| iree_hal_dim_t k_end; | |
| } matrix_debug_window_t; | |
| // Given a position (m, n) in the result matrix given by `results`, populates | |
| // a `window` structure determining the region of the result matrix to be | |
| // printed in diagnostics. | |
| static void get_matrix_diagnostics_window_coords( | |
| iree_hal_dim_t m, iree_hal_dim_t n, const matmul_results_t* results, | |
| matrix_debug_window_t* window) { | |
| auto getIntEnv = [](const char* env_var, int default_val) { | |
| const char* env = getenv(env_var); | |
| return env ? std::atoi(env) : default_val; | |
| }; | |
| iree_hal_dim_t context_m = getIntEnv("IREE_MATMUL_TEST_SHOW_CONTEXT_M", 16); | |
| iree_hal_dim_t context_n = getIntEnv("IREE_MATMUL_TEST_SHOW_CONTEXT_N", 16); | |
| iree_hal_dim_t context_k = getIntEnv("IREE_MATMUL_TEST_SHOW_CONTEXT_K", 16); | |
| window->m_start = | |
| (iree_hal_dim_t)iree_max(0, (int64_t)m - (int64_t)context_m / 2); | |
| window->m_end = iree_min(results->m, window->m_start + context_m); | |
| window->n_start = | |
| (iree_hal_dim_t)iree_max(0, (int64_t)n - (int64_t)context_n / 2); | |
| window->n_end = iree_min(results->n, window->n_start + context_n); | |
| window->k_start = 0; | |
| window->k_end = iree_min(results->k, context_k); | |
| } | |
| // Helper for check_matmul_results: handler for the failure case. | |
| // If |file| is not NULL, detailed logging is written to it. | |
| static iree_status_t check_matmul_failure( | |
| FILE* file, const matmul_results_t* results, | |
| iree_test_utils_e2e_value_t actual_value, | |
| iree_test_utils_e2e_value_t expected_value, iree_hal_dim_t row, | |
| iree_hal_dim_t col, const matrix_debug_window_t* window) { | |
| if (!file) { | |
| // No logging of errors with check_every>1 as most of the reference matrix | |
| // elements have not been computed. The caller is expected to retry with | |
| // check_every=1. | |
| return iree_make_status(IREE_STATUS_ABORTED); | |
| } | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| fprintf(file, | |
| "\n\nerror: the actual and expected result matrices disagree " | |
| "at row %" PRIdim ", column %" PRIdim ".\n\n", | |
| row, col); | |
| char actual_value_buf[32]; | |
| char expected_value_buf[32]; | |
| iree_test_utils_snprintf_value(actual_value_buf, sizeof(actual_value_buf), | |
| actual_value); | |
| iree_test_utils_snprintf_value(expected_value_buf, sizeof(expected_value_buf), | |
| expected_value); | |
| fprintf(file, "actual value: %s\n", actual_value_buf); | |
| fprintf(file, "expected value: %s\n", expected_value_buf); | |
| fprintf(file, "\n"); | |
| print_matrix(file, "left-hand side", results->m, window->m_start, | |
| window->m_end, results->k, window->k_start, window->k_end, | |
| results->lhs_type, results->lhs_contents.data, NULL, NULL); | |
| fprintf(file, "\n"); | |
| print_matrix(file, "right-hand side", results->k, window->k_start, | |
| window->k_end, results->n, window->n_start, window->n_end, | |
| results->rhs_type, results->rhs_contents.data, NULL, NULL); | |
| fprintf(file, "\n"); | |
| if (results->acc_contents.data) { | |
| print_matrix(file, "input accumulator", results->m, window->m_start, | |
| window->m_end, results->n, window->n_start, window->n_end, | |
| results->acc_type, results->acc_contents.data, NULL, NULL); | |
| fprintf(file, "\n"); | |
| } | |
| print_matrix(file, "expected result", results->m, window->m_start, | |
| window->m_end, results->n, window->n_start, window->n_end, | |
| results->result_type, results->expected_contents.data, | |
| results->actual_contents.data, iree_test_utils_emoji(true)); | |
| fprintf(file, "\n"); | |
| print_matrix(file, "actual result", results->m, window->m_start, | |
| window->m_end, results->n, window->n_start, window->n_end, | |
| results->result_type, results->actual_contents.data, | |
| results->expected_contents.data, iree_test_utils_emoji(false)); | |
| fprintf(file, "\n"); | |
| IREE_TRACE_ZONE_END(z0); | |
| return iree_make_status(IREE_STATUS_ABORTED); | |
| } | |
| // Helper for check_matmul_results: the actual interesting part once we've | |
| // obtained and validated the {m,k,n}_size values. On error, detailed logging is | |
| // written to |file| if it is not NULL. | |
| static iree_status_t check_matmul_results_impl(FILE* file, | |
| const matmul_results_t* results, | |
| int check_every) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, reference_matmul(results->m, results->k, results->n, 0, results->m, 0, | |
| results->n, results->lhs_type, results->rhs_type, | |
| results->acc_type, results->transpose_rhs, | |
| results->lhs_contents, results->rhs_contents, | |
| results->acc_contents, results->expected_contents, | |
| check_every)); | |
| int count = 0; | |
| for (iree_hal_dim_t m = 0; m < results->m; ++m) { | |
| for (iree_hal_dim_t n = 0; n < results->n; ++n) { | |
| if (++count < check_every) continue; | |
| count = 0; | |
| iree_hal_dim_t idx = m * results->n + n; | |
| iree_test_utils_e2e_value_t actual_value = | |
| iree_test_utils_read_buffer_element(idx, results->result_type, | |
| results->actual_contents.data); | |
| iree_test_utils_e2e_value_t expected_value = | |
| iree_test_utils_read_buffer_element(idx, results->result_type, | |
| results->expected_contents.data); | |
| if (!iree_test_utils_result_elements_agree(actual_value, | |
| expected_value)) { | |
| // Numerical error detected. Determine a window into the matrices to | |
| // generate a diagnostic for. | |
| matrix_debug_window_t window; | |
| get_matrix_diagnostics_window_coords(m, n, results, &window); | |
| // Compute all remaining matrix elements in that window, as we are going | |
| // to print them. | |
| reference_matmul( | |
| results->m, results->k, results->n, window.m_start, window.m_end, | |
| window.n_start, window.n_end, results->lhs_type, results->rhs_type, | |
| results->acc_type, results->transpose_rhs, results->lhs_contents, | |
| results->rhs_contents, results->acc_contents, | |
| results->expected_contents, /*check_every=*/1); | |
| iree_status_t status = check_matmul_failure( | |
| file, results, actual_value, expected_value, m, n, &window); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| } | |
| } | |
| IREE_TRACE_ZONE_END(z0); | |
| return iree_ok_status(); | |
| } | |
| // Given an actual matmul's inputs and output (all host-local), uses a reference | |
| // matmul 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_matmul_results(FILE* file, | |
| const matmul_results_t* results) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| int check_every = iree_test_utils_calculate_check_every( | |
| results->m * results->n, results->n); | |
| iree_status_t status = check_matmul_results_impl(file, results, check_every); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // `matmul_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 MatmulTestModuleState final { | |
| public: | |
| explicit MatmulTestModuleState(iree_allocator_t host_allocator) | |
| : host_allocator_(host_allocator) {} | |
| ~MatmulTestModuleState() = 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>> GenerateRandomMatrix( | |
| iree_hal_device_t* device, int64_t dim0, int64_t dim1, | |
| iree_hal_element_type_t element_type, int32_t seed) { | |
| iree_hal_dim_t dims[2] = { | |
| (iree_hal_dim_t)dim0, | |
| (iree_hal_dim_t)dim1, | |
| }; | |
| 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); | |
| } | |
| Status CheckMatmulResults(iree_hal_device_t* device, int64_t m, int64_t k, | |
| int64_t n, int32_t transpose_rhs, | |
| iree_hal_buffer_view_t* lhs, | |
| iree_hal_buffer_view_t* rhs, | |
| iree_hal_buffer_view_t* acc, | |
| iree_hal_buffer_view_t* actual_result) { | |
| matmul_results_t results = {}; | |
| IREE_RETURN_IF_ERROR(matmul_results_initialize( | |
| device, (iree_hal_dim_t)m, (iree_hal_dim_t)k, (iree_hal_dim_t)n, | |
| transpose_rhs, lhs, rhs, acc, actual_result, host_allocator_, | |
| &results)); | |
| iree_status_t status = check_matmul_results(stderr, &results); | |
| matmul_results_deinitialize(&results); | |
| return status; | |
| } | |
| // Generates a 3D matrix with pseudorandom values. | |
| StatusOr<vm::ref<iree_hal_buffer_view_t>> GenerateRandomMatrix3D( | |
| 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; | |
| 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); | |
| } | |
| // Checks batch matmul results by verifying each batch element independently. | |
| Status CheckBatchMatmulResults(iree_hal_device_t* device, int64_t batch, | |
| int64_t m, int64_t k, int64_t n, | |
| int32_t transpose_rhs, | |
| iree_hal_buffer_view_t* lhs, | |
| iree_hal_buffer_view_t* rhs, | |
| iree_hal_buffer_view_t* acc, | |
| iree_hal_buffer_view_t* actual_result) { | |
| // Download the full 3D buffers once. | |
| iree_hal_element_type_t lhs_type = iree_hal_buffer_view_element_type(lhs); | |
| iree_hal_element_type_t rhs_type = iree_hal_buffer_view_element_type(rhs); | |
| iree_hal_element_type_t acc_type = | |
| iree_hal_buffer_view_element_type(actual_result); | |
| iree_host_size_t lhs_elem_size = | |
| iree_hal_element_dense_byte_count(lhs_type); | |
| iree_host_size_t rhs_elem_size = | |
| iree_hal_element_dense_byte_count(rhs_type); | |
| iree_host_size_t acc_elem_size = | |
| iree_hal_element_dense_byte_count(acc_type); | |
| iree_host_size_t lhs_batch_bytes = m * k * lhs_elem_size; | |
| iree_host_size_t rhs_batch_bytes = k * n * rhs_elem_size; | |
| iree_host_size_t acc_batch_bytes = m * n * acc_elem_size; | |
| // Allocate host buffers and download. | |
| uint8_t* lhs_data = nullptr; | |
| uint8_t* rhs_data = nullptr; | |
| uint8_t* acc_data = nullptr; | |
| uint8_t* actual_data = nullptr; | |
| uint8_t* expected_data = nullptr; | |
| iree_host_size_t lhs_total = batch * lhs_batch_bytes; | |
| iree_host_size_t rhs_total = batch * rhs_batch_bytes; | |
| iree_host_size_t acc_total = batch * acc_batch_bytes; | |
| IREE_RETURN_IF_ERROR( | |
| iree_allocator_malloc(host_allocator_, lhs_total, (void**)&lhs_data)); | |
| IREE_RETURN_IF_ERROR( | |
| iree_allocator_malloc(host_allocator_, rhs_total, (void**)&rhs_data)); | |
| IREE_RETURN_IF_ERROR( | |
| iree_allocator_malloc(host_allocator_, acc_total, (void**)&acc_data)); | |
| IREE_RETURN_IF_ERROR(iree_allocator_malloc(host_allocator_, acc_total, | |
| (void**)&actual_data)); | |
| IREE_RETURN_IF_ERROR(iree_allocator_malloc(host_allocator_, acc_total, | |
| (void**)&expected_data)); | |
| iree_status_t status = iree_ok_status(); | |
| if (iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, iree_hal_buffer_view_buffer(lhs), 0, lhs_data, lhs_total, | |
| IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, iree_hal_buffer_view_buffer(rhs), 0, rhs_data, rhs_total, | |
| IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); | |
| } | |
| if (acc && iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, iree_hal_buffer_view_buffer(acc), 0, acc_data, acc_total, | |
| IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); | |
| } else { | |
| memset(acc_data, 0, acc_total); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = iree_hal_device_transfer_d2h( | |
| device, iree_hal_buffer_view_buffer(actual_result), 0, actual_data, | |
| acc_total, IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, | |
| iree_infinite_timeout()); | |
| } | |
| // Check each batch element independently. | |
| for (int64_t b = 0; b < batch && iree_status_is_ok(status); ++b) { | |
| iree_byte_span_t lhs_span = {lhs_data + b * lhs_batch_bytes, | |
| lhs_batch_bytes}; | |
| iree_byte_span_t rhs_span = {rhs_data + b * rhs_batch_bytes, | |
| rhs_batch_bytes}; | |
| iree_byte_span_t acc_span = {acc_data + b * acc_batch_bytes, | |
| acc_batch_bytes}; | |
| iree_byte_span_t actual_span = {actual_data + b * acc_batch_bytes, | |
| acc_batch_bytes}; | |
| iree_byte_span_t expected_span = {expected_data + b * acc_batch_bytes, | |
| acc_batch_bytes}; | |
| matmul_results_t batch_results = {}; | |
| batch_results.host_allocator = host_allocator_; | |
| batch_results.m = (iree_hal_dim_t)m; | |
| batch_results.k = (iree_hal_dim_t)k; | |
| batch_results.n = (iree_hal_dim_t)n; | |
| batch_results.lhs_type = lhs_type; | |
| batch_results.rhs_type = rhs_type; | |
| batch_results.acc_type = acc_type; | |
| batch_results.result_type = acc_type; | |
| batch_results.transpose_rhs = transpose_rhs != 0; | |
| batch_results.lhs_contents = lhs_span; | |
| batch_results.rhs_contents = rhs_span; | |
| batch_results.acc_contents = acc_span; | |
| batch_results.actual_contents = actual_span; | |
| batch_results.expected_contents = expected_span; | |
| status = check_matmul_results(stderr, &batch_results); | |
| } | |
| iree_allocator_free(host_allocator_, lhs_data); | |
| iree_allocator_free(host_allocator_, rhs_data); | |
| iree_allocator_free(host_allocator_, acc_data); | |
| iree_allocator_free(host_allocator_, actual_data); | |
| iree_allocator_free(host_allocator_, expected_data); | |
| return status; | |
| } | |
| private: | |
| iree_allocator_t host_allocator_; | |
| }; | |
| static const vm::NativeFunction<MatmulTestModuleState> | |
| kMatmulTestModuleFunctions[] = { | |
| vm::MakeNativeFunction("generate_random_matrix", | |
| &MatmulTestModuleState::GenerateRandomMatrix), | |
| vm::MakeNativeFunction("generate_random_matrix_3d", | |
| &MatmulTestModuleState::GenerateRandomMatrix3D), | |
| vm::MakeNativeFunction("check_matmul_results", | |
| &MatmulTestModuleState::CheckMatmulResults), | |
| vm::MakeNativeFunction("check_batch_matmul_results", | |
| &MatmulTestModuleState::CheckBatchMatmulResults), | |
| }; | |
| struct MatmulTestModule final : public vm::NativeModule<MatmulTestModuleState> { | |
| using vm::NativeModule<MatmulTestModuleState>::NativeModule; | |
| StatusOr<std::unique_ptr<MatmulTestModuleState>> CreateState( | |
| iree_allocator_t host_allocator) override { | |
| return std::make_unique<MatmulTestModuleState>(host_allocator); | |
| } | |
| StatusOr<std::unique_ptr<MatmulTestModuleState>> ForkState( | |
| MatmulTestModuleState* parent_state, | |
| iree_allocator_t host_allocator) override { | |
| return CreateState(host_allocator); | |
| } | |
| }; | |
| } // namespace iree | |
| static iree_status_t matmul_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::MatmulTestModule>( | |
| "matmul_test", /*version=*/0, instance, host_allocator, | |
| iree::span<const iree::vm::NativeFunction<iree::MatmulTestModuleState>>( | |
| iree::kMatmulTestModuleFunctions)); | |
| *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; | |
| } | |
| // Run the tests. Note that some modules may be compiled for other platforms | |
| // and not have the required architectures for execution within them - to keep | |
| // the test runner dumber we gracefully fail those cases by returning success. | |
| iree_status_t status = iree_test_utils_load_and_run_e2e_tests( | |
| iree_allocator_system(), matmul_test_module_create); | |
| int exit_code = EXIT_SUCCESS; | |
| if (!iree_status_is_ok(status)) { | |
| iree_status_fprint(stderr, status); | |
| bool is_device_unavailable = iree_status_is_not_found(status); | |
| iree_status_free(status); | |
| exit_code = is_device_unavailable ? EXIT_SUCCESS : EXIT_FAILURE; | |
| } | |
| IREE_TRACE_APP_EXIT(exit_code); | |
| return exit_code; | |
| } |