blob: 4bab96e6a2149994cb3ba0895c714bebb7c5882a [file] [log] [blame]
/*
* Copyright 2023 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "iree/builtins/ukernel/query_tile_sizes.h"
static bool iree_uk_query_tile_sizes_operation_is_matmul(
iree_uk_uint32_t flags) {
iree_uk_uint32_t op = iree_uk_query_tile_sizes_operation(flags);
return op == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_F32F32F32 ||
op == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_I8I8I32;
}
static void iree_uk_query_tile_sizes_2d_validate(
const iree_uk_query_tile_sizes_2d_params_t* params) {
#ifdef IREE_UK_ENABLE_ASSERTS
IREE_UK_ASSERT(iree_uk_query_tile_sizes_operation_is_matmul(params->flags));
iree_uk_uint32_t role = iree_uk_query_tile_sizes_operand_role(params->flags);
IREE_UK_ASSERT(role == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_LHS ||
role == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RHS ||
role == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RESULT);
const iree_uk_int64_t kDynamic = IREE_UK_INT64_MIN;
IREE_UK_ASSERT((params->size0 >= 0 || params->size0 == kDynamic) ||
(params->size1 >= 0 || params->size1 == kDynamic));
#endif // IREE_UK_ENABLE_ASSERTS
}
static iree_uk_matmul_tile_sizes_t iree_uk_query_matmul_tile_sizes_generic(
const iree_uk_query_tile_sizes_2d_params_t* params) {
// Dummy values, originally taken from what was used on ARM_64 +dotprod for
// i8i8i32. Not particularly meaningful outside of that case, just is what
// some tests have been written against.
(void)params;
return (iree_uk_matmul_tile_sizes_t){.M = 8, .K = 4, .N = 8};
}
// Experimental tile sizes for RVV
static iree_uk_matmul_tile_sizes_t iree_uk_query_matmul_tile_sizes_rvv(
const iree_uk_query_tile_sizes_2d_params_t* params) {
(void)params;
return (iree_uk_matmul_tile_sizes_t){.M = 16, .K = 16, .N = 16};
}
static void iree_uk_query_tile_sizes_2d_matmul(
const iree_uk_query_tile_sizes_2d_params_t* params,
iree_uk_query_tile_sizes_2d_out_params_t* out_params) {
iree_uk_matmul_tile_sizes_t matmul_tile_sizes;
if (iree_uk_query_tile_sizes_operation(params->flags) ==
IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_I8I8I32) {
matmul_tile_sizes = iree_uk_query_matmul_tile_sizes_rvv(params);
} else {
matmul_tile_sizes = iree_uk_query_matmul_tile_sizes_generic(params);
}
iree_uk_uint32_t role = iree_uk_query_tile_sizes_operand_role(params->flags);
if (role == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_LHS) {
out_params->tile_size0 = matmul_tile_sizes.M;
out_params->tile_size1 = matmul_tile_sizes.K;
} else if (role == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RHS) {
out_params->tile_size0 = matmul_tile_sizes.N;
out_params->tile_size1 = matmul_tile_sizes.K;
} else if (role == IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RESULT) {
out_params->tile_size0 = matmul_tile_sizes.M;
out_params->tile_size1 = matmul_tile_sizes.N;
} else {
// Can't happen, validated earlier.
IREE_UK_ASSUME_UNREACHABLE;
}
}
IREE_UK_EXPORT void iree_uk_query_tile_sizes_2d(
const iree_uk_query_tile_sizes_2d_params_t* params,
iree_uk_query_tile_sizes_2d_out_params_t* out_params) {
iree_uk_query_tile_sizes_2d_validate(params);
if (iree_uk_query_tile_sizes_operation_is_matmul(params->flags)) {
iree_uk_query_tile_sizes_2d_matmul(params, out_params);
} else {
// Can't happen, validated earlier.
IREE_UK_ASSUME_UNREACHABLE;
}
}