Towards a Java Virtual Machine for Processing-In-Memory

Processing-in-Memory (PIM) [1] is a computing paradigm in which computation takes place in memory devices. It attracts much attention because of its high memory bandwidth and energy-efficient data-parallel processing capabilities. The recent emergence of the first real-world commercial PIM system, UPMEM PIM [11], has driven attention even more. UPMEM PIM can be modeled as a distributed system in a single computer, consisting of around 2500 PIM processors called DPUs, each of which has its own DRAM called MRAM of 64 MB though DMA. CPU can read/write this MRAM and let DPU execute small programs. Unfortunately, software systems for PIM are still limited. Although an MPI-like library [2] for low-level C programming in the SPMD (single program, multiple data) style was developed, high-level object-oriented programming in managed languages has never been studied.

For high productivity, we pursue high-level programming for offloading Java applications to DPUs. Specifically, Figure 1 shows our motivating example, which consists of three key concepts: 1) single user-defined classes are executable on both CPUs and DPUs; 2) objects are created and kept in DPU memory; 3) handlers for results on DPUs are controlled on CPUs. These three concepts are natural demands on using PIM systems for Java applications. However, they pose non-trivial challenges in Java because the existing efforts on computation offloading in Java were made for GPU computing[5, 3] and mobile computing [9], of which the characteristics differ from UPMEM PIM. For example, in GPU computing, it is unusual to keep data in GPU memory after computation, and rich computing capabilities are available; in mobile computing, thousands of processor cores do not work cooperatively in parallel for a single computing task – UPMEM PIM systems are designed for these very things.

This paper is a status report on the development of our Java framework for PIM offloading. We design and implement a prototype Java VM that enables offloading Java methods written in the SPMD style to DPUs on UPMEM PIM to perform scalable data-parallel processing (Section 2). We also experimentally demonstrate the scalability of our Java VM with a simple benchmark on the UPMEM PIM system (Section 3).

Our programming framework consists of the compiler, the runtime system for the CPU and DPU sides (Figure 2). The compiler identifies the entry methods of Java programs (app.java in Figure 2), which are directly called from the CPU, and compiles each of them into separate binary for the DPUs. Each binary also contains the methods that may be directly or indirectly called from the entry method. For example, MV.mult and MV.sqnorm in Figure 1 are entry methods, and their binaries respectively contain Vector.dot as it is called therein. The runtime system for the DPU side is statically linked with each binary. The main program runs on the CPU. When the program calls an entry method through the host APIs, such as ctx.kernel and ctx.invoke in Figure 1, the CPU-side runtime system loads the corresponding binary to the DPUs and launches the DPUs.

We implemented the prototype of the compiler and the runtime system for the DPU side. The implementation of the runtime system for the CPU side is limited. It only supports a single DPU, and the context is not supported. In the prototype, SPMD support for DPU programming is low-level; a high-level API will be designed on top of it in the future.

We evaluate the performance and the scalability of our prototype framework. We develop a matrix-vector multiplication program, which is offloaded to a single DPU. The matrix is a $24\times 16,384$ matrix of non-negative integers. This is a submatrix obtained by dividing a $60,000\times 16,384$ matrix evenly for 2500 DPUs. In this experimentation, we use the $(𝒩,\text{max},+)$-algebra instead of the classic $(𝒩,+,\times )$-algebra. This is because the DPU does not support multiplication. This program was compiled with UPMEM SDK version 2023.2.0 and executed in a single DPU running at 350 MHz of the UPMEM PIM system [11].

Figure 3 shows the execution time. The “total time” indicates the entire execution time of the multiplication. This shows that performance was improved up to 4 threads. One of the reasons why it did not scale up to 11 threads was that the DPU had a single DMA controller. Because objects were in MRAM, all access to objects involved DMA operations. Thus, the DMA controller was a bottleneck.

We also measured the execution time of the program that did not perform array-boundary checks, whose results are the “no boundary check” in Figure 3. This shows that the overhead of array-boundary checks was significant. Array boundary check accesses the size field of the array, which is in MRAM. The fact that the performance scaled up to 6 threads supports that the bottleneck was the DMA controller.

The dashed line in Figure 3 shows the execution time of the same computation using the standard Java 17 HotSpot JVM, but the size of the matrix was $60,000\times 16,384$ because around 2500 DPUs are available in the UPMEM PIM system. For this measurement, we used two Intel Xeon Gold 6354 CPUs running at 3.0 GHz. This system has 72 hardware threads in total. This shows that offloading to DPU was $6.3\times$ slower even if we had used all available DPUs. However, considering our implementation is straightforward, we believe that there is a large room for improvement. For example, eliminating the array boundary checks [12] would improve the performance significantly.

This paper has reported our prototype Java VM for PIM offloading. Although our Java VM has not yet sufficiently reached our goal of high-level Java programming, the experimental results showed its potential on UPMEM PIM. We plan to implement higher-level APIs like Java 8 Stream as part of our software stack by using Babylon [7] and PIM-oriented in-memory databases [6] as the primary application. Development of a more sophisticated compiler and runtime systems to achieve high performance is also listed in our future work.