  Achieving high performance for Sparse MatrixMatrix Multiplication (SpMM) has
received increasing research attention, especially on multi-core CPUs, due to
the large input data size in applications such as graph neural networks (GNNs).
Most existing solutions for SpMM computation follow the aheadof-time (AOT)
compilation approach, which compiles a program entirely before it is executed.
AOT compilation for SpMM faces three key limitations: unnecessary memory
access, additional branch overhead, and redundant instructions. These
limitations stem from the fact that crucial information pertaining to SpMM is
not known until runtime. In this paper, we propose JITSPMM, a just-in-time
(JIT) assembly code generation framework to accelerated SpMM computation on
multi-core CPUs with SIMD extensions. First, JITSPMM integrates the JIT
assembly code generation technique into three widely-used workload division
methods for SpMM to achieve balanced workload distribution among CPU threads.
Next, with the availability of runtime information, JITSPMM employs a novel
technique, coarse-grain column merging, to maximize instruction-level
parallelism by unrolling the performance-critical loop. Furthermore, JITSPMM
intelligently allocates registers to cache frequently accessed data to
minimizing memory accesses, and employs selected SIMD instructions to enhance
arithmetic throughput. We conduct a performance evaluation of JITSPMM and
compare it two AOT baselines. The first involves existing SpMM implementations
compiled using the Intel icc compiler with auto-vectorization. The second
utilizes the highly-optimized SpMM routine provided by Intel MKL. Our results
show that JITSPMM provides an average improvement of 3.8x and 1.4x,
respectively.
