Integrate JIT generated GEMM microkernels into create_convolution2d_nhwc
Introduce a new field, generator, into struct gemm_parameters, which contains JIT code generators for gemm, igemm, gemm1, igemm1. When set, the convolution operator creation will try to generate the microkernel using the JIT. (Right now only gemm is supported, the rest will follow in future patches.)
The xnn_ukernel_gemm and xnn_ukernel_igemm structs also has a new field, struct xnn_code_buffer general_code_buffer and mr1_code_buffer, where the generated code will be kept, and is released when the operator is deleted.
The generator field is only set in the e2e benchmarks, where we update F32 E2E benchmarks to support testing JIT generated microkernels.
PiperOrigin-RevId: 425700057
diff --git a/BUILD.bazel b/BUILD.bazel
index a744f2f..fa77684 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -9134,6 +9134,7 @@
msvc_copts = xnnpack_msvc_std_copts(),
deps = [
":indirection",
+ ":jit_memory",
":logging_utils",
":operator_run",
":packing",
@@ -9167,6 +9168,7 @@
msvc_copts = xnnpack_msvc_std_copts(),
deps = [
":indirection_test_mode",
+ ":jit_memory_test_mode",
":logging_utils",
":operator_run_test_mode",
":packing_test_mode",
@@ -9178,12 +9180,39 @@
)
xnnpack_cc_library(
+ name = "jit_memory",
+ srcs = [
+ "src/jit/memory.c",
+ ],
+ hdrs = INTERNAL_HDRS,
+ msvc_copts = xnnpack_msvc_std_copts(),
+ deps = [
+ ":logging_utils",
+ ],
+)
+
+xnnpack_cc_library(
+ name = "jit_memory_test_mode",
+ srcs = [
+ "src/jit/memory.c",
+ ],
+ hdrs = INTERNAL_HDRS,
+ copts = [
+ "-UNDEBUG",
+ "-DXNN_TEST_MODE=1",
+ ],
+ msvc_copts = xnnpack_msvc_std_copts(),
+ deps = [
+ ":logging_utils",
+ ],
+)
+
+xnnpack_cc_library(
name = "jit",
srcs = [
"src/jit/aarch32-assembler.cc",
"src/jit/aarch64-assembler.cc",
"src/jit/assembler.cc",
- "src/jit/memory.c",
],
hdrs = INTERNAL_HDRS + [
"src/xnnpack/aarch32-assembler.h",
@@ -9194,6 +9223,7 @@
aarch64_srcs = JIT_AARCH64_SRCS,
msvc_copts = xnnpack_msvc_std_copts(),
deps = [
+ ":jit_memory",
":logging_utils",
],
)
@@ -9204,7 +9234,6 @@
"src/jit/aarch32-assembler.cc",
"src/jit/aarch64-assembler.cc",
"src/jit/assembler.cc",
- "src/jit/memory.c",
],
hdrs = INTERNAL_HDRS + [
"src/xnnpack/aarch32-assembler.h",
@@ -9219,6 +9248,7 @@
],
msvc_copts = xnnpack_msvc_std_copts(),
deps = [
+ ":jit_memory_test_mode",
":logging_utils",
],
)