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