Simplify how mmt4d ukernels deal with the K=0 case. (#16137)
The case of K=0 is vacuous when accumulating, so we just early-return.
This doesn't change.
When not accumulating, it amounts to clearing the accumulator. We used
to have dedicated code for that in `mmt4d.c`. That then allowed our
ukernels tile-functions to `IREE_UK_ASSUME(K >= 1);`. That makes the asm
a little simpler. However, that only really mattered when we thought we
would be hand-writing that assembly. Now that we're not, it doesn't
really matter, and the slightly smaller tile function code size is
outweighed by the additional code in `mmt4d.c`.
Overall, this PR drops 40 lines of source code and makes bitcode 2%
smaller.
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64.c
index 2c52225..cdadc7d 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64.c
@@ -25,7 +25,6 @@
acc[i] = vdupq_n_f32(0);
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
float32x4_t rhs[2];
for (int i = 0; i < 2; ++i) {
@@ -110,7 +109,6 @@
acc[i] = vdupq_n_f32(0);
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
float32x4_t rhs[2];
for (int i = 0; i < 2; ++i) {
@@ -218,7 +216,6 @@
acc[i] = vdupq_n_s32(0);
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
int16x8_t rhs = vmovl_s8(vld1_s8(rhs_ptr));
rhs_ptr += 8;
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_bf16.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_bf16.c
index 4425ec5..2425500 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_bf16.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_bf16.c
@@ -75,7 +75,6 @@
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
bfloat16x8_t rhs[4];
for (int i = 0; i < 4; ++i) {
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_dotprod.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_dotprod.c
index dce5200..8bf713c 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_dotprod.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_dotprod.c
@@ -25,7 +25,6 @@
acc[i] = vdupq_n_s32(0);
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
int8x16_t rhs[2];
for (int i = 0; i < 2; ++i) {
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fp16fml.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fp16fml.c
index 919613b..0b574e4 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fp16fml.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fp16fml.c
@@ -51,7 +51,6 @@
acc[i] = vdupq_n_f32(0);
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
float16x8_t rhs = vld1q_f16(rhs_ptr);
rhs_ptr += 8;
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fullfp16.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fullfp16.c
index 6d1660d..e869d78 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fullfp16.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fullfp16.c
@@ -25,7 +25,6 @@
acc[i] = vdupq_n_f16(0);
}
}
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
float16x8_t rhs = vld1q_f16(rhs_ptr);
rhs_ptr += 8;
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c
index 97c6a52..49f4116 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c
@@ -58,8 +58,6 @@
}
}
}
-
- IREE_UK_ASSUME(params->K >= 1);
for (int k = 0; k < params->K; ++k) {
int8x16_t rhs[4];
for (int i = 0; i < 4; ++i) {
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/mmt4d_x86_64_avx512_vnni.c b/runtime/src/iree/builtins/ukernel/arch/x86_64/mmt4d_x86_64_avx512_vnni.c
index 5389c82..ffd38c7 100644
--- a/runtime/src/iree/builtins/ukernel/arch/x86_64/mmt4d_x86_64_avx512_vnni.c
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/mmt4d_x86_64_avx512_vnni.c
@@ -307,7 +307,6 @@
const __m128i idx_2_mod_4 = _mm_set1_epi32(0x0e0a0602);
const __m128i idx_3_mod_4 = _mm_set1_epi32(0x0f0b0703);
const __m512i mask_0f = _mm512_set1_epi8(0x0f);
- IREE_UK_ASSUME(params->K >= 1);
for (iree_uk_int32_t k = 0; k < params->K; ++k) {
// Load 8xs16 LHS data.
__m128i lhs = _mm_loadu_si128((const __m128i*)lhs_ptr);
diff --git a/runtime/src/iree/builtins/ukernel/common.h b/runtime/src/iree/builtins/ukernel/common.h
index 450c546..21f0342 100644
--- a/runtime/src/iree/builtins/ukernel/common.h
+++ b/runtime/src/iree/builtins/ukernel/common.h
@@ -155,13 +155,6 @@
#define IREE_UK_ASSUME_UNREACHABLE
#endif // IREE_UK_HAVE_BUILTIN(__builtin_unreachable)
-#define IREE_UK_ASSUME(condition) \
- do { \
- if (!(condition)) { \
- IREE_UK_ASSUME_UNREACHABLE; \
- } \
- } while (false)
-
#if IREE_UK_HAVE_ATTRIBUTE(noinline) || defined(IREE_UK_COMPILER_GCC)
#define IREE_UK_ATTRIBUTE_NOINLINE __attribute__((noinline))
#else
diff --git a/runtime/src/iree/builtins/ukernel/mmt4d.c b/runtime/src/iree/builtins/ukernel/mmt4d.c
index e775ca1..e582085 100644
--- a/runtime/src/iree/builtins/ukernel/mmt4d.c
+++ b/runtime/src/iree/builtins/ukernel/mmt4d.c
@@ -101,42 +101,17 @@
}
}
-// Helper for early-return path when K==0 and we just need to clear the output.
-static void iree_uk_mmt4d_zero_out(const iree_uk_mmt4d_params_t* params) {
- iree_uk_mmt4d_type_t mmt4d_type = iree_uk_mmt4d_type(params->flags);
- iree_uk_type_t out_type = iree_uk_mmt4d_out_type(mmt4d_type);
- int out_elem_size_log2 = iree_uk_type_size_log2(out_type);
- iree_uk_index_t contiguous_size = params->N * params->M0 * params->N0
- << out_elem_size_log2;
- iree_uk_index_t stride = params->out_stride0 << out_elem_size_log2;
- char* out_ptr =
- (char*)params->out_buffer + (params->out_offset << out_elem_size_log2);
- for (iree_uk_index_t i = 0; i < params->M; ++i) {
- iree_uk_memset(out_ptr, 0, contiguous_size);
- out_ptr += stride;
- }
-}
-
// Early-return code paths, including trivial or near-trivial cases (when one
// of the dimensions is 0) and in the future, hardware ports that specialize
// the entire loop nest.
// Returns true if already done.
static bool iree_uk_mmt4d_early(const iree_uk_mmt4d_params_t* params) {
// Trivial cases
- if (params->M == 0 || params->N == 0) {
+ if (params->M == 0 || params->N == 0 ||
+ (params->K == 0 && params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE)) {
return true;
}
- if (params->K == 0) {
- if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
- // Nothing to do!
- } else {
- iree_uk_mmt4d_zero_out(params);
- }
- return true;
- }
-
// Targets that want to specialize the entire loop nest can do so here.
-
return false;
}
