vec_iree: Add support for LLVM ukernel via plugin importer For Kelvin, we used efficient "aconv" op to implement custom mmt4d ukernel and achieved substantial performance improvement for many models. For Springbok, LLVM ukernel also works, but we disable it for now as it currently does not perform better than current version with RVV_ON. We also did clean-ups on the folder "vmvx_ukernel" to accommodate latest VMVX and LLVM ukernel APIs and design. Some further notes: - We renamed folder vmvx_ukernel to ukernel as it now supports for both VMVX and LLVM. - We removed custom ukerenl/query_tile_sizes.c to use IREE default for VMVX ukernel, as it was arbitrarily chosen before and VMVX ukernel now uses dynamic tile sizes. - We only enabled LLVM ukernel for Kelvin for now, as explained above. - For our custom ukernel implementation in ukernel/mmt4d_tile.c, we only considered the case of using Kelvin intrincics or using RVV. Generic ukernel case was omitted as it's not our interest and is already available as IREE builtin default ukernel. Change-Id: Ia96bcc6a676067cab412c8bc89788faa73403b6d
diff --git a/CMakeLists.txt b/CMakeLists.txt index 4fbba55..ffd43bd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -102,6 +102,7 @@ set(KELVIN_LINKER_SCRIPT "$ENV{ROOTDIR}/sw/kelvin/crt/kelvin.ld" CACHE PATH "Kelvin linker script path (default: kelvin.ld)") add_subdirectory($ENV{ROOTDIR}/sw/kelvin/crt kelvin) include(riscv_kelvin) + include_directories($ENV{ROOTDIR}/sw/kelvin) include_directories($ENV{ROOTDIR}/sw/kelvin/crt) endif() @@ -143,7 +144,7 @@ add_subdirectory(samples) add_subdirectory(ssd_postprocess) -add_subdirectory(vmvx_ukernel) +add_subdirectory(ukernel) # Add pigweed support include($ENV{ROOTDIR}/sw/pigweed/pw_build/pigweed.cmake)
diff --git a/README.md b/README.md index 32a081c..6eb4a2d 100644 --- a/README.md +++ b/README.md
@@ -35,7 +35,7 @@ * simple_vec_mul: Point-wise vector multiplication examples * ssd_postprocess: Vision postprocessing library for Single-Shot Detectors (SSD) -* vmvx_ukernel: Custom microkernel library for VMVX +* ukernel: Custom microkernel library for VMVX and LLVM ## Build the project
diff --git a/cmake/sparrow_static_module.cmake b/cmake/sparrow_static_module.cmake index 6fabab9..501adfd 100644 --- a/cmake/sparrow_static_module.cmake +++ b/cmake/sparrow_static_module.cmake
@@ -107,6 +107,12 @@ list(APPEND _COMPILER_ARGS "--iree-llvmcpu-target-abi=ilp32") list(APPEND _COMPILER_ARGS "--iree-llvmcpu-link-embedded=false") list(APPEND _COMPILER_ARGS "--iree-llvmcpu-use-fast-min-max-ops") + if (${BUILD_WITH_KELVIN}) + # Enable LLVM ukernel only in Kelvin for now + list(APPEND _COMPILER_ARGS "--iree-llvmcpu-enable-ukernels=all") + list(APPEND _COMPILER_ARGS "--iree-opt-data-tiling") + list(APPEND _COMPILER_ARGS "--iree-llvmcpu-link-ukernel-bitcode=false") + endif() if (${_RULE_INLINE_HAL}) list(APPEND _COMPILER_ARGS "--iree-execution-model=inline-dynamic") endif()
diff --git a/device/CMakeLists.txt b/device/CMakeLists.txt index c8f9222..46fa554 100644 --- a/device/CMakeLists.txt +++ b/device/CMakeLists.txt
@@ -21,9 +21,15 @@ "device_static_loader.c" DEPS iree::hal::drivers::local_sync::sync_driver + iree::hal::local::executable_plugin_manager iree::hal::local::loaders::static_library_loader + ukernel::ukernel_importer ) +if (${BUILD_WITH_KELVIN}) + target_compile_definitions(device_device_static_loader PUBLIC BUILD_KELVIN) +endif() + iree_cc_library( NAME device_vmvx_loader @@ -33,5 +39,5 @@ "device_vmvx_loader.c" DEPS iree::hal::drivers::local_sync::sync_driver - vmvx_ukernel::vmvx_module_loader + ukernel::vmvx_module_loader )
diff --git a/device/device_static_loader.c b/device/device_static_loader.c index 0d783f6..cff099f 100644 --- a/device/device_static_loader.c +++ b/device/device_static_loader.c
@@ -18,9 +18,12 @@ #include "device/device.h" #include "iree/hal/drivers/local_sync/sync_device.h" +#include "iree/hal/local/executable_plugin_manager.h" #include "iree/hal/local/loaders/static_library_loader.h" #include "model_util/model_api.h" +extern iree_hal_executable_import_provider_t importer_query(void); + // A function to create the HAL device from the different backend targets. // The HAL device and loader are returned based on the implementation, and they // must be released by the caller. @@ -33,13 +36,31 @@ iree_hal_sync_device_params_t params; iree_hal_sync_device_params_initialize(¶ms); + // Create plugin manager + iree_hal_executable_plugin_manager_t* plugin_manager = NULL; + +#ifdef BUILD_KELVIN + // Enable plugin manager only in Kelvin for now + if (iree_status_is_ok(status)) { + status = iree_hal_executable_plugin_manager_create( + /*capacity=*/1, host_allocator, &plugin_manager); + } + + // Register import provider + if (iree_status_is_ok(status)) { + status = iree_hal_executable_plugin_manager_register_provider( + plugin_manager, importer_query()); + } +#endif // # ifdef BUILD_KELVIN + // Load the statically embedded library const iree_hal_executable_library_query_fn_t libraries[] = {library_query()}; if (iree_status_is_ok(status)) { status = iree_hal_static_library_loader_create( IREE_ARRAYSIZE(libraries), libraries, - iree_hal_executable_import_provider_null(), host_allocator, loader); + iree_hal_executable_plugin_manager_provider(plugin_manager), + host_allocator, loader); } // Use the default host allocator for buffer allocations.
diff --git a/vmvx_ukernel/CMakeLists.txt b/ukernel/CMakeLists.txt similarity index 86% rename from vmvx_ukernel/CMakeLists.txt rename to ukernel/CMakeLists.txt index 86de17c..f69f8eb 100644 --- a/vmvx_ukernel/CMakeLists.txt +++ b/ukernel/CMakeLists.txt
@@ -20,6 +20,8 @@ if (${BUILD_WITH_RVV}) set(PREFIX_UKERNEL "${CMAKE_CURRENT_SOURCE_DIR}") set(PREFIX_VMVX "${CMAKE_CURRENT_SOURCE_DIR}") +elseif (${BUILD_WITH_KELVIN}) + set(PREFIX_UKERNEL "${CMAKE_CURRENT_SOURCE_DIR}") endif() iree_cc_library( @@ -31,16 +33,20 @@ "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/mmt4d.c" "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/pack.c" "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/pack_tile.c" + "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/query_tile_sizes.c" "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/unpack.c" "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/unpack_tile.c" "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/weak.c" "${PREFIX_UKERNEL}/mmt4d_tile.c" - "${PREFIX_UKERNEL}/query_tile_sizes.c" DEPS iree::builtins::ukernel::internal_headers PUBLIC ) +if (${BUILD_WITH_KELVIN}) + target_compile_definitions(ukernel_ukernel PUBLIC BUILD_KELVIN) +endif() + iree_cc_library( NAME vmvx @@ -73,3 +79,12 @@ "IREE_HAVE_HAL_EXECUTABLE_LOADER_VMVX_MODULE=1" PUBLIC ) + +iree_cc_library( + NAME + ukernel_importer + SRCS + "importer.c" + DEPS + ::ukernel +)
diff --git a/vmvx_ukernel/elementwise.c b/ukernel/elementwise.c similarity index 100% rename from vmvx_ukernel/elementwise.c rename to ukernel/elementwise.c
diff --git a/ukernel/importer.c b/ukernel/importer.c new file mode 100644 index 0000000..cf0215d --- /dev/null +++ b/ukernel/importer.c
@@ -0,0 +1,56 @@ +/* + * 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/api.h" +#include "iree/hal/local/executable_loader.h" +#include "iree/hal/local/executable_plugin.h" + +// Importer entry point wrapping the actual ukernel. +static int iree_uk_importer_mmt4d(void* params_ptr, void* context, + void* reserved) { + iree_uk_mmt4d((const iree_uk_mmt4d_params_t*)params_ptr); + return 0; +} + +// Called to resolve the iree_uk_mmt4d import by symbol name. +static iree_status_t iree_uk_importer_resolve( + void* self, iree_host_size_t count, const char* const* symbol_names, + void** out_fn_ptrs, void** out_fn_contexts, + iree_hal_executable_import_resolution_t* out_resolution) { + *out_resolution = 0; + bool any_required_not_found = true; + for (size_t i = 0; i < count; ++i) { + if (out_fn_ptrs[i]) continue; + bool found = iree_hal_executable_plugin_strcmp(symbol_names[i], + "iree_uk_mmt4d") == 0; + if (found) { + out_fn_ptrs[i] = (void*)iree_uk_importer_mmt4d; + out_fn_contexts[i] = NULL; + any_required_not_found = false; + } else { + break; + } + } + return any_required_not_found ? iree_make_status(IREE_STATUS_NOT_FOUND) + : iree_ok_status(); +} + +iree_hal_executable_import_provider_t importer_query(void) { + static iree_hal_executable_import_provider_t importer = { + .resolve = iree_uk_importer_resolve, + }; + return importer; +}
diff --git a/ukernel/mmt4d_tile.c b/ukernel/mmt4d_tile.c new file mode 100644 index 0000000..a45df27 --- /dev/null +++ b/ukernel/mmt4d_tile.c
@@ -0,0 +1,531 @@ +/* + * 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/mmt4d_internal.h" + +#ifdef BUILD_KELVIN + +#include <kelvin.h> + +static_assert(sizeof(struct vconv_u8_t) == 4); +union { + struct vconv_u8_t conv; + uint32_t raw; +} cmds; + +static iree_uk_int32_t buffer[8] __attribute__((aligned)); + +// dot product s8xx8->s32 +static iree_uk_int32_t dot_product_s8s8(const iree_uk_int8_t* u, + const iree_uk_int8_t* w, int n) { + int vl; + getmaxvl_w(vl); + vdup_w_x_m(v0, 0); + + while (n) { + int count = n < vl ? n : vl; + + vld_b_lp_xx(v8, u, count); + vaddw_h_vx(v8, v8, 0); + vzip_h_vv(v10, v8, v9); + + vld_b_lp_xx(v16, w, count); + vaddw_h_vx(v16, v16, 0); + vzip_h_vv(v18, v16, v17); + + vmulw_w_vv(v4, v10, v18); + vzip_w_vv(v1, v4, v5); + + vadd_w_vv(v0, v0, v1); + n -= count; + } + + vst_w_l_xx(v0, buffer, vl); + + iree_uk_int32_t sum = 0; + for (int i = 0; i < vl; ++i) sum += buffer[i]; + return sum; +} + +// dot product s16xs8->s32 +static iree_uk_int32_t dot_product_s16s8(const iree_uk_int16_t* u, + const iree_uk_int8_t* w, int n) { + int vl; + getmaxvl_w(vl); + vdup_w_x_m(v0, 0); + + while (n) { + int count = n < vl ? n : vl; + + vld_h_lp_xx(v8, u, count); + + vld_b_lp_xx(v16, w, count); + vaddw_h_vx(v16, v16, 0); + vzip_h_vv(v18, v16, v17); + + vmulw_w_vv(v4, v8, v18); + vzip_w_vv(v1, v4, v5); + + vadd_w_vv(v0, v0, v1); + n -= count; + } + + vst_w_l_xx(v0, buffer, vl); + + iree_uk_int32_t sum = 0; + for (int i = 0; i < vl; ++i) sum += buffer[i]; + return sum; +} + +// Matrix multiplication s8xs8->s32. +// lhs (row-major): 8/4/2/1x4, rhs (column-major): 4x8 +// out (row-major): 8/4/2/1x8 +static void matmul_s8s8(const iree_uk_int8_t* lhs, const iree_uk_int8_t* rhs, + iree_uk_int32_t* out, + const iree_uk_mmt4d_params_t* params) { + iree_uk_int16_t M0 = params->M0; + iree_uk_int16_t K = params->K; + IREE_UK_ASSERT(params->N0 == 8 && params->K0 == 4); + IREE_UK_ASSERT(M0 == 8 || M0 == 4 || M0 == 2 || M0 == 1); + IREE_UK_ASSERT(K == 1 || K == 2 || K == 4 || (K % 8 == 0)); + + cmds.conv.mode = 0; + cmds.conv.start = 0; + cmds.conv.stop = K > 8 ? 7 : K - 1; + cmds.conv.sbias1 = 0; + cmds.conv.sdata1 = true; + cmds.conv.sbias2 = 0; + cmds.conv.sdata2 = true; + + const iree_uk_int8_t* p_in = lhs; + const iree_uk_int8_t* p_flt = rhs; + iree_uk_int32_t* p_out = out; + + // load initial output tile values + if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) { + if (M0 <= 2) { + vld_w_p_x(v56, p_out); + if (M0 == 2) { + vld_w_p_x(v57, p_out); + } + } else { + vld_w_p_x_m(v56, p_out); + } + if (M0 == 8) { + vld_w_p_x_m(v60, p_out); + } + p_out = out; + } else { + vdup_w_x_m(v56, 0); + vdup_w_x_m(v60, 0); + } + + iree_uk_int16_t k = K; + while (k > 0) { + k -= 8; + // load LHS and RHS + if (K <= 2) { + vld_b_lp_xx(v0, p_in, 4 * M0); + vld_b_p_x(v8, p_flt); + if (K == 2) { + vld_b_lp_xx(v1, p_in, 4 * M0); + vld_b_p_x(v9, p_flt); + } + } else { + vld_b_sp_xx_m(v0, p_in, 4 * M0); + vld_b_p_x_m(v8, p_flt); + } + if (K >= 8) { + vld_b_sp_xx_m(v4, p_in, 4 * M0); + vld_b_p_x_m(v12, p_flt); + } + + // re-arrange LHS + vzip_w_vv_m(v20, v0, v4); + vzip_w_vv_m(v28, v20, v24); + vzip_w_vv_m(v0, v28, v32); + + // matrix multiplication + aconv_vxv(v48, v0, cmds, v8); + } + + vcget(v48); + + // de-interleaving + vzip_w_vv(v20, v48, v49); + vzip_w_vv(v22, v50, v51); + vzip_w_vv(v24, v20, v22); + vzip_w_vv(v26, v21, v23); + + if (M0 <= 2) { + vadd_w_vv(v56, v56, v24); + vst_w_p_x(v56, p_out); + if (M0 == 2) { + vadd_w_vv(v57, v57, v25); + vst_w_p_x(v57, p_out); + } + } else { + vadd_w_vv_m(v56, v56, v24); + vst_w_p_x_m(v56, p_out); + } + + if (M0 == 8) { + // de-interleaving + vzip_w_vv(v28, v52, v53); + vzip_w_vv(v30, v54, v55); + vzip_w_vv(v32, v28, v30); + vzip_w_vv(v34, v29, v31); + + vadd_w_vv_m(v60, v60, v32); + vst_w_p_x_m(v60, p_out); + } +} + +// Matrix multiplication s16xs8->s32. +// lhs (row-major): 8/4/2/1x4, rhs (column-major): 4x8 +// out (row-major): 8/4/2/1x8 +// s16 * s8 = (s8_hi * s8) << 8 + (u8_lo * s8) +// where s8_hi = s16[15:8]; u8_lo = s16[7:0] +static void matmul_s16s8(const iree_uk_int16_t* lhs, const iree_uk_int8_t* rhs, + iree_uk_int32_t* out, + const iree_uk_mmt4d_params_t* params) { + iree_uk_int16_t M0 = params->M0; + iree_uk_int16_t K = params->K; + IREE_UK_ASSERT(params->N0 == 8 && params->K0 == 4); + IREE_UK_ASSERT(M0 == 8 || M0 == 4 || M0 == 2 || M0 == 1); + IREE_UK_ASSERT(K == 1 || K == 2 || K == 4 || (K % 8 == 0)); + + cmds.conv.mode = 0; + cmds.conv.start = 0; + cmds.conv.stop = K > 8 ? 7 : K - 1; + cmds.conv.sbias1 = 0; + cmds.conv.sdata1 = true; + cmds.conv.sbias2 = 0; + cmds.conv.sdata2 = true; + + const iree_uk_int16_t* p_in = lhs; + const iree_uk_int8_t* p_flt = rhs; + iree_uk_int32_t* p_out = out; + + // load initial output tile values + if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) { + if (M0 <= 2) { + vld_w_p_x(v56, p_out); + if (M0 == 2) { + vld_w_p_x(v57, p_out); + } + } else { + vld_w_p_x_m(v56, p_out); + } + if (M0 == 8) { + vld_w_p_x_m(v60, p_out); + } + } else { + vdup_w_x_m(v56, 0); + vdup_w_x_m(v60, 0); + } + + // first pass - high + iree_uk_int16_t k = K; + while (k > 0) { + k -= 8; + // load LHS high bits and RHS + if (K == 1) { + vld_h_lp_xx(v12, p_in, 2 * M0); + vld_h_lp_xx(v13, p_in, 2 * M0); + vodd_b_vv(v0, v12, v13); + vld_b_p_x(v8, p_flt); + } else if (K == 2) { + vld_h_sp_xx_m(v12, p_in, 2 * M0); + vodd_b_vv(v0, v12, v13); + vodd_b_vv(v1, v14, v15); + vld_b_p_x(v8, p_flt); + vld_b_p_x(v9, p_flt); + } + if (K >= 4) { + vld_h_sp_xx_m(v12, p_in, 2 * M0); + vld_h_sp_xx_m(v16, p_in, 2 * M0); + vodd_b_vv_m(v0, v12, v16); + vld_b_p_x_m(v8, p_flt); + } + if (K >= 8) { + vld_h_sp_xx_m(v20, p_in, 2 * M0); + vld_h_sp_xx_m(v24, p_in, 2 * M0); + vodd_b_vv_m(v4, v20, v24); + vld_b_p_x_m(v12, p_flt); + } + + // re-arrange LHS + vzip_w_vv_m(v20, v0, v4); + vzip_w_vv_m(v28, v20, v24); + vzip_w_vv_m(v0, v28, v32); + + // matrix multiplication + aconv_vxv(v48, v0, cmds, v8); + } + + vcget(v48); + + // left shift 8 bits and store to accumulator + vsll_w_vx_m(v48, v48, 8); + vsll_w_vx_m(v52, v52, 8); + acset_v(v48, v48); + + // second pass - low + cmds.conv.sdata1 = false; + p_in = lhs; + p_flt = rhs; + p_out = out; + k = K; + while (k > 0) { + k -= 8; + // load LHS low bits and RHS + if (K == 1) { + vld_h_lp_xx(v12, p_in, 2 * M0); + vld_h_lp_xx(v13, p_in, 2 * M0); + vevn_b_vv(v0, v12, v13); + vld_b_p_x(v8, p_flt); + } else if (K == 2) { + vld_h_sp_xx_m(v12, p_in, 2 * M0); + vevn_b_vv(v0, v12, v13); + vevn_b_vv(v1, v14, v15); + vld_b_p_x(v8, p_flt); + vld_b_p_x(v9, p_flt); + } + if (K >= 4) { + vld_h_sp_xx_m(v12, p_in, 2 * M0); + vld_h_sp_xx_m(v16, p_in, 2 * M0); + vevn_b_vv_m(v0, v12, v16); + vld_b_p_x_m(v8, p_flt); + } + if (K >= 8) { + vld_h_sp_xx_m(v20, p_in, 2 * M0); + vld_h_sp_xx_m(v24, p_in, 2 * M0); + vevn_b_vv_m(v4, v20, v24); + vld_b_p_x_m(v12, p_flt); + } + + // re-arrange LHS + vzip_w_vv_m(v20, v0, v4); + vzip_w_vv_m(v28, v20, v24); + vzip_w_vv_m(v0, v28, v32); + + // matrix multiplication + aconv_vxv(v48, v0, cmds, v8); + } + + vcget(v48); + + // de-interleaving + vzip_w_vv(v20, v48, v49); + vzip_w_vv(v22, v50, v51); + vzip_w_vv(v24, v20, v22); + vzip_w_vv(v26, v21, v23); + + if (M0 <= 2) { + vadd_w_vv(v56, v56, v24); + vst_w_p_x(v56, p_out); + if (M0 == 2) { + vadd_w_vv(v57, v57, v25); + vst_w_p_x(v57, p_out); + } + } else { + vadd_w_vv_m(v56, v56, v24); + vst_w_p_x_m(v56, p_out); + } + + if (M0 == 8) { + // de-interleaving + vzip_w_vv(v28, v52, v53); + vzip_w_vv(v30, v54, v55); + vzip_w_vv(v32, v28, v30); + vzip_w_vv(v34, v29, v31); + + vadd_w_vv_m(v60, v60, v32); + vst_w_p_x_m(v60, p_out); + } +} + +#else // RVV implementation in Springbok + +#include <riscv_vector.h> +// Calculate the dot product of two int8 vectors using RVV +static iree_uk_int32_t dot_product_s8s8(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 = __riscv_vsetvl_e8m4(n - i); + vu = __riscv_vle8_v_i8m4(u + i, vl); // load + vw = __riscv_vle8_v_i8m4(w + i, vl); // load + vx = __riscv_vwmul(vu, vw, vl); // multiply + v_sum = __riscv_vmv_v_x_i32m1(0, vl); // init + v_sum = __riscv_vwredsum(vx, v_sum, vl); // sum + sum += __riscv_vmv_x(v_sum); + } + return sum; +} + +// Calculate the dot product of int16 x int8 vectors using RVV +static iree_uk_int32_t dot_product_s16s8(const iree_uk_int16_t* u, + const iree_uk_int8_t* w, int n) { + size_t vl; + // auxiliary variables + vint16m4_t vu, vw; + vint8m2_t vy; + vint32m8_t vx; + vint32m1_t v_sum; + iree_uk_int32_t sum = 0; + for (size_t i = 0; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vu = __riscv_vle16_v_i16m4(u + i, vl); // load + vy = __riscv_vle8_v_i8m2(w + i, vl); // load + vw = __riscv_vwadd_vx(vy, 0, vl); // widen + vx = __riscv_vwmul(vu, vw, vl); // multiply + v_sum = __riscv_vmv_v_x_i32m1(0, vl); // init + v_sum = __riscv_vredsum(vx, v_sum, vl); // sum + sum += __riscv_vmv_x(v_sum); + } + return sum; +} + +#endif // # ifdef BUILD_KELVIN + +// RVV implementation of matmul tile, s8*s8->s32 case. +static void iree_uk_mmt4d_tile_s8s8s32_custom( + void* out_tile_untyped, const void* lhs_panel_untyped, + const void* rhs_panel_untyped, 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; + +#ifdef BUILD_KELVIN + if (N0 == 8 && K0 == 4 && (M0 == 8 || M0 == 4 || M0 == 2 || M0 == 1)) { + matmul_s8s8(lhs_panel, rhs_panel, out_tile, params); + return; + } +#endif // # ifdef BUILD_KELVIN + + // Initialize the accumulator tile. + if (!(params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE)) { + for (iree_uk_int32_t i = 0; i < M0 * N0; i++) out_tile[i] = 0; + } + + // Accumulation loop. + for (iree_uk_index_t k = 0; k < params->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_s8s8(lhs_panel + i0 * K0, rhs_panel + j0 * K0, K0); + } + } + lhs_panel += M0 * K0; + rhs_panel += N0 * K0; + } +} + +// implementation of matmul tile, s16*s8->s32 case. +static void iree_uk_mmt4d_tile_s16s8s32_custom( + void* out_tile_untyped, const void* lhs_panel_untyped, + const void* rhs_panel_untyped, const iree_uk_mmt4d_params_t* params) { + iree_uk_int32_t* out_tile = out_tile_untyped; + const iree_uk_int16_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; + +#ifdef BUILD_KELVIN + if (N0 == 8 && K0 == 4 && (M0 == 8 || M0 == 4 || M0 == 2 || M0 == 1)) { + matmul_s16s8(lhs_panel, rhs_panel, out_tile, params); + return; + } +#endif // # ifdef BUILD_KELVIN + + // Initialize the accumulator tile. + if (!(params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE)) { + for (iree_uk_int32_t i = 0; i < M0 * N0; i++) out_tile[i] = 0; + } + + // Accumulation loop. + for (iree_uk_index_t k = 0; k < params->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_s16s8(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, 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 (params->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 < params->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_f32 = lhs_panel[i0 * K0 + k0]; + float rhs_f32 = rhs_panel[j0 * K0 + k0]; + acc[i0 * N0 + j0] += lhs_f32 * rhs_f32; + } + } + } + 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) { + switch (iree_uk_mmt4d_type(params->flags)) { + case iree_uk_mmt4d_type_s8s8s32: + return iree_uk_mmt4d_tile_s8s8s32_custom; + case iree_uk_mmt4d_type_s16s8s32: + return iree_uk_mmt4d_tile_s16s8s32_custom; + case iree_uk_mmt4d_type_f32f32f32: + return iree_uk_mmt4d_tile_f32f32f32_generic; + default: + IREE_UK_ASSUME_UNREACHABLE; + return 0; + } +}
diff --git a/vmvx_ukernel/mmt4d_tile.c b/vmvx_ukernel/mmt4d_tile.c deleted file mode 100644 index c200cf0..0000000 --- a/vmvx_ukernel/mmt4d_tile.c +++ /dev/null
@@ -1,118 +0,0 @@ -/* - * 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 = __riscv_vsetvl_e8m4(n - i); - vu = __riscv_vle8_v_i8m4(u + i, vl); // load - vw = __riscv_vle8_v_i8m4(w + i, vl); // load - vx = __riscv_vwmul(vu, vw, vl); // multiply - v_sum = __riscv_vmv_v_x_i32m1(0, vl); // init - v_sum = __riscv_vwredsum(vx, v_sum, vl); // sum - sum += __riscv_vmv_x(v_sum); - } - return sum; -} - -// RVV implementation of matmul tile, i8*i8->i32 case. -static void iree_uk_mmt4d_tile_s8s8s32_rvv( - void* out_tile_untyped, const void* lhs_panel_untyped, - const void* rhs_panel_untyped, 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 (!(params->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 < params->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, 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 (params->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 < params->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_s8s8s32: - return iree_uk_mmt4d_tile_s8s8s32_rvv; - default: - // shouldn't happen, validated earlier. - IREE_UK_ASSUME_UNREACHABLE; - return 0; - } -}
diff --git a/vmvx_ukernel/query_tile_sizes.c b/vmvx_ukernel/query_tile_sizes.c deleted file mode 100644 index ae3ae19..0000000 --- a/vmvx_ukernel/query_tile_sizes.c +++ /dev/null
@@ -1,94 +0,0 @@ -/* - * 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_internal.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 int 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; - } - return 0; -}