Skip to main content

Unleashing Java's GPU Power: A Rust-Powered Solution for NVIDIA Optimus

·5 mins

Hey 👋 this is Kosma! Welcome back to my personal blog where I share insights on JVM internals, performance optimization, and more. If you’re new here, I’m an experienced software engineer with a passion for scalable applications and clean code.

🚀 Introduction #

Today, we’re diving into a niche but fascinating topic: how to ensure your Java application harnesses the full power of your NVIDIA GPU. If you’ve ever found your Java app sluggishly running on integrated graphics when you know there’s a beast of a GPU sitting idle, this post is for you.

We’ll explore how to build a JVM executable that tells NVIDIA Optimus to use the dedicated GPU by exporting a specific symbol. It’s a bit of a journey involving Rust, JNI, and some low-level tinkering, but I promise it’ll be worth it.

🌟 The NVIDIA Optimus Conundrum #

First things first, let’s talk about NVIDIA Optimus. It’s a clever piece of tech designed to balance performance and battery life in laptops. The idea is simple: use the power-efficient integrated GPU for everyday tasks, and switch to the high-performance NVIDIA GPU when needed.

Sounds great, right? Well, it is, except when it isn’t. Sometimes, Optimus doesn’t recognize that your Java application is hungry for GPU power. The result? Your performance-intensive Java app ends up crawling along on integrated graphics.

For more on how Optimus works, check out the NVIDIA Optimus Technology Overview.

❓ The Problem: Java and GPU Detection #

Here’s where things get tricky. By default, Java applications often don’t trigger the switch to the dedicated GPU. This can be a huge issue if you’re doing anything graphics-intensive or computationally heavy.

The solution seems simple: just tell Optimus to use the dedicated GPU. And you can do that by exporting a specific symbol, NvOptimusEnablement, with a value of 0x00000001.

But here’s the catch: you can’t just do this from within your Java code. The symbol needs to be present when the executable is loaded, which happens before your Java code even starts running.

At this point, you might be thinking, “Why not just dynamically link to a library that exports this symbol?” Good thinking, but unfortunately, it won’t work. Dynamic linking happens too late in the process. By the time your Java app is up and running, Optimus has already made its decision about which GPU to use.

This is where our adventure with Rust begins.

🛠 The Rust-Powered Solution #

We’re going to use Rust to create a program that does two things:

  1. Exports the NvOptimusEnablement symbol
  2. Initializes the JVM and runs our Java application

Here’s how we’ll structure our project:

  • Cargo.toml: This will set up our Rust project and dependencies.
  • .cargo/config.toml: We’ll use this for some special build configurations.
  • src/main.rs: This is where the magic happens - our Rust code to export the symbol and fire up the JVM.
  • src/main/java/org/kosmadunikowski/nvidiaoptimusjvm/MainClass.java: A simple Java class to test our setup.

Let’s break it down:

Cargo.toml #

[package]
name = "nvidia_optimus_jvm"
version = "0.1.0"
edition = "2021"

[dependencies]
jni = { version = "0.21.1", features = ["invocation"] }

This sets up our project and pulls in the JNI crate, which we’ll use to interact with Java.

.cargo/config.toml #

[build]
rustflags = ["-C", "link-args=-uNvOptimusEnablement"]

This little snippet is crucial. It tells the Rust compiler to include our NvOptimusEnablement symbol in the final executable.

src/main.rs #

extern crate jni;

use jni::{InitArgsBuilder, JNIEnv, JNIVersion, JavaVM};

// Export the symbol for NVIDIA Optimus
#[no_mangle]
#[link_section = ".data"]
pub static NvOptimusEnablement: u32 = 0x00000001;

fn main() {
    // Initialize the JVM and run the Java application
    init_jvm_and_run_java();
}

// Function to initialize the JVM and run a Java application
fn init_jvm_and_run_java() {
    // Set up JVM arguments
    let jvm_args = InitArgsBuilder::new()
        .version(JNIVersion::V8)
        .option("-Djava.class.path=.") // Adjust the class path as necessary
        .build()
        .unwrap();

    // Create the JVM
    let jvm = JavaVM::new(jvm_args).expect("Failed to create JVM");

    // Attach the current thread to call into Java
    let mut env = jvm
        .attach_current_thread()
        .expect("Failed to attach current thread");

    // Call the main method of the Java application
    run_java_application(&mut env, "org/kosmadunikowski/nvidiaoptimusjvm/MainClass", "main");
}

// Function to call the main method of a Java class
fn run_java_application(env: &mut JNIEnv, class_name: &str, method_name: &str) {
    // Find the Java class
    let class = env.find_class(class_name).expect("Failed to find class");

    // Get the method ID for the main method
    let method_id = env
        .get_static_method_id(&class, method_name, "()V")
        .expect("Failed to get method ID");

    unsafe {
        // Call the main method
        env.call_static_method_unchecked(
            &class,
            method_id,
            jni::signature::ReturnType::Primitive(jni::signature::Primitive::Void),
            &[],
        )
        .expect("Failed to call method");
    }
}

This is where the real work happens. We export our NvOptimusEnablement symbol, set up the JVM, and call into our Java code.

src/main/java/org/kosmadunikowski/nvidiaoptimusjvm/MainClass.java #

package org.kosmadunikowski.nvidiaoptimusjvm;

public class MainClass {
    public static void main() {
        System.out.println("Hello, World!");
    }
}

A simple Java class to test our setup. In a real-world scenario, this would be your actual Java application.

🏗️ Building and Running #

To bring this all together:

  1. Make sure you have Rust installed. If not, grab it from rust-lang.org.
  2. Create a new Rust project: cargo new nvidia_optimus_jvm
  3. Replace the default main.rs with our version and add the Java file.
  4. Build the project: cargo build --release
  5. Run the executable from target/release. Make sure your Java class path is set correctly.

And voila! Your Java application should now be running with the full power of your NVIDIA GPU.

📝 Wrapping Up #

This Rust-based solution is a neat way to ensure your Java application taps into the power of your dedicated GPU. It’s not without its trade-offs, though. By implementing the binary ourselves, we lose some of the features of the original Java CLI. If you need those features, you might have to reimplement them in Rust.

For more on the Java CLI and what you might be missing, check out the Java CLI Command Documentation.

If you found this deep dive interesting, you might enjoy my previous post on JNI and region pinning in G1 GC. It’s another exploration of the intricate world of JNI and performance optimization.

As always, I’d love to hear your thoughts and experiences. Have you run into GPU issues with Java before? How did you solve them? Let me know in the comments!

Until next time, keep optimizing!

Kosma