Debugging JNI calls to the GPU

I'm playing aroung with a Java based LLM (code here). When running a JVM that calls the GPU using the TornadoVM, it crashed and in the log, I saw:

Native frames: (J=compiled Java code, j=interpreted, Vv=VM code, C=native code)

C  [libcuda.so.1+0x1302b0]

C  [libcuda.so.1+0x332420]

C  [libtornado-ptx.so+0x64b8]  Java_uk_ac_manchester_tornado_drivers_ptx_PTXStream_cuLaunchKernel+0x198

j  uk.ac.manchester.tornado.drivers.ptx.PTXStream.cuLaunchKernel([BLjava/lang/String;IIIIIIJ[B[B)[[B+0 tornado.drivers.ptx@2.2.1-dev

...

Now, finding the Shared Object files (*.so), I called: 

objdump -d /usr/lib/x86_64-linux-gnu/libcuda.so.1 

objdump -d /usr/local/bin/Java/tornadovm-2.2.1-dev-ptx/lib/libtornado-ptx.so

and looked at the addresses in the stack dump.

First, libtornado-ptx.so. Note that the address (0x64b8) is the return address from a call, that is the next line after the call that went Pete Tong. 

    64b3:       e8 b8 e1 ff ff          call   4670 <cuLaunchKernel@plt>

    64b8:       48 83 c4 30             add    $0x30,%rsp

So, it's the call to cuLaunchKernel that is interesting.

  33241b:       e8 00 de df ff          call   130220 <exit@plt+0x4e460>

  332420:       5a                      pop    %rdx

and the final (top most) stack frame:

  1302ab:       4d 85 e4                test   %r12,%r12

  1302ae:       74 58                   je     130308 <exit@plt+0x4e548>

  1302b0:       41 8b 04 24             mov    (%r12),%eax

The instruction test %x,%y is a common idiom in null checks (basically, it's x and y are ANDed and the je jumps if the Zero Flag is set - note that this flag is set if the result of the AND is non-zero or both x and y are zero).

So, it looks like we've essentially got what's equivalent to a NullPointerException in the machine code. Still looking at what's null... [Solved: had to use a model that is compatible with GPULlama3.java)

Java and the GPU

Java was always built to abstract away the hardware on which it runs but it's approach to GPU has been somewhat late to the game [Gary Frost, YouTube].

There are projects out there that promise to give Java access to the GPU. I looked at Aparapi but it appears to be moribud. So, I gravitated to TornadoVM which Frost describes as "state of the art".

The trouble is that TornadoVM runs everything in a Docker image that has all the shared objects built in. This is fine for a quick demo - this is the result of running on my Quadro T2000:

docker-tornadovm$ ./run_nvidia_openjdk.sh tornado -cp example/target/example-1.0-SNAPSHOT.jar example.MatrixMultiplication
Computing MxM of 512x512
CPU Execution: 1.17 GFlops, Total time = 230 ms
GPU Execution: 268.44 GFlops, Total Time = 1 ms
Speedup: 230x

This demonstrates how the GPU runs a nested for-loop doing matrix multiplication much faster than the same code on the CPU. But it runs it all in a Docker container and I need to package a JAR everytime I make a change. How do I run it outside the container?

To work this out, I opened a shell in the Docker image and saw that the TornadoVM build it uses was built from Git branch d3062accc. So, the first thing was to checkout that branch of TornadoVM and build it.

I built with:

mvn clean install -Pgraal-jdk-11-plus

using the graalvm-ee-java11-22.3.4 JDK.

Note that you'll need Graal as the TornadoVM code has dependencies on it. I built my own Graal JDK by following the instructions here but using a different branch as I couldn't find the download for the graal.version defined in the TornadoVM pom.xml. Note, you'll also need mx and a bootstrapping JDK that has the right compiler interface (JVMCI), in my case labsjdk-ce-21.0.1-jvmci-23.1-b19.

So far, so good. I ran the tornado script which is just a wrapper around a call to the java executable (don't forget to set your JAVA_HOME environment variable to point at the Graal JDK) but it complained it could not see a tornado.backend file.

Again, a sneaky look at the Docker container indicated that we have to tell it which driver to use. So, I created the file and told it tornado.backends=opencl-backend but then tornado complained it didn't have the OpenCL drivers. Oops. 

You have to build the drivers you want seperately it seems. But if you try to build Tornado drivers without the native OpenCL dev library, you'll see:

TornadoVM/tornado-drivers/opencl-jni$ mvn clean install # yes, Maven cmake via cmake-maven-plugin
....
/usr/bin/ld: cannot find -lOpenCL
...

The Docker image saves you from having to install the OpenCL libraries on your machine. To get it working on bare metal, I played it safe and got an old Ubuntu box and installed them there. You'll need to install them with:

sudo apt install ocl-icd-opencl-dev

and then ran Maven in the opencl* sub directories. This time, the Maven build completed successfully.

However, running tornado in the subsequent dist folder still pukes but with something like:

Caused by: uk.ac.manchester.tornado.api.exceptions.TornadoRuntimeException: OpenCL JNI Library not found

at tornado.drivers.opencl@0.15.1/uk.ac.manchester.tornado.drivers.opencl.OpenCL.<clinit>(OpenCL.java:68)

... 11 more

Not what I was expecting. I found I needed to:

cp ./tornado-drivers/opencl-jni/target/linux-amd64-release/cmake/libtornado-opencl.so $TORNADO_SDK/lib

Where TORNADO_SDK is pointing at the relevent dist folder.

Now, finally, you can run on the bare metal:

$ tornado -cp target/classes/ example.MatrixMultiplication

Computing MxM of 512x512

CPU Execution: 1.21 GFlops, Total time = 222 ms

GPU Execution: 17.90 GFlops, Total Time = 15 ms

Speedup: 14x

(Results from an old NVIDIA GeForce GTX 650)

Note, you'll need to also run it with the Graal JVM. Set both the PATH and JAVA_HOME environment variables to point to it.

Where now?

This is a nice introduction to running Java on the GPU but it's just the start. There are many caveats. Example: what if your Java code throws an Exception? GPUs have no equivalent of exceptions so what happens then? More to come.

Low Level Locks

Deep down in the JVM you'll find something like this in Atomic::cmpxchg:

  __asm__ volatile ("lock; 1: cmpxchgl %1,(%3)"
                    : "=a" (exchange_value)
                    : "r" (exchange_value), "a" (compare_value), "r" (dest)
                    : "cc", "memory");
  return exchange_value;

(Note: I've slightly modified this for reasons of clarity to ignore the case where the code might be running on a CPU that is not a multi-processor).

This code is hit often when the JVM tries to lock something. It is one of its many techniques to arbitrate contention (what happens in the event of contention I'll leave to another post).

So, what's it doing? It's basically doing what Java programmers would know from the AtomicXXX classes using low-level assembler. This is how GCC makes explicit assembler code. A good reference is here but briefly:

• asm (or alternatively __asm__) are keywords that allow the programmer to directly write assembler. Volatile means don't move our code as part of an optimization. It must execute where we placed it.
• The contents of the parentheses are delimited by colons and their meanings are ASSEMBLY_CODE:OUTPUT:INPUT:CLOBBERED.
• The CLOBBERED list tells the compiler what we have fiddled with (other than input and output) and that it should bear this mind before and after executing it. In this case, we're telling it we fiddled with the memory and the condition flags ("cc").
• The INPUT and OUTPUT fields map our C variables to particular CPU registers. In this case, we're saying use any registers for the values input via exchange_value and dest because "r" means "any register". For the input in compare_value and output in exchange_value we're saying explicitly use the EAX register ("a" and "=a" respectively. The equals sign means it is a write-only output.

Why the EAX register is used for input and output is given in Intel's Software Developer's Manual:

"The CMPXCHG instruction requires three operands: a source operand in a register, another source operand in the EAX register, and a destination operand. If the values contained in the destination operand and the EAX register are equal, the destination operand is replaced with the value of the other source operand (the value not in the EAX register). Otherwise, the original value of the destination operand is loaded in the EAX register." [1]

If you put this code in a simple .c file and compile it with:

$ gcc -S YOUR_FILE.c

you'll see the raw assembly output to YOUR_FILE.s confirms the above.

[1] Intel® 64 and IA-32 Architectures Software Developer’s Manual Combined Volumes: 1, 2A, 2B, 2C, 3A, 3B, and 3C

TLABs

I sometimes peruse the hotspoc-gc-dev lists for grins and saw this post on G1 garbage collecting:

"Unlike the other collectors G1 did not use dynamic sizes for the TLABs that the Java threads use for allocation. That means that the TLABs when using G1 were all fixed size and very small. This patch adds dynamic sizing of TLABs to G1 too..."

What are these mysterious TLABs?

"A Thread Local Allocation Buffer (TLAB) is a region of Eden that is used for allocation by a single thread.  It enables a thread to do object allocation using thread local top and limit pointers, which is faster than doing an atomic operation on a top pointer that is shared across threads." [1]

Or, to put it another way:

"In HotSpot, almost all application allocation is performed in the eden space of the young generation, using contiguous (aka "pointer-bumping") allocation. That is, there's a current pointer and an end address, and as long the next allocation fits we return the current pointer value and increment it by the allocation size.

"But on a multiprocessor, this isn't safe: two threads could both read the current alloc pointer value, and both update, and think they were using the same memory for different objects. So some form of synchronization is necessary. We could just take a lock, but we can do better by using an atomic hardware instruction, a compare-and-swap, or CAS (in the SPARC architecture; the same thing is called compare-and-exchange on x86). So in our race scenario, one thread's CAS succeeds and the other fails; the latter retries.

"This still has two problems: atomic hardware operations are expensive on most architectures, and on machines with many processors this could be a source of cache contention, making them even more expensive.

"So we avoid this with TLABs: each thread allocates a medium-sized chunk and saves it, allocating within it with no synchronization required. Only when its TLAB is used does it go back to allocating from the shared space." [2]

Detlef points out: the bigger the TLAB, the less often you must contend for shared memory but also the more young generation you will use [2].

There are JVM flags to manipulate them (-XX:TLABSize, -XX:-ResizeTLAB) and even to print summaries (-XX:+PrintTLAB) [1].

[1] Jon Masamitsu's Weblog.
[2] David Detlef's Weblog.

FastLocks in the JVM

I mentioned how AtomicIntegerArray and its kin can be slow. I'm now a little closer to finding why.

I decompiled my code that has one thread writing to and one thread reading from an AtomicIntegerArray using the JVM arguments:

-Xmx32g -server -XX:+PrintOptoAssembly -XX:CompileCommand=print,*Runnable -XX:CompileThreshold=10000

Using Java version:

$ /usr/java/jdk1.7.0_10/bin/java -version
java version "1.7.0_10"
Java(TM) SE Runtime Environment (build 1.7.0_10-b18)
Java HotSpot(TM) 64-Bit Server VM (build 23.6-b04, mixed mode)

My code is testing reading and writing to an array structure. Initially, all values are 0. 

The writer thread goes through the array setting the values to 1. 

The reader thread follows the writer as best it can. If an element is 0 (that is, uninitialized) it spins waiting for the value to change to 1.

The Java code for reading looks something like this in the super class:

public abstract class AbstractReaderRunnable implements Runnable {    

    protected int index = 0;

    protected int lap = 0;

    @Override

    public void run() {

        while (index < ARRAY_SIZE) {

            lap = elementAtIndex();

            while (lap == 0) {

                lap = elementAtIndex();

            }

            index++;

        }

    }

    protected abstract int elementAtIndex();

and in my particular subclass (that is, a class for which the array structure is a AtomicIntegerArray) looks like this:

class AtomicReaderRunnable extends AbstractReaderRunnable {

    private final AtomicIntegerArray array;

    protected int elementAtIndex() {

        return array.get(index);

    }

.

.

Fairly simple. Now comes the hard part: the machine code.

Note: "In Intel syntax the first operand is the destination, and the second operand is the source whereas in AT&T syntax the first operand is the source and the second operand is the destination." [1]

HotSpot by default appear to print the assembly in Intel format, although you can change this [2]

069     movl    R10, [R13 + #24 (8-bit)]        # int ! Field com/phenry/concurrent/lockless/AbstractReaderRunnable.index

06d     movq    R8, [R11 + #16 (8-bit)] # ptr ! Field java/util/concurrent/atomic/AtomicIntegerArray.array

.

.

So, the R10 registry is our index field and R8 points to the int[] array in the Java AtomicIntegerArray class.

.

.

094     movl    R11, [R8 + #24 + R10 << #2]     # int

099     MEMBAR-acquire ! (empty encoding)

099

099     movl    [R13 + #28 (8-bit)], R11        # int ! Field com/phenry/concurrent/lockless/AbstractReaderRunnable.lap

In English: here we multiply R10 by 4 by shifting it twice to the left (<< #2) - because ints are 4 bytes wide - and apparently add this to the start of the array to find our memory address. The contents of this we put into the R11 registry. In turn, this is put in the memory address where we store our lap field in the Java world.

Note that HotSpot appears to have inlined methods [6]. The code generated by Hotspot is not necessarily isomorphic to the original Java code. Also note that it might recompile it to machine code more than once, apparently depending on the usage profile.

Now, the thing of interest is the MEMBAR-acquire ! (empty encoding) line. This is not an x86 instruction and doesn't even take an address (the left-most column). So, why is it there?

"Each CPU has its own peculiar memory-barrier instructions, which can make portability a challenge" [3] The x86 hardware has a NoOp for most membar operations:

void LIR_Assembler::membar() {
  // QQQ sparc TSO uses this,
  __ membar( Assembler::Membar_mask_bits(Assembler::StoreLoad));
}

void LIR_Assembler::membar_acquire() {
  // No x86 machines currently require load fences
  // __ load_fence();
}

void LIR_Assembler::membar_release() {
  // No x86 machines currently require store fences
  // __ store_fence();
}

void LIR_Assembler::membar_loadload() {
  // no-op
  //__ membar(Assembler::Membar_mask_bits(Assembler::loadload));
}

void LIR_Assembler::membar_storestore() {
  // no-op
  //__ membar(Assembler::Membar_mask_bits(Assembler::storestore));
}

void LIR_Assembler::membar_loadstore() {
  // no-op
  //__ membar(Assembler::Membar_mask_bits(Assembler::loadstore));
}

void LIR_Assembler::membar_storeload() {
  __ membar(Assembler::Membar_mask_bits(Assembler::StoreLoad));
}

(from Hotspot's source code that can be found in hotspot/src/cpu/x86/vm/c1_LIRAssembler_x86.cpp).

The membar method itself looks like this:

  // Serializes memory and blows flags
  void membar(Membar_mask_bits order_constraint) {
    if (os::is_MP()) {
      // We only have to handle StoreLoad
      if (order_constraint & StoreLoad) {
        // All usable chips support "locked" instructions which suffice
        // as barriers, and are much faster than the alternative of
        // using cpuid instruction. We use here a locked add [esp],0.
        // This is conveniently otherwise a no-op except for blowing
        // flags.
        // Any change to this code may need to revisit other places in
        // the code where this idiom is used, in particular the
        // orderAccess code.
        lock();
        addl(Address(rsp, 0), 0);// Assert the lock# signal here
      }
    }
  }

(from assembler_x86.hpp in the HotSpot source code).

And true to the comment, we do see lines output like:

     lock addl [rsp + #0], 0 ! membar_volatile

but it's in the writer code! The writer code looks like this:

070     movl    R8, [RCX + #28 (8-bit)] # int ! Field com/phenry/concurrent/lockless/AbstractWriterRunnable.lap

.

.

078     movl    R11, [RCX + #24 (8-bit)]        # int ! Field com/phenry/concurrent/lockless/AbstractWriterRunnable.index

07c     movq    R10, [R10 + #16 (8-bit)]        # ptr ! Field java/util/concurrent/atomic/AtomicIntegerArray.array

.

.

Similar to before we have a register that's our index field (R11) and  the int[] array in the Java AtomicIntegerArray class (R10). Also similar to before, we calculate the address that this index points to in the array but this time we populate its value with R8 (our lap field in our Java code).

.

.

08e   B8: #     B4 B9 <- 287336="" b7="" font="" nbsp="" req:="">

08e     MEMBAR-release ! (empty encoding)

08e

08e     movslq  R11, R11        # i2l

091     movl    [R10 + #24 + R11 << #2], R8     # int

096     lock addl [rsp + #0], 0 ! membar_volatile

As before we have the (redundant) MEMBAR pseudo instruction but we also have the very significant lock prefix.

The reader code does not need to worry since the membar has flushed the store buffers (which "enable our fast processors to run without blocking while data is transferred to and from the cache sub-system" [7]. These fences are more related to instruction ordering - another side effect of the synchronize semantics - than the cache subsystem). Furthermore:

"If a core wants to read some memory, and it does not have it ... then it must make a read on the ring bus.  It will then either be read from main-memory if not in the cache sub-systems, or read from L3 if clean, or snooped from another core if Modified.  In any case the read will never return a stale copy from the cache sub-system, it is guaranteed to be coherent." [7]

This seems to be a standard way of synchronizing outside the JVM [5]. So, why is it so slow? The JVM has some very clever code to avoid this heavy-weight approach that I will go into in a future post.

[1] http://www.ibiblio.org/gferg/ldp/GCC-Inline-Assembly-HOWTO.html

[2] http://stackoverflow.com/questions/9337670/hotspot7-hsdis-printassembly-intel-syntax

[3] Memory Barriers: a Hardware View for Software Hackers

[4] x86-64 Machine-Level Programming .

[5] StackOverflow.
[6] "Deep dive into assembly code from Java" -Kohsuke Kawaguchi's blog.

[7] "CPU Cache Flushing Fallacy" - Martin Thompson's blog.

The Case of the Missing Memory

If you start Java with -Xmx512m, you'd expect the max heap size to be 512 megs, right? So, when you connect via JConsole (or JMX) why is the heap size smaller? My JConsole shows a committed and maximum size of the "Heap Memory Usage" as about 491MiB.

I've been fiddling with the OpenJDK source and been putting log statements into it to see how the GC is behaving. So, for instance, I changed:

hotspot/src/share/vm/gc_implementation/parallelScavenge/parallelScavengeHeap.cpp

added:

#include <stdio.h>

and liberally sprinkled the code with statements like:


  printf("PH: ParallelScavengeHeap::initialize: size %d: \n", _reserved.byte_size());

And sure enough, when I run my application, I see the JVM spit out:

PH: ParallelScavengeHeap::initialize: size 536870912: 

So, it's using the collector that is the "default on certain machine types" [1] with the full 512MiB allocated to it. This collector is "targeted towards large young generation heaps and to scale with more CPUs. It works very well with large young generation heap sizes that are in gigabytes, like 12GB to 80GB or more, and scales very well with increase in CPUs, 8 CPUs or more" [2] (I have 16 CPUs and 64GBs, although obviously have not set my heap that high).

Now, when I connect to the JVM and look at the memory usage, I see statements like:

PH: EdenMutableSpacePool::get_memory_usage: size 134217728: 

PH: SurvivorMutableSpacePool::get_memory_usage: size 22347776: 

PH: PSGenerationPool::get_memory_usage: size 357957632: 

since I have changed these C++ classes that can all be found in

hotspot/src/share/vm/services/psMemoryPool.cpp

(Note that these three classes are all the subtypes of CollectedMemoryPool for the PS collector.)

This memory usage totals the 491MiB usage we saw in JConsole but is still exactly 22 347 776 bytes short of the 512MiB we defined at start up time.

The "missing" memory is of course because there are two survivor spaces [3] but since only one has live objects at any point, we only count its size.

[1] Dr Richard Warburton's blog.
[2] Oracle's Improving Java Application Performance and Scalability by Reducing Garbage Collection Times and Sizing Memory Using JDK 1.4.1.
[3] Oracle's Virtual Machine Garbage Collection Tuning.