blob: 10105b66b72e192ff57329404fe8fe64a5a191b7 [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 <riscv_vector.h>
#include <string.h>
#include "iree/builtins/ukernel/mmt4d_internal.h"
// Calculate the dot product of two int8 vectors using RVV
static iree_uk_int32_t dot_product_rvv(const iree_uk_int8_t* u,
const iree_uk_int8_t* w, int n) {
size_t vl;
// auxiliary variables
vint8m4_t vu, vw;
vint16m8_t vx;
vint32m1_t v_sum;
iree_uk_int32_t sum = 0;
for (size_t i = 0; i < n; i += vl) {
vl = vsetvl_e8m4(n - i);
vu = vle8_v_i8m4(u + i, vl); // load
vw = vle8_v_i8m4(w + i, vl); // load
vx = vwmul(vu, vw, vl); // multiply
v_sum = vmv_s(v_sum, 0, vl); // init
v_sum = vwredsum(v_sum, vx, v_sum, vl); // sum
sum += vmv_x(v_sum);
}
return sum;
}
// RVV implementation of matmul tile, i8*i8->i32 case.
static void iree_uk_mmt4d_tile_i8i8i32_rvv(
void* out_tile_untyped, const void* lhs_panel_untyped,
const void* rhs_panel_untyped, iree_uk_int32_t K, iree_uk_uint32_t flags,
const iree_uk_mmt4d_params_t* params) {
iree_uk_int32_t* out_tile = out_tile_untyped;
const iree_uk_int8_t* lhs_panel = lhs_panel_untyped;
const iree_uk_int8_t* rhs_panel = rhs_panel_untyped;
iree_uk_int16_t M0 = params->M0;
iree_uk_int16_t N0 = params->N0;
iree_uk_int16_t K0 = params->K0;
// Initialize the accumulator tile.
if (!(flags & IREE_UK_FLAG_MMT4D_ACCUMULATE)) {
memset(out_tile, 0, M0 * N0 * sizeof(iree_uk_int32_t));
}
// Accumulation loop.
for (iree_uk_index_t k = 0; k < K; ++k) {
for (iree_uk_index_t i0 = 0; i0 < M0; ++i0) {
for (iree_uk_index_t j0 = 0; j0 < N0; ++j0) {
out_tile[i0 * N0 + j0] +=
dot_product_rvv(lhs_panel + i0 * K0, rhs_panel + j0 * K0, K0);
}
}
lhs_panel += M0 * K0;
rhs_panel += N0 * K0;
}
}
// Generic implementation of matmul tile, f32*f32->f32 case.
static void iree_uk_mmt4d_tile_f32f32f32_generic(
void* out_tile_untyped, const void* lhs_panel_untyped,
const void* rhs_panel_untyped, iree_uk_int32_t K, iree_uk_uint32_t flags,
const iree_uk_mmt4d_params_t* params) {
float* out_tile = out_tile_untyped;
const float* lhs_panel = lhs_panel_untyped;
const float* rhs_panel = rhs_panel_untyped;
iree_uk_int16_t M0 = params->M0;
iree_uk_int16_t N0 = params->N0;
iree_uk_int16_t K0 = params->K0;
// Initialize the local accumulator tile.
float acc[iree_uk_mmt4d_tile_generic_max_bytes / sizeof(*out_tile)];
if (flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
for (int i = 0; i < M0 * N0; ++i) acc[i] = out_tile[i];
} else {
for (int i = 0; i < M0 * N0; ++i) acc[i] = 0;
}
// Accumulation loop.
for (iree_uk_index_t k = 0; k < K; ++k) {
for (iree_uk_index_t i0 = 0; i0 < M0; ++i0) {
for (iree_uk_index_t j0 = 0; j0 < N0; ++j0) {
for (iree_uk_index_t k0 = 0; k0 < K0; ++k0) {
float lhs_val = lhs_panel[i0 * K0 + k0];
float rhs_val = rhs_panel[j0 * K0 + k0];
acc[i0 * N0 + j0] += lhs_val * rhs_val;
}
}
}
lhs_panel += M0 * K0;
rhs_panel += N0 * K0;
}
// Store the local accumulator tile to the destination.
for (int i = 0; i < M0 * N0; ++i) out_tile[i] = acc[i];
}
iree_uk_mmt4d_tile_func_t iree_uk_mmt4d_select_tile_func(
const iree_uk_mmt4d_params_t* params) {
// TODO(lundong): to be replaced with Kelvin
switch (iree_uk_mmt4d_type(params->flags)) {
case iree_uk_mmt4d_type_f32f32f32:
return iree_uk_mmt4d_tile_f32f32f32_generic;
case iree_uk_mmt4d_type_i8i8i32:
return iree_uk_mmt4d_tile_i8i8i32_rvv;
default:
// shouldn't happen, validated earlier.
IREE_UK_ASSUME_UNREACHABLE;
return 0;
}
}