sw:vec_iree: Fix microkernel breakage
Fix the ukernel breakage due to
https://github.com/iree-org/iree/pull/12015.
Change-Id: Id34c67b06f25ebd13445a78e29b277833d105493
diff --git a/vmvx_ukernel/mmt4d_tile.c b/vmvx_ukernel/mmt4d_tile.c
new file mode 100644
index 0000000..a591c66
--- /dev/null
+++ b/vmvx_ukernel/mmt4d_tile.c
@@ -0,0 +1,120 @@
+/*
+ * 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_tile.h"
+
+#include <riscv_vector.h>
+#include <string.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_ACCUMULATE)) {
+ memset(out_tile, 0, M0 * N0 * sizeof(iree_uk_int32_t));
+ }
+ // Accumulation loop.
+ for (iree_uk_ssize_t k = 0; k < K; ++k) {
+ for (iree_uk_ssize_t i0 = 0; i0 < M0; ++i0) {
+ for (iree_uk_ssize_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_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_ssize_t k = 0; k < K; ++k) {
+ for (iree_uk_ssize_t i0 = 0; i0 < M0; ++i0) {
+ for (iree_uk_ssize_t j0 = 0; j0 < N0; ++j0) {
+ for (iree_uk_ssize_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 (params->type) {
+ 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;
+ }
+}
